替换后

app.py

@@ -1,55 +1,64 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-# 加载模型和分词器（选择 Qwen-1.8B-Chat 适配你的配置）
-model_name = "Qwen/Qwen-1.8B-Chat"
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+# 加载模型和分词器（适配 Qwen1.5-1.8B-Chat 并保留优化配置）
+model_name = "Qwen/Qwen1.5-1.8B-Chat"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
-    trust_remote_code=True,
-    torch_dtype=torch.float16,  # 用半精度减少内存占用
+    torch_dtype=torch.float16,  # 半精度减少内存占用
     device_map="auto",
-    load_in_4bit=True,  # 4-bit 量化，进一步降低内存压力
+    load_in_4bit=True,  # 4-bit 量化降低内存压力
     bnb_4bit_compute_dtype=torch.float16
 )
 
-# 创建优化后的文本生成管道
-chat_pipeline = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=150,  # 平衡回复长度和速度
-    temperature=0.7,
-    do_sample=True,
-    num_return_sequences=1,
-    repetition_penalty=1.1  # 减少重复回复
-)
-
-# 聊天函数（优化历史拼接逻辑，减少计算量）
+# 优化后的聊天函数（适配 Qwen 的对话模板）
 def chat_with_model(message, history):
-    # 精简对话历史，只保留最近 3 轮，降低内存占用
+    # 只保留最近 3 轮历史，减少计算量
     history = history[-3:]
-    prompt = "用户：你好，我是你的助手。\n"
+    messages = []
+    # 拼接历史对话
     for user_msg, bot_msg in history:
-        prompt += f"用户：{user_msg}\n助手：{bot_msg}\n"
-    prompt += f"用户：{message}\n助手："
+        messages.append({"role": "user", "content": user_msg})
+        messages.append({"role": "assistant", "content": bot_msg})
+    # 加入当前用户消息
+    messages.append({"role": "user", "content": message})
+    
+    # 生成模型输入
+    inputs = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_dict=True,
+        return_tensors="pt",
+    ).to(model.device)
     
-    # 推理时限制 batch 大小，适配 2 核 CPU
+    # 推理生成回复
     with torch.no_grad():
-        response = chat_pipeline(prompt, batch_size=1)[0]["generated_text"]
-    bot_response = response.split("助手：")[-1].strip()
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=150,
+            temperature=0.7,
+            repetition_penalty=1.1,
+            do_sample=True
+        )
+    # 解码并提取回复
+    bot_response = tokenizer.decode(
+        outputs[0][inputs["input_ids"].shape[-1]:],
+        skip_special_tokens=True
+    ).strip()
     return bot_response
 
-# 启动Gradio界面（关闭多余组件，减少资源占用）
+# 启动 Gradio 界面（保留资源优化配置）
 if __name__ == "__main__":
     gr.ChatInterface(
         fn=chat_with_model,
         title="轻量聊天助手",
-        description="基于 Qwen-1.8B-Chat 适配 2 核 16G 配置"
+        description="基于 Qwen1.5-1.8B-Chat 适配 2 核 16G 配置"
     ).launch(
         server_name="0.0.0.0",
         server_port=7860,
-        share=False,  # 关闭分享功能，减少资源消耗
+        share=False,
         inline=False
     )