Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-import json
 
 # 全局变量，避免重复加载
 model = None
@@ -10,30 +9,57 @@ tokenizer = None
 def load_model():
     """加载模型和分词器"""
     global model, tokenizer
-    model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"  # 建议使用蒸馏版节省资源
+    model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
     try:
         tokenizer = AutoTokenizer.from_pretrained(model_name)
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
-            torch_dtype=torch.float16,  # 使用半精度减少内存占用
-            device_map="auto",           # 自动分配设备
-            low_cpu_mem_usage=True       # 优化CPU内存使用
+            torch_dtype=torch.float16,
+            device_map="auto",
+            low_cpu_mem_usage=True
         )
+        # 确保tokenizer有pad_token
+        if tokenizer.pad_token is None:
+            tokenizer.pad_token = tokenizer.eos_token
         print("模型加载成功！")
     except Exception as e:
         print(f"模型加载失败: {e}")
 
-def openai_compatible_api(message, history):
-    """处理OpenAI格式的请求"""
+def chat_with_deepseek(message, history):
+    """与DeepSeek模型聊天 - 修正版"""
+    global model, tokenizer
+    
     if model is None:
         load_model()
-
-    # 构建符合DeepSeek模型要求的对话格式
-    # 注意：请根据您使用的具体模型调整提示词模板
-    prompt = f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
     
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)
-
+    # 构建对话历史
+    conversation = []
+    for user_msg, assistant_msg in history:
+        conversation.append({"role": "user", "content": user_msg})
+        conversation.append({"role": "assistant", "content": assistant_msg})
+    conversation.append({"role": "user", "content": message})
+    
+    # 使用tokenizer的apply_chat_template方法（如果支持）
+    try:
+        prompt = tokenizer.apply_chat_template(
+            conversation, 
+            tokenize=False, 
+            add_generation_prompt=True
+        )
+    except:
+        # 如果不支持apply_chat_template，使用简单格式
+        prompt = ""
+        for msg in conversation:
+            if msg["role"] == "user":
+                prompt += f"<|im_start|>user\n{msg['content']}<|im_end|>\n"
+            else:
+                prompt += f"<|im_start|>assistant\n{msg['content']}<|im_end|>\n"
+        prompt += "<|im_start|>assistant\n"
+    
+    # 编码输入
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048)
+    
+    # 生成回复
     with torch.no_grad():
         outputs = model.generate(
             **inputs,
@@ -42,34 +68,49 @@ def openai_compatible_api(message, history):
             top_p=0.9,
             do_sample=True,
             pad_token_id=tokenizer.eos_token_id,
-            eos_token_id=tokenizer.eos_token_id
+            eos_token_id=tokenizer.eos_token_id,
+            repetition_penalty=1.1
         )
+    
+    # 解码回复
+    response = tokenizer.decode(outputs[0], skip_special_tokens=False)
+    
+    # 关键修正：提取助理的回复部分
+    if "<|im_start|>assistant" in response:
+        # 找到最后一个assistant标记开始的位置
+        assistant_start = response.rfind("<|im_start|>assistant")
+        if assistant_start != -1:
+            assistant_content = response[assistant_start:]
+            # 提取assistant标记后的内容
+            if "\n" in assistant_content:
+                content_start = assistant_content.find("\n") + 1
+                generated_text = assistant_content[content_start:].split("<|im_end|>")[0].strip()
+            else:
+                generated_text = assistant_content.split("<|im_start|>assistant")[-1].split("<|im_end|>")[0].strip()
+        else:
+            generated_text = "抱歉，我无法生成合适的回复。"
+    else:
+        # 如果找不到标记，返回整个响应（去除提示部分）
+        generated_text = response.replace(prompt, "").strip()
+    
+    # 关键修改：直接返回字符串，而不是OpenAI格式的字典
+    return generated_text
 
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    # 提取助理的回复部分
-    generated_text = response.split("<|im_start|>assistant\n")[-1].strip()
-
-    # 返回OpenAI兼容格式
-    return {
-        "choices": [{
-            "message": {
-                "role": "assistant",
-                "content": generated_text
-            }
-        }]
-    }
-
-# 在Gradio界面启动前加载模型（可选）
-load_model()
+# 预先加载模型（可选，会延长启动时间但减少第一次请求的延迟）
+# load_model()
 
-# 创建Gradio聊天界面
+# 创建Gradio界面
 demo = gr.ChatInterface(
-    fn=openai_compatible_api,
-    title="DeepSeek API Service",
-    description="OpenAI-compatible API for DeepSeek-R1",
-    examples=["你好，请介绍一下你自己", "写一个Python函数计算斐波那契数列"]
+    fn=chat_with_deepseek,
+    title="DeepSeek-R1 聊天助手",
+    description="基于DeepSeek-R1-Distill-Qwen-1.5B的聊天机器人",
+    examples=["你好！", "请介绍一下你自己", "写一个Python函数计算斐波那契数列"],
+    cache_examples=False  # 禁用缓存，避免格式问题
 )
 
-# 修正后的launch调用 - 移除了show_api参数
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False
+    )