Update app.py

app.py

@@ -1,70 +1,84 @@
+
+
+
+
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
-import threading
 
 MODEL_ID = "caobin/llm-caobin"
 
 # 加载 tokenizer 和模型
 tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+
+# 根据是否有 GPU 自动设置 dtype
+dtype = torch.float16 if torch.cuda.is_available() else torch.float32
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.float16,
+    torch_dtype=dtype,
     trust_remote_code=True
 )
-
-# 判断是否有 GPU
-device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 model.eval()
 
-# 边生成边输出的函数
-def generate_stream(prompt, max_new_tokens=512, temperature=0.7, top_p=0.9, max_history=3, history=[]):
-    # 只保留最近 max_history 轮对话
-    recent_history = history[-max_history:]
+MAX_HISTORY = 3  # 只保留最近几轮对话
+
+def chat_fn(message, history):
+    """
+    message: 用户最新输入
+    history: [{"role": "user"/"assistant", "content": str}, ...]
+    """
+    # 只保留最近 MAX_HISTORY 轮
+    recent_history = history[-MAX_HISTORY*2:]  # user+assistant = 2 条消息一轮
+
+    # 拼接 prompt
     full_prompt = ""
-    for user_msg, bot_msg in recent_history:
-        full_prompt += f"<|user|>{user_msg}<|assistant|>{bot_msg}"
-    full_prompt += f"<|user|>{prompt}<|assistant|>"
+    for msg in recent_history:
+        if msg["role"] == "user":
+            full_prompt += f"<|user|>{msg['content']}"
+        elif msg["role"] == "assistant":
+            full_prompt += f"<|assistant|>{msg['content']}"
+    full_prompt += f"<|user|>{message}<|assistant|>"
 
     inputs = tokenizer(full_prompt, return_tensors="pt").to(device)
 
-    # 使用流式输出
-    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
-    generate_kwargs = dict(
+    # 生成回复
+    output_ids = model.generate(
         **inputs,
-        streamer=streamer,
-        max_new_tokens=max_new_tokens,
-        temperature=temperature,
-        top_p=top_p,
+        max_new_tokens=512,
+        temperature=0.7,
+        top_p=0.9,
         do_sample=True,
+        pad_token_id=tokenizer.eos_token_id
     )
 
-    thread = threading.Thread(target=model.generate, kwargs=generate_kwargs)
-    thread.start()
+    # 只 decode 新生成部分
+    generated_text = tokenizer.decode(output_ids[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True)
 
-    # 边生成边返回文本
-    output_text = ""
-    for new_text in streamer:
-        output_text += new_text
-        yield output_text.strip()
+    return generated_text.strip()
 
-# Gradio 回调函数
 def respond(message, chat_history):
-    # 返回一个生成器，用于流式更新
-    generator = generate_stream(message, history=chat_history)
-    bot_response = ""
-    for partial in generator:
-        bot_response = partial
-        yield "", chat_history + [(message, bot_response)]
-
-# 创建 Gradio 界面
+    # chat_history 是 Gradio 最新格式 [{"role":..., "content":...}, ...]
+    response = chat_fn(message, chat_history)
+    # 更新聊天历史
+    new_history = chat_history + [
+        {"role": "user", "content": message},
+        {"role": "assistant", "content": response}
+    ]
+    return "", new_history
+
+# Gradio 界面
 with gr.Blocks(title="caobin LLM Chatbot") as demo:
-    gr.Markdown("# 🤖 caobin‘s AI assistant")
+    gr.Markdown("# 🤖 caobin's AI assistant")
 
-    chatbot = gr.Chatbot(height=450)
+    chatbot = gr.Chatbot([], height=450)  # 初始化为空列表
     msg = gr.Textbox(label="输入你的问题")
 
     msg.submit(respond, [msg, chatbot], [msg, chatbot])
 
 demo.launch()
+
+
+