Update app.py

app.py

@@ -2,44 +2,39 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
-# 模型 ID
 MODEL_ID = "caobin/llm-caobin"
 
-# 加载 tokenizer 和模型
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    device_map="auto",  # CPU 上会自动映射到 CPU
+    device_map="auto",  # CPU 上也可以
     trust_remote_code=True
 )
 
-# 聊天函数
 def chat_fn(message, history):
     # 只保留最近 3 轮历史
-    history = history[-3:]
+    recent_history = history[-6:]  # 3轮对话，每轮2条消息
     full_prompt = ""
-    for user_msg, bot_msg in history:
-        full_prompt += f"<|user|>{user_msg}<|assistant|>{bot_msg}"
+    for msg in recent_history:
+        if msg["role"] == "user":
+            full_prompt += f"<|user|>{msg['content']}<|assistant|>"
+        elif msg["role"] == "assistant":
+            full_prompt += msg['content']
     full_prompt += f"<|user|>{message}<|assistant|>"
 
-    # tokenizer 转 tensor
     inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
-
-    # 生成回答
     output_ids = model.generate(
         **inputs,
-        max_new_tokens=256,   
+        max_new_tokens=256,
         temperature=0.7,
         top_p=0.9,
         do_sample=True,
     )
-
     output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
     if "<|assistant|>" in output_text:
         output_text = output_text.split("<|assistant|>")[-1]
     return output_text.strip()
 
-# Gradio UI
 with gr.Blocks(title="caobin LLM Chatbot") as demo:
     gr.Markdown("# 🤖 caobin's AI assistant")
     chatbot = gr.Chatbot(height=450)
@@ -47,7 +42,9 @@ with gr.Blocks(title="caobin LLM Chatbot") as demo:
 
     def respond(message, chat_history):
         response = chat_fn(message, chat_history)
-        chat_history.append((message, response))
+        # 用字典格式添加消息
+        chat_history.append({"role": "user", "content": message})
+        chat_history.append({"role": "assistant", "content": response})
         return "", chat_history
 
     msg.submit(respond, [msg, chatbot], [msg, chatbot])
@@ -57,3 +54,4 @@ demo.launch()
 
 
 
+