[THUDM/ChatGLM-6B]添加了流式解码器,做到更好的控制台体验和更高的解码效率

2024-05-21 380 views
4
流式解码

目前的流式输出方式是每次都完整decode一遍整个输出内容,并且刷新屏幕。这样效率不是很高,同时也不利于通过网络来流式传输数据。最理想的方式应该是做到流式解码,每次只解码增量部分,于是我添加了一个SentencePiece的流式解码器以及一个配套的demo。

代码说明

解码器的原始实现参考的是SentencePieceProcessor的CPP代码,重写为python代码的同时为流式解码做了一些更改,主要集中在多字节编码的处理和句首空格(bos_ws)处理上。可以避免现有流式输出中有可能出现的因为多字节编码未完整输出而产生的ReplacementCharacter(�)。同时也可以实现无闪烁的控制台打字机输出以及更方便的基于websocket的传输。

流式解码器方法

一共三个方法,分别是put()``end()以及get(),put用于输入index并解码,end用于结束流式解码器,get用于从解码缓冲区中获取内容并清空缓冲区,详细使用方法见demo。

一些希望

希望这个解码器未来可以进入到模型代码仓库中(HuggingFace那个),并且整合进stream_chat方法中。 如果有什么疑问和建议,欢迎和我讨论。

回答

9

很有价值的实现,我需要检验一下流式解码和原 tokenizer 是否是等价的。可能存在的问题是现在 stream_generate 会对模型的输出做一些后处理,https://huggingface.co/THUDM/chatglm-6b/blob/main/modeling_chatglm.py#L1251

确实还没有考虑后处理,目前我只做了<n>的替换,而且是在demo里面实现的。我目前这个StreamDecoder只是根据SentencePiece修改的,所以还没研究具体GLM是怎么处理的。后续我试试能不能在这个Decoder外面再封装一层,以达到与GLM Decoder相同的效果。

3

@duzx16 我通过单元测试发现了一些bug,修好之前我先关闭这个pr

4

我把输出进行了对齐,并且直接在stream_utils.py中添加了单元测试。 目前已知没有做到输出对齐的情况是这样的:模型输出:" 1 ",应该打印:"1",实际打印" 1 ",没有实现strip()效果。 我想不到什么能简单实现strip()函数效果的流式输出方法,而且感觉必要性不是很强,就先没有做这个情况下的输出对齐了。