Update app.py

app.py

@@ -3,7 +3,7 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import spaces
 
-
+# 加载模型和分词器
 model_name = "Zhihu-ai/Zhi-writing-dsr1-14b"
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
@@ -12,10 +12,9 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map="auto",
     trust_remote_code=True
 )
-
 @spaces.GPU()
 def predict(message, history):
-    
+    # 构建输入
     history_text = ""
     for human, assistant in history:
         history_text += f"Human: {human}\nAssistant: {assistant}\n"
@@ -23,40 +22,29 @@ def predict(message, history):
     
     # 生成回复
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    
-    # 使用流式生成
-    for response in model.generate(
+    outputs = model.generate(
         **inputs,
         max_new_tokens=10000,
         do_sample=True,
         temperature=0.7,
         top_p=0.9,
         repetition_penalty=1.1,
-        pad_token_id=tokenizer.eos_token_id,
-        streamer=gr.TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
-    ):
-        yield response.strip()
-
-
-
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# Zhi-writing-dsr1-14")
-    gr.Markdown("这是一个基于Zhi-writing-dsr1-14的文章生成器")
+        pad_token_id=tokenizer.eos_token_id
+    )
+    response = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
     
-    chatbot = gr.Chatbot()
-    msg = gr.Textbox(label="输入消息")
-    clear = gr.Button("清除对话")
-    
-    def respond(message, chat_history):
-        bot_message = ""
-        for response in predict(message, chat_history):
-            bot_message = response
-            chat_history.append((message, bot_message))
-            yield chat_history
-        return "", chat_history
-    
-    msg.submit(respond, [msg, chatbot], [msg, chatbot])
-    clear.click(lambda: None, None, chatbot, queue=False)
+    return response.strip()
+
+# 创建Gradio界面
+demo = gr.ChatInterface(
+    predict,
+    title="测试Zhi-writing-dsr1-14b",
+    description="Zhihu-ai/Zhi-writing-dsr1-14b",
+    examples=["鲁迅口吻写五百字，描述桔猫的可爱！", "桔了个仔是谁", "介绍自己"],
+    theme=gr.themes.Soft()
+)
 
 if __name__ == "__main__":
     demo.launch(share=True)
+
+