mirror of https://github.com/THUDM/ChatGLM-6B
Fix model path
parent
25356d6441
commit
343e7bc7b6
|
@ -3,14 +3,15 @@ import platform
|
|||
import signal
|
||||
from transformers import AutoTokenizer, AutoModel
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained("./model", trust_remote_code=True)
|
||||
model = AutoModel.from_pretrained("./model", trust_remote_code=True).half().cuda()
|
||||
tokenizer = AutoTokenizer.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True)
|
||||
model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().cuda()
|
||||
model = model.eval()
|
||||
|
||||
os_name = platform.system()
|
||||
clear_command = 'cls' if os_name == 'Windows' else 'clear'
|
||||
stop_stream = False
|
||||
|
||||
|
||||
def build_prompt(history):
|
||||
prompt = "欢迎使用 ChatGLM-6B 模型,输入内容即可进行对话,clear 清空对话历史,stop 终止程序"
|
||||
for query, response in history:
|
||||
|
@ -18,10 +19,12 @@ def build_prompt(history):
|
|||
prompt += f"\n\nChatGLM-6B:{response}"
|
||||
return prompt
|
||||
|
||||
|
||||
def signal_handler(signal, frame):
|
||||
global stop_stream
|
||||
stop_stream = True
|
||||
|
||||
|
||||
def main():
|
||||
history = []
|
||||
global stop_stream
|
||||
|
@ -45,7 +48,7 @@ def main():
|
|||
if count % 8 == 0:
|
||||
os.system(clear_command)
|
||||
print(build_prompt(history), flush=True)
|
||||
signal.signal(signal.SIGINT,signal_handler)
|
||||
signal.signal(signal.SIGINT, signal_handler)
|
||||
os.system(clear_command)
|
||||
print(build_prompt(history), flush=True)
|
||||
|
||||
|
|
Loading…
Reference in New Issue