Spaces:
Build error
Build error
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 = AutoModel.from_pretrained("THUDM/chatglm-6b-int4",trust_remote_code=True).float() | |
model = model.eval() | |
MAX_TURNS = 20 | |
MAX_BOXES = MAX_TURNS * 2 | |
def predict(input, max_length, top_p, temperature, history=None): | |
if history is None: | |
history = [] | |
for response, history in model.stream_chat(tokenizer, input, history, max_length=max_length, top_p=top_p, | |
temperature=temperature): | |
updates = [] | |
for query, response in history: | |
updates.append(gr.update(visible=True, value="用户:" + query)) | |
updates.append(gr.update(visible=True, value="ChatGLM-6B:" + response)) | |
if len(updates) < MAX_BOXES: | |
updates = updates + [gr.Textbox.update(visible=False)] * (MAX_BOXES - len(updates)) | |
yield [history] + updates | |
with gr.Blocks() as demo: | |
state = gr.State([]) | |
text_boxes = [] | |
for i in range(MAX_BOXES): | |
if i % 2 == 0: | |
text_boxes.append(gr.Markdown(visible=False, label="提问:")) | |
else: | |
text_boxes.append(gr.Markdown(visible=False, label="回复:")) | |
with gr.Row(): | |
with gr.Column(scale=4): | |
txt = gr.Textbox(show_label=False, placeholder="Enter text and press enter", lines=11).style( | |
container=False) | |
with gr.Column(scale=1): | |
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) | |
button = gr.Button("Generate") | |
button.click(predict, [txt, max_length, top_p, temperature, state], [state] + text_boxes) | |
demo.queue().launch() | |