Update app.py

app.py

@@ -1,24 +1,29 @@
 from transformers import pipeline
 import gradio as gr
 
-# 中文 GPT2 模型
-generator = pipeline("text-generation", model="IDEA-CCNL/Wenzhong2.0-GPT2-110M")
+# 使用開放的中文 GPT-2 模型
+generator = pipeline("text-generation", model="ckiplab/gpt2-base-chinese", 
+                     tokenizer="ckiplab/gpt2-base-chinese")
 
-def chat_fn(user_input, history):
+def chat_fn(message, history):
     history = history or []
+    input_text = "\n".join(history + [f"你: {message}", "AI:"])
     
-    # 建立 prompt
-    prompt = ""
-    for user_msg, bot_msg in history:
-        prompt += f"用戶：{user_msg}\n機器人：{bot_msg}\n"
-    prompt += f"用戶：{user_input}\n機器人："
+    output = generator(input_text, max_new_tokens=80, pad_token_id=0)[0]["generated_text"]
+    # 取出 AI 回應部分
+    response = output.split("AI:")[-1].strip().split("你:")[0].strip()
+    
+    history.append(f"你: {message}")
+    history.append(f"AI: {response}")
+    
+    return history, history
 
-    # 模型生成
-    output = generator(prompt, max_new_tokens=100, pad_token_id=50256)[0]["generated_text"]
-    reply = output.split("機器人：")[-1].strip()
+with gr.Blocks() as demo:
+    chatbot = gr.Chatbot()
+    state = gr.State([])
+    textbox = gr.Textbox(show_label=False, placeholder="請輸入訊息").style(container=False)
 
-    history.append((user_input, reply))
-    return history, history
+    textbox.submit(chat_fn, [textbox, state], [chatbot, state])
+    textbox.submit(lambda: "", None, textbox)  # 清空輸入框
 
-# ✅ 用 gr.ChatInterface() 就會自動長得像聊天室
-gr.ChatInterface(chat_fn).launch()
+demo.launch()