Update app.py

app.py

@@ -1,33 +1,31 @@
+# app.py
+
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import gradio as gr
 
-# 模型名称（可以换成你自己的Qwen2模型）
-model_name = "tosei0000/chatbot"
+model_name = "Qwen/Qwen2-1.5B"  # 替换为你的模型路径或名称
 
 # 加载 tokenizer 和 model
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.bfloat16, device_map="auto", trust_remote_code=True)
-
-# 设置 pad_token_id（避免警告和生成错误）
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",          # 依赖 accelerate
+    trust_remote_code=True
+)
+
+# 明确设置 pad_token_id
 tokenizer.pad_token_id = tokenizer.eos_token_id
 model.config.pad_token_id = tokenizer.eos_token_id
 
-# 聊天历史存储
-chat_history = []
-
-# 多轮对话生成函数
+# 多轮对话函数
 def chat(user_input, history):
-    # 构造 prompt（把历史拼接起来）
-    prompt = ""
-    for i, (user_msg, bot_msg) in enumerate(history):
-        prompt += f"User: {user_msg}\nAssistant: {bot_msg}\n"
-    prompt += f"User: {user_input}\nAssistant:"
-
-    # 编码输入
+    # 拼接历史
+    prompt = "".join(
+        f"User: {u}\nAssistant: {a}\n" for u, a in history
+    ) + f"User: {user_input}\nAssistant:"
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-
-    # 生成
     output = model.generate(
         **inputs,
         max_new_tokens=256,
@@ -37,30 +35,23 @@ def chat(user_input, history):
         pad_token_id=tokenizer.pad_token_id,
         eos_token_id=tokenizer.eos_token_id
     )
-
-    # 解码
-    decoded = tokenizer.decode(output[0], skip_special_tokens=True)
-
-    # 提取模型最新回复部分（去掉前面的prompt）
-    response = decoded[len(prompt):].strip().split("\n")[0]
-
-    # 更新历史
-    history.append((user_input, response))
+    text = tokenizer.decode(output[0], skip_special_tokens=True)
+    # 提取回复
+    reply = text[len(prompt):].strip().split("\n")[0]
+    history.append((user_input, reply))
     return history, history
 
-# 创建 Gradio 接口
-with gr.Blocks(title="Qwen2 聊天机器人") as demo:
-    gr.Markdown("## 🤖 Qwen2 Chatbot")
+# Gradio 界面
+with gr.Blocks(title="Qwen2 Chatbot") as demo:
+    gr.Markdown("## 🤖 Qwen2 聊天机器人")
     chatbot = gr.Chatbot()
     msg = gr.Textbox(label="输入你的问题")
     clear = gr.Button("清除对话")
-
-    state = gr.State([])  # 存储历史
+    state = gr.State([])
 
     msg.submit(chat, [msg, state], [chatbot, state])
     clear.click(lambda: ([], []), None, [chatbot, state])
 
-# 启动 Gradio
 if __name__ == "__main__":
     demo.launch()