Update app.py

app.py

@@ -2,9 +2,6 @@ import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# ===============================
-# 模型加载
-# ===============================
 MODEL_ID = "caobin/llm-caobin"
 
 tokenizer = AutoTokenizer.from_pretrained(
@@ -20,30 +17,19 @@ model = AutoModelForCausalLM.from_pretrained(
 
 model.eval()
 
-# ===============================
-# 聊天核心（messages schema）
-# ===============================
 def chat_fn(message, history):
     """
-    history: List[Dict] with keys: role, content
+    history: List[Tuple[user, assistant]]
     """
-
-    # 只保留最近 3 轮（6 条 message）
-    history = history[-6:]
+    history = history[-3:]
 
     prompt = ""
-    for msg in history:
-        role = msg["role"]
-        content = msg["content"]
-        prompt += f"<|{role}|>{content}"
+    for user, assistant in history:
+        prompt += f"<|user|>{user}<|assistant|>{assistant}"
 
-    # 当前用户问题
     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(
@@ -60,53 +46,23 @@ def chat_fn(message, history):
         skip_special_tokens=True
     )
 
-    # 只取 assistant 新生成的部分
     if "<|assistant|>" in output_text:
         output_text = output_text.split("<|assistant|>")[-1]
 
     return output_text.strip()
 
 
-# ===============================
-# Gradio UI（messages 模式）
-# ===============================
 with gr.Blocks(title="caobin LLM Chatbot") as demo:
     gr.Markdown("# 🤖 caobin's AI Assistant")
 
-    chatbot = gr.Chatbot(
-        type="messages",
-        height=450
-    )
-
-    msg = gr.Textbox(
-        label="输入你的问题",
-        placeholder="请输入你的问题，支持多轮对话"
-    )
+    chatbot = gr.Chatbot(height=450)
+    msg = gr.Textbox(label="输入你的问题")
 
     def respond(message, chat_history):
-        # 用户消息
-        chat_history.append({
-            "role": "user",
-            "content": message
-        })
-
-        # 模型回复
         response = chat_fn(message, chat_history)
-
-        chat_history.append({
-            "role": "assistant",
-            "content": response
-        })
-
+        chat_history.append((message, response))
         return "", chat_history
 
-    msg.submit(
-        respond,
-        inputs=[msg, chatbot],
-        outputs=[msg, chatbot]
-    )
+    msg.submit(respond, [msg, chatbot], [msg, chatbot])
 
-# ===============================
-# 启动
-# ===============================
 demo.launch()