Update app.py

app.py

@@ -1,125 +1,217 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
-import json
+import logging
+import time
 
-# 全局变量，避免重复加载
+# 配置日志
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# 全局变量
 model = None
 tokenizer = None
+last_load_time = 0
 
-def load_model():
-    """加载模型和分词器"""
-    global model, tokenizer
+def safe_load_model():
+    """安全加载模型，带错误重试机制"""
+    global model, tokenizer, last_load_time
+    
+    # 避免频繁重载模型
+    if model is not None and time.time() - last_load_time < 300:  # 5分钟内不重载
+        return True
+    
     model_name = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
+    max_retries = 3
+    
+    for attempt in range(max_retries):
+        try:
+            logger.info(f"尝试加载模型，第{attempt + 1}次...")
+            
+            # 清理GPU缓存（如果有）
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            
+            tokenizer = AutoTokenizer.from_pretrained(
+                model_name, 
+                trust_remote_code=True,
+                resume_download=True  # 支持断点续传
+            )
+            
+            # 确保tokenizer有pad_token
+            if tokenizer.pad_token is None:
+                tokenizer.pad_token = tokenizer.eos_token
+            
+            model = AutoModelForCausalLM.from_pretrained(
+                model_name,
+                torch_dtype=torch.float16,
+                device_map="auto",
+                low_cpu_mem_usage=True,
+                trust_remote_code=True,
+                resume_download=True
+            )
+            
+            last_load_time = time.time()
+            logger.info("模型加载成功！")
+            return True
+            
+        except Exception as e:
+            logger.error(f"模型加载失败（尝试{attempt + 1}/{max_retries}）: {e}")
+            if attempt < max_retries - 1:
+                time.sleep(5)  # 等待5秒后重试
+            else:
+                return False
+
+def generate_response_safe(message, max_retries=2):
+    """安全的响应生成函数，带重试机制"""
+    for attempt in range(max_retries):
+        try:
+            if not safe_load_model():
+                return "模型加载失败，请稍后重试"
+            
+            # 限制输入长度，避免内存溢出
+            if len(message) > 2000:
+                message = message[:2000] + "...(内容过长已截断)"
+            
+            # 构建提示词
+            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
+            )
+            
+            # 生成回复，限制生成长度
+            with torch.no_grad():
+                outputs = model.generate(
+                    **inputs,
+                    max_new_tokens=256,  # 减少生成长度
+                    temperature=0.7,
+                    top_p=0.9,
+                    do_sample=True,
+                    pad_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_start = response.rfind("<|im_start|>assistant")
+                if assistant_start != -1:
+                    assistant_content = response[assistant_start:]
+                    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()
+            
+            # 清理缓存
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+                
+            return generated_text
+            
+        except Exception as e:
+            logger.error(f"生成响应失败（尝试{attempt + 1}/{max_retries}）: {e}")
+            if attempt < max_retries - 1:
+                time.sleep(2)
+            else:
+                return f"生成失败: {str(e)}"
+
+def process_chat(message, history):
+    """处理聊天请求"""
     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
-        )
-        # 确保tokenizer有pad_token
-        if tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
-        print("模型加载成功！")
+        response = generate_response_safe(message)
+        return response
     except Exception as e:
-        print(f"模型加载失败: {e}")
+        logger.error(f"聊天处理异常: {e}")
+        return "抱歉，处理您的请求时出现了问题，请稍后重试。"
 
-def predict_api(message):
-    """API专用预测函数"""
-    if model is None:
-        load_model()
+# 创建简化版的Gradio界面
+with gr.Blocks(
+    theme=gr.themes.Soft(),
+    css="""
+    .gradio-container {
+        max-width: 800px !important;
+    }
+    """
+) as demo:
     
-    # 构建对话提示
-    prompt = f"<|im_start|>user\n{message}<|im_end|>\n<|im_start|>assistant\n"
+    gr.Markdown("""
+    # DeepSeek-R1 聊天助手
+    *基于DeepSeek-R1-Distill-Qwen-1.5B模型*
+    """)
     
-    # 编码输入
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048)
+    # 简化聊天组件
+    chatbot = gr.Chatbot(
+        label="对话历史",
+        height=400,
+        show_copy_button=True
+    )
     
-    # 生成回复
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=512,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id,
-            eos_token_id=tokenizer.eos_token_id,
-            repetition_penalty=1.1
+    with gr.Row():
+        msg = gr.Textbox(
+            label="输入消息",
+            placeholder="请输入您的问题...",
+            scale=4,
+            max_lines=3
         )
+        submit_btn = gr.Button("发送", variant="primary", scale=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 {
-        "choices": [{
-            "message": {
-                "role": "assistant",
-                "content": generated_text
-            }
-        }]
-    }
-
-# 创建Gradio界面
-with gr.Blocks() as demo:
-    gr.Markdown("# DeepSeek-R1 API 服务")
-    
-    # 聊天界面
-    chatbot = gr.Chatbot(label="DeepSeek-R1")
-    msg = gr.Textbox(label="输入消息")
-    clear = gr.Button("清除")
+    clear_btn = gr.Button("清除对话")
     
+    # 处理函数
     def respond(message, chat_history):
-        """处理聊天请求"""
-        # 调用预测函数
-        response = predict_api(message)
-        # 提取内容
-        bot_message = response["choices"][0]["message"]["content"]
-        # 更新聊天历史
-        chat_history.append((message, bot_message))
+        if not message.strip():
+            return "", chat_history
+        
+        chat_history.append([message, ""])
+        response = process_chat(message, chat_history)
+        chat_history[-1][1] = response
+        
         return "", chat_history
     
-    # 设置界面交互
-    msg.submit(respond, [msg, chatbot], [msg, chatbot])
-    clear.click(lambda: None, None, chatbot, queue=False)
+    # 事件绑定
+    msg_submit = msg.submit(
+        respond, [msg, chatbot], [msg, chatbot],
+        queue=True,
+        show_progress="hidden"
+    )
+    
+    btn_click = submit_btn.click(
+        respond, [msg, chatbot], [msg, chatbot],
+        queue=True,
+        show_progress="hidden"
+    )
     
-    # 添加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"
+    clear_btn.click(
+        lambda: None, None, chatbot, queue=False
     )
 
-# 预先加载模型（可选）
-# load_model()
+# 预加载模型（可选）
+try:
+    safe_load_model()
+except Exception as e:
+    logger.warning(f"预加载模型失败: {e}")
 
 if __name__ == "__main__":
+    # 优化启动配置
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
-        share=False
+        share=False,
+        show_error=True,
+        debug=False,  # 关闭调试模式减少输出
+        max_threads=2,  # 限制线程数
+        quiet=True  # 减少日志输出
     )