ChatGLM3-6B对话流式输出

ChatGLM3-6B被太多框架使用，基本的流式输出函数竟然一时找不到了。
目前看到的有两个地方有相关的：
第一个：
ChatGLM3/basic_demo/cli_demo.py at main · THUDM/ChatGLM3
https://github.com/THUDM/ChatGLM3/blob/main/basic_demo/cli_demo.py

stop_stream = False
def main():
    past_key_values, history = None, []
    global stop_stream
    print(welcome_prompt)
    while True:
        query = input("\n用户：")
        if query.strip() == "stop":
            break
        if query.strip() == "clear":
            past_key_values, history = None, []
            os.system(clear_command)
            print(welcome_prompt)
            continue
        print("\nChatGLM：", end="")
        current_length = 0
        for response, history, past_key_values in model.stream_chat(tokenizer, query, history=history, top_p=1,
                                                                    temperature=0.01,
                                                                    past_key_values=past_key_values,
                                                                    return_past_key_values=True):
            if stop_stream:
                stop_stream = False
                break
            else:
                print(response[current_length:], end="", flush=True)
                current_length = len(response)
        print("")

这是一个命令行运行大模型的例子，是流式输出的，但是用户和ai的历史输出是直接放在字符串中拼接的，我倾向于一个json结构的对话历史。

第二个：
THUDM/ChatGLM3: ChatGLM3 series: Open Bilingual Chat LLMs | 开源双语对话语言模型 https://github.com/THUDM/ChatGLM3

>> from transformers import AutoTokenizer, AutoModel
>> tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm3-6b", trust_remote_code=True)
>> model = AutoModel.from_pretrained("THUDM/chatglm3-6b", trust_remote_code=True, device='cuda')
>> model = model.eval()
>> response, history = model.chat(tokenizer, "你好", history=[])
>> print(response)

你好👋!我是人工智能助手 ChatGLM3 - 6B, 很高兴见到你, 欢迎问我任何问题。
>> response, history = model.chat(tokenizer, "晚上睡不着应该怎么办", history=history)
>> print(response)

这是官网给的例子，不是流式输出，不过输出历史都是用[{‘role’: ‘user’, ‘content’:…}]格式保存的。
这两个能结合起来吗？注意到方法model.stream_chat和model.chat形式很像，尝试拼接了一下，发现居然可行。
封装为函数时，history作为外部变量保存历史，出现了报错：UnboundLocalError: cannot access local variable ‘history’ where it is not associated with a value。
不想用global关键字，也不想把history作为参数传入函数。
把内部循环的history变量修改为inner_history解决了问题。

from transformers import AutoTokenizer, AutoModel

tokenizer = AutoTokenizer.from_pretrained("/new/workspace/models/ChatGLM3/chatglm3-6b", trust_remote_code=True)
model = AutoModel.from_pretrained("/new/workspace/models/ChatGLM3/chatglm3-6b", trust_remote_code=True, device='cuda')
model = model.eval()

def ask(query='你好',new_chat=False):
    past_key_values=None
    stop_stream = False
    current_length = 0
    for response, inner_history, past_key_values in model.stream_chat(tokenizer, query, history=history, top_p=1,
                                                                temperature=0.01,
                                                                past_key_values=past_key_values,
                                                                return_past_key_values=True):
        if stop_stream:
            stop_stream = False
            break
        else:
            print(response[current_length:], end="", flush=True)
            current_length = len(response)
    print("")
    history.clear()
    history.extend(inner_history)

history=[]
ask('一年有几个季节')
ask('第四个有什么特征')

输出效果：
一年通常分为四个季节,分别是春季、夏季、秋季和冬季。但不同地区的气候和季节变化可能会略有不同。
冬季通常出现在北半球的冬季,而在南半球则是夏季。在冬季,气温通常较低,天气也较为寒冷,天气也较为阴沉,降水也较为少。植物也会进入休眠状态,动物也会采取不同的生存策略以适应寒冷的环境。

history内容：

[{'role': 'user', 'content': '一年有几个季节'},
 {'role': 'assistant',
  'metadata': '',
  'content': '一年通常分为四个季节,分别是春季、夏季、秋季和冬季。但不同地区的气候和季节变化可能会略有不同。'},
 {'role': 'user', 'content': '第四个有什么特征'},
 {'role': 'assistant',
  'metadata': '',
  'content': '冬季通常出现在北半球的冬季,而在南半球则是夏季。在冬季,气温通常较低,天气也较为寒冷,天气也较为阴沉,降水也较为少。植物也会进入休眠状态,动物也会采取不同的生存策略以适应寒冷的环境。'}]