mirror of https://github.com/THUDM/ChatGLM-6B
57 lines
2.2 KiB
Python
57 lines
2.2 KiB
Python
from transformers import AutoModel, AutoTokenizer
|
|
import gradio as gr
|
|
|
|
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()
|
|
|
|
MAX_TURNS = 20
|
|
MAX_BOXES = MAX_TURNS * 2
|
|
|
|
|
|
def btn_is_clickable(txt):
|
|
if txt is not None and txt.strip() != '':
|
|
return gr.update(interactive=True)
|
|
else:
|
|
return gr.update(interactive=False)
|
|
|
|
|
|
def user_message(user_message, history):
|
|
return "", history + [[user_message, None]]
|
|
|
|
|
|
def predict(input, max_length, top_p, temperature, history):
|
|
if len(history) > MAX_TURNS:
|
|
history = history[-20:]
|
|
for response, const_history in model.stream_chat(tokenizer, input, history, max_length=max_length, top_p=top_p,
|
|
temperature=temperature):
|
|
history[-1][1] = response
|
|
yield history
|
|
|
|
|
|
with gr.Blocks() as demo:
|
|
with gr.Row():
|
|
with gr.Column(scale=1):
|
|
gr.Markdown("https://github.com/THUDM/ChatGLM-6B")
|
|
max_length = gr.Slider(0, 4096, value=2048, step=1.0, label="Maximum length", interactive=True)
|
|
top_p = gr.Slider(0, 1, value=0.7, step=0.01, label="Top P", interactive=True)
|
|
temperature = gr.Slider(0, 1, value=0.95, step=0.01, label="Temperature", interactive=True)
|
|
with gr.Column(scale=4):
|
|
chatbot = gr.Chatbot([], elem_id="chatbot").style(height=750)
|
|
with gr.Row():
|
|
with gr.Column(scale=4):
|
|
txt = gr.Textbox(
|
|
show_label=False,
|
|
placeholder="有问题就会有答案"
|
|
).style(container=False)
|
|
with gr.Column(scale=1, min_width=0):
|
|
btn = gr.Button("发送", interactive=False)
|
|
# 控制按钮是否可以点击
|
|
txt.change(btn_is_clickable, txt, btn)
|
|
# 发送消息
|
|
btn.click(user_message, [txt, chatbot], [txt, chatbot], queue=False).then(
|
|
predict, [txt, max_length, top_p, temperature, chatbot], chatbot
|
|
)
|
|
|
|
demo.queue().launch(share=False, inbrowser=True)
|