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

support stream chat

pull/192/head
AdamBear 2023-03-22 15:11:21 +08:00
parent 905aa26b91
commit ef6b3ff35e
1 changed files with 17 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

23
web_demo2.py
View File

@ -11,8 +11,8 @@ st.set_page_config(
@st.cache_resource @st.cache_resource
def get_model(): def get_model():
tokenizer = AutoTokenizer.from_pretrained("/THUDM/chatglm-6b", trust_remote_code=True) tokenizer = AutoTokenizer.from_pretrained("/data/chatglm-6b", trust_remote_code=True)
model = AutoModel.from_pretrained("THUDM/chatglm-6b", trust_remote_code=True).half().cuda() model = AutoModel.from_pretrained("/data/chatglm-6b", trust_remote_code=True).half().cuda()
model = model.eval() model = model.eval()
return tokenizer, model return tokenizer, model
@ -25,26 +25,31 @@ def predict(input, history=None):
tokenizer, model = get_model() tokenizer, model = get_model()
if history is None: if history is None:
history = [] history = []
response, history = model.chat(tokenizer, input, history)
#updates = [] with container:
if len(history) > 0:
for i, (query, response) in enumerate(history): for i, (query, response) in enumerate(history):
#updates.append("用户：" + query)
message(query, avatar_style="big-smile", key=str(i) + "_user") message(query, avatar_style="big-smile", key=str(i) + "_user")
#updates.append("ChatGLM-6B：" + response)
message(response, avatar_style="bottts", key=str(i)) message(response, avatar_style="bottts", key=str(i))
# if len(updates) < MAX_BOXES: message(input, avatar_style="big-smile", key=str(len(history)) + "_user")
# updates = updates + [""] * (MAX_BOXES - len(updates)) st.write("AI正在回复:")
with st.empty():
for response, history in model.stream_chat(tokenizer, input, history):
query, response = history[-1]
st.write(response)
return history return history
container = st.container()
# create a prompt text for the text generation # create a prompt text for the text generation
prompt_text = st.text_area(label="用户命令输入", prompt_text = st.text_area(label="用户命令输入",
height = 100, height = 100,
placeholder="请在这儿输入您的命令") placeholder="请在这儿输入您的命令")
if 'state' not in st.session_state: if 'state' not in st.session_state:
st.session_state['state'] = [] st.session_state['state'] = []
@ -53,4 +58,4 @@ if st.button("发送", key="predict"):
# text generation # text generation
st.session_state["state"] = predict(prompt_text, st.session_state["state"]) st.session_state["state"] = predict(prompt_text, st.session_state["state"])
st.balloons() st.session_state["state"]