Update app.py

app.py

@@ -1,18 +1,85 @@
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+import gradio as gr
 
-model_path = "tosei0000/chatbot"
+# 模型名称（可以换成你自己的Qwen2模型）
+model_name = "tosei0000/chatbot"
 
-tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
+# 加载 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)
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
+# 设置 pad_token_id（避免警告和生成错误）
+tokenizer.pad_token_id = tokenizer.eos_token_id
+model.config.pad_token_id = tokenizer.eos_token_id
 
-def chat(prompt, max_new_tokens=100):
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
-    outputs = model.generate(**inputs, max_new_tokens=max_new_tokens)
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+# 聊天历史存储
+chat_history = []
 
-response = chat("こんにちは！")
-print(response)
+# 多轮对话生成函数
+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:"
+
+    # 编码输入
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+    # 生成
+    output = model.generate(
+        **inputs,
+        max_new_tokens=256,
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9,
+        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))
+    return history, history
+
+# 创建 Gradio 接口
+with gr.Blocks(title="Qwen2 聊天机器人") as demo:
+    gr.Markdown("## 🤖 Qwen2 Chatbot")
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox(label="输入你的问题")
+    clear = gr.Button("清除对话")
+
+    state = gr.State([])  # 存储历史
+
+    msg.submit(chat, [msg, state], [chatbot, state])
+    clear.click(lambda: ([], []), None, [chatbot, state])
+
+# 启动 Gradio
+if __name__ == "__main__":
+    demo.launch()
+
+
+# from transformers import AutoTokenizer, AutoModelForCausalLM
+# import torch
+
+# model_path = "tosei0000/chatbot"
+
+# tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+# model = AutoModelForCausalLM.from_pretrained(model_path, trust_remote_code=True)
+
+# device = "cuda" if torch.cuda.is_available() else "cpu"
+# model.to(device)
+
+# def chat(prompt, max_new_tokens=100):
+#     inputs = tokenizer(prompt, return_tensors="pt").to(device)
+#     outputs = model.generate(**inputs, max_new_tokens=max_new_tokens)
+#     return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+# response = chat("こんにちは！")
+# print(response)