Update app.py

app.py

@@ -1,31 +1,27 @@
-# app.py
-
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 import gradio as gr
 
-model_name = "tosei0000/tosei"  # 替换为你的模型路径或名称
+model_name = "tosei0000/chatbot"
 
-# 加载 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",          # 依赖 accelerate
     trust_remote_code=True
 )
 
-# 明确设置 pad_token_id
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+
 tokenizer.pad_token_id = tokenizer.eos_token_id
 model.config.pad_token_id = tokenizer.eos_token_id
 
-# 多轮对话函数
 def chat(user_input, history):
-    # 拼接历史
     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)
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
     output = model.generate(
         **inputs,
         max_new_tokens=256,
@@ -36,12 +32,10 @@ def chat(user_input, history):
         eos_token_id=tokenizer.eos_token_id
     )
     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 Chatbot") as demo:
     gr.Markdown("## 🤖 Qwen2 聊天机器人")
     chatbot = gr.Chatbot()
@@ -56,6 +50,7 @@ if __name__ == "__main__":
     demo.launch()
 
 
+
 # from transformers import AutoTokenizer, AutoModelForCausalLM
 # import torch