support chat

app.py

@@ -1,11 +1,43 @@
 from peft import PeftModel, PeftConfig
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers_stream_generator.main import NewGenerationMixin, StreamGenerationConfig
 import gradio as gr
+import torch
+
+config = PeftConfig.from_pretrained("Junity/Genshin-World-Model")
+model = AutoModelForCausalLM.from_pretrained("baichuan-inc/Baichuan-13B-Base")
+model = PeftModel.from_pretrained(model, "Junity/Genshin-World-Model")
+tokenizer = AutoTokenizer.from_pretrained("Junity/Genshin-World-Model")
+
+history = []
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def respond(role_name, msg, chatbot, character):
+    global history
+    if role_name is not None:
+        history.append(role_name + "：" + msg)
+    else:
+        history.append(msg)
+    total_input = []
+    for i, message in enumerate(history[::-1]):
+        content_tokens = tokenizer.encode(message + '\n')
+        total_input = content_tokens + total_input
+        if content_tokens + total_input > 4096:
+            break
+    total_input = total_input[-4096:]
+    input_ids = torch.LongTensor([total_input]).to(device)
+    generation_config = model.generation_config
+    stream_config = StreamGenerationConfig(**generation_config.to_dict(), do_stream=True)
+
+    def stream_generator():
+        outputs = []
+        for token in model.generate(input_ids, generation_config=stream_config):
+            outputs.append(token.item())
+            yield None, tokenizer.decode(outputs, skip_special_tokens=True)
+
+    return stream_generator()
 
-# config = PeftConfig.from_pretrained("Junity/Genshin-World-Model")
-# model = AutoModelForCausalLM.from_pretrained("baichuan-inc/Baichuan-13B-Base")
-# model = PeftModel.from_pretrained(model, "Junity/Genshin-World-Model")
-# tokenizer = AutoTokenizer.from_pretrained("Junity/Genshin-World-Model")
 
 with gr.Blocks() as demo:
     gr.Markdown(
@@ -15,4 +47,13 @@ with gr.Blocks() as demo:
         - 此模型不支持要求对方回答什么，只支持续写。
         """
     )
-    demo.queue().launch()
+    with gr.Tab("聊天") as chat:
+        role_name = gr.Textbox(label="你将扮演的角色")
+        msg = gr.Textbox(label="输入")
+    with gr.Row():
+        clear = gr.Button("Clear")
+        sub = gr.Button("Submit")
+    chatbot = gr.Chatbot()
+    sub.click(fn=respond, inputs=[role_name, msg, chatbot], outputs=[msg, chatbot])
+    clear.click(lambda: None, None, chatbot, queue=False)
+    demo.queue().launch()