Update app.py

app.py

@@ -1,25 +1,45 @@
 from transformers import pipeline
 import gradio as gr
 
-generator = pipeline("text-generation", model="ckiplab/gpt2-base-chinese", 
-                     tokenizer="ckiplab/gpt2-base-chinese")
+# 使用中文 GPT2 對話模型（支援 CPU）
+generator = pipeline(
+    "text-generation",
+    model="thu-coai/CDial-GPT2_LCCC-base",
+    tokenizer="thu-coai/CDial-GPT2_LCCC-base",
+    device=-1  # 使用 CPU
+)
 
+# 對話處理函式
 def chat_fn(message, history):
     history = history or []
-    input_text = "\n".join(history + [f"你: {message}", "AI:"])
-    output = generator(input_text, max_new_tokens=80, pad_token_id=0)[0]["generated_text"]
-    response = output.split("AI:")[-1].strip().split("你:")[0].strip()
-    history.append(f"你: {message}")
-    history.append(f"AI: {response}")
-    messages = [(history[i], history[i+1]) for i in range(0, len(history)-1, 2)]
-    return messages, history
+    
+    # 將所有歷史訊息合併為 prompt
+    prompt = ""
+    for user_msg, bot_msg in history:
+        prompt += f"你說：{user_msg}\nAI說：{bot_msg}\n"
+    prompt += f"你說：{message}\nAI說："
 
+    # 生成新回應
+    output = generator(prompt, max_new_tokens=80, pad_token_id=0)[0]["generated_text"]
+
+    # 從模型輸出中擷取 AI 回覆
+    response = output.split("AI說：")[-1].split("你說：")[-1].strip()
+
+    # 更新歷史
+    history.append((message, response))
+    return history, history
+
+# 建立 Gradio 介面
 with gr.Blocks() as demo:
-    chatbot = gr.Chatbot(label="中文聊天機器人", type="tuples")
-    state = gr.State([])
-    textbox = gr.Textbox(placeholder="請輸入訊息")
+    gr.Markdown("## 🧠 中文聊天機器人（記住上下文）")
+
+    chatbot = gr.Chatbot(label="GPT2 中文對話")
+    msg = gr.Textbox(show_label=False, placeholder="請輸入訊息，Enter 送出")
+    clear = gr.Button("🧹 清除對話")
+
+    state = gr.State([])  # 儲存對話歷史
 
-    textbox.submit(chat_fn, [textbox, state], [chatbot, state])
-    textbox.submit(lambda: "", None, textbox)
+    msg.submit(chat_fn, inputs=[msg, state], outputs=[chatbot, state])
+    clear.click(lambda: ([], []), outputs=[chatbot, state])
 
 demo.launch()