Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+import json
 
 # 全局变量，避免重复加载
 model = None
@@ -25,36 +26,13 @@ def load_model():
     except Exception as e:
         print(f"模型加载失败: {e}")
 
-def chat_with_deepseek(message, history):
-    """与DeepSeek模型聊天 - 修正版"""
-    global model, tokenizer
-    
+def predict_api(message):
+    """API专用预测函数"""
     if model is None:
         load_model()
     
-    # 构建对话历史
-    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"
+    # 构建对话提示
+    prompt = f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
     
     # 编码输入
     inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048)
@@ -75,7 +53,7 @@ def chat_with_deepseek(message, history):
     # 解码回复
     response = tokenizer.decode(outputs[0], skip_special_tokens=False)
     
-    # 关键修正：提取助理的回复部分
+    # 提取助理的回复部分
     if "<|im_start|>assistant" in response:
         # 找到最后一个assistant标记开始的位置
         assistant_start = response.rfind("<|im_start|>assistant")
@@ -93,20 +71,51 @@ def chat_with_deepseek(message, history):
         # 如果找不到标记，返回整个响应（去除提示部分）
         generated_text = response.replace(prompt, "").strip()
     
-    # 关键修改：直接返回字符串，而不是OpenAI格式的字典
-    return generated_text
-
-# 预先加载模型（可选，会延长启动时间但减少第一次请求的延迟）
-# load_model()
+    # 返回OpenAI兼容格式
+    return {
+        "choices": [{
+            "message": {
+                "role": "assistant",
+                "content": generated_text
+            }
+        }]
+    }
 
 # 创建Gradio界面
-demo = gr.ChatInterface(
-    fn=chat_with_deepseek,
-    title="DeepSeek-R1 聊天助手",
-    description="基于DeepSeek-R1-Distill-Qwen-1.5B的聊天机器人",
-    examples=["你好！", "请介绍一下你自己", "写一个Python函数计算斐波那契数列"],
-    cache_examples=False  # 禁用缓存，避免格式问题
-)
+with gr.Blocks() as demo:
+    gr.Markdown("# DeepSeek-R1 API 服务")
+    
+    # 聊天界面
+    chatbot = gr.Chatbot(label="DeepSeek-R1")
+    msg = gr.Textbox(label="输入消息")
+    clear = gr.Button("清除")
+    
+    def respond(message, chat_history):
+        """处理聊天请求"""
+        # 调用预测函数
+        response = predict_api(message)
+        # 提取内容
+        bot_message = response["choices"][0]["message"]["content"]
+        # 更新聊天历史
+        chat_history.append((message, bot_message))
+        return "", chat_history
+    
+    # 设置界面交互
+    msg.submit(respond, [msg, chatbot], [msg, chatbot])
+    clear.click(lambda: None, None, chatbot, queue=False)
+    
+    # 添加API端点
+    gr.Interface(
+        fn=predict_api,
+        inputs=gr.Textbox(label="输入消息", lines=2),
+        outputs=gr.JSON(label="API响应"),
+        title="OpenAI兼容API",
+        description="使用此端点进行API调用",
+        api_name="predict"
+    )
+
+# 预先加载模型（可选）
+# load_model()
 
 if __name__ == "__main__":
     demo.launch(