Update app.py

app.py

@@ -1,10 +1,24 @@
 from transformers import pipeline
 import gradio as gr
 
-generator = pipeline("text-generation", model="distilgpt2")
+# 中文 GPT2 模型
+generator = pipeline("text-generation", model="IDEA-CCNL/Wenzhong2.0-GPT2-110M")
 
-def chat_fn(message, history):
-    output = generator(message, max_new_tokens=50)[0]["generated_text"]
-    return output, history
+def chat_fn(user_input, history):
+    history = history or []
 
-gr.Interface(fn=chat_fn, inputs="text", outputs="text").launch()
+    # 構造中文 prompt
+    prompt = ""
+    for user_msg, bot_msg in history:
+        prompt += f"用戶：{user_msg}\n機器人：{bot_msg}\n"
+    prompt += f"用戶：{user_input}\n機器人："
+
+    # 生成中文回答
+    output = generator(prompt, max_new_tokens=100, pad_token_id=50256)[0]["generated_text"]
+    reply = output.split("機器人：")[-1].strip()
+
+    # 更新歷史
+    history.append((user_input, reply))
+    return history, history
+
+gr.ChatInterface(chat_fn).launch()