github
/
ChatGLM-6B
mirror of https://github.com/THUDM/ChatGLM-6B
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Merge pull request #225 from holk-h/main 

Support stream out interruption by using Ctrl+C
pull/266/head
Zhengxiao Du 2023-03-28 19:50:29 +08:00 committed by GitHub
parent 6fc8141a9c bf39dac067
commit 25356d6441
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 17 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
cli_demo.py
View File

@ -1,14 +1,15 @@
import os import os
import platform import platform
import signal
from transformers import AutoTokenizer, AutoModel from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True) tokenizer = AutoTokenizer.from_pretrained("./model", trust_remote_code=True)
model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().cuda() model = AutoModel.from_pretrained("./model", trust_remote_code=True).half().cuda()
model = model.eval() model = model.eval()
os_name = platform.system() os_name = platform.system()
clear_command = 'cls' if os_name == 'Windows' else 'clear' clear_command = 'cls' if os_name == 'Windows' else 'clear'
stop_stream = False
def build_prompt(history): def build_prompt(history):
prompt = "欢迎使用 ChatGLM-6B 模型，输入内容即可进行对话，clear 清空对话历史，stop 终止程序" prompt = "欢迎使用 ChatGLM-6B 模型，输入内容即可进行对话，clear 清空对话历史，stop 终止程序"
@ -17,9 +18,13 @@ def build_prompt(history):
prompt += f"\n\nChatGLM-6B：{response}" prompt += f"\n\nChatGLM-6B：{response}"
return prompt return prompt
def signal_handler(signal, frame):
global stop_stream
stop_stream = True
def main(): def main():
history = [] history = []
global stop_stream
print("欢迎使用 ChatGLM-6B 模型，输入内容即可进行对话，clear 清空对话历史，stop 终止程序") print("欢迎使用 ChatGLM-6B 模型，输入内容即可进行对话，clear 清空对话历史，stop 终止程序")
while True: while True:
query = input("\n用户：") query = input("\n用户：")
@ -32,10 +37,15 @@ def main():
continue continue
count = 0 count = 0
for response, history in model.stream_chat(tokenizer, query, history=history): for response, history in model.stream_chat(tokenizer, query, history=history):
if stop_stream:
stop_stream = False
break
else:
count += 1 count += 1
if count % 8 == 0: if count % 8 == 0:
os.system(clear_command) os.system(clear_command)
print(build_prompt(history), flush=True) print(build_prompt(history), flush=True)
signal.signal(signal.SIGINT,signal_handler)
os.system(clear_command) os.system(clear_command)
print(build_prompt(history), flush=True) print(build_prompt(history), flush=True)