Update app.py

app.py

@@ -21,22 +21,29 @@ model = AutoModelForCausalLM.from_pretrained(
 model.eval()
 
 # ===============================
-# 核心聊天逻辑（使用 tuple history）
+# 聊天核心（messages schema）
 # ===============================
 def chat_fn(message, history):
     """
-    history: List[Tuple[user, assistant]]
+    history: List[Dict] with keys: role, content
     """
-    # 只保留最近 3 轮
-    history = history[-3:]
+
+    # 只保留最近 3 轮（6 条 message）
+    history = history[-6:]
 
     prompt = ""
-    for user, assistant in history:
-        prompt += f"<|user|>{user}<|assistant|>{assistant}"
+    for msg in history:
+        role = msg["role"]
+        content = msg["content"]
+        prompt += f"<|{role}|>{content}"
 
+    # 当前用户问题
     prompt += f"<|user|>{message}<|assistant|>"
 
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+    inputs = tokenizer(
+        prompt,
+        return_tensors="pt"
+    ).to(model.device)
 
     with torch.no_grad():
         output_ids = model.generate(
@@ -53,7 +60,7 @@ def chat_fn(message, history):
         skip_special_tokens=True
     )
 
-    # 只取 assistant 的新回答
+    # 只取 assistant 新生成的部分
     if "<|assistant|>" in output_text:
         output_text = output_text.split("<|assistant|>")[-1]
 
@@ -61,20 +68,36 @@ def chat_fn(message, history):
 
 
 # ===============================
-# Gradio UI
+# Gradio UI（messages 模式）
 # ===============================
 with gr.Blocks(title="caobin LLM Chatbot") as demo:
     gr.Markdown("# 🤖 caobin's AI Assistant")
 
-    chatbot = gr.Chatbot(height=450)
+    chatbot = gr.Chatbot(
+        type="messages",
+        height=450
+    )
+
     msg = gr.Textbox(
         label="输入你的问题",
         placeholder="请输入你的问题，支持多轮对话"
     )
 
     def respond(message, chat_history):
+        # 用户消息
+        chat_history.append({
+            "role": "user",
+            "content": message
+        })
+
+        # 模型回复
         response = chat_fn(message, chat_history)
-        chat_history.append((message, response)) 
+
+        chat_history.append({
+            "role": "assistant",
+            "content": response
+        })
+
         return "", chat_history
 
     msg.submit(