Update app.py

app.py

@@ -11,7 +11,7 @@ from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
 import json
 
 # 初始化FastAPI应用
@@ -29,6 +29,7 @@ app.add_middleware(
 # 全局变量
 model = None
 tokenizer = None
+model_name = None
 
 # Pydantic模型定义
 class Message(BaseModel):
@@ -73,82 +74,158 @@ class ModelListResponse(BaseModel):
 
 def load_model():
     """加载Qwen Coder模型"""
-    global model, tokenizer
+    global model, tokenizer, model_name
     
-    model_name = "Qwen/Qwen2.5-Coder-7B-Instruct"
+    # 模型选择优先级列表
+    model_candidates = [
+        "Qwen/Qwen2.5-Coder-7B-Instruct",
+        "Qwen/Qwen2.5-Coder-3B-Instruct", 
+        "Qwen/Qwen2.5-Coder-1.5B-Instruct"
+    ]
     
-    print("Loading tokenizer...")
-    tokenizer = AutoTokenizer.from_pretrained(
-        model_name,
-        trust_remote_code=True
-    )
+    for candidate_model in model_candidates:
+        try:
+            print(f"Attempting to load model: {candidate_model}")
+            
+            # 先测试tokenizer加载
+            print("Loading tokenizer...")
+            test_tokenizer = AutoTokenizer.from_pretrained(
+                candidate_model,
+                trust_remote_code=True,
+                use_fast=False,
+                revision="main"
+            )
+            
+            # 如果tokenizer加载成功，继续加载模型
+            print("Loading model...")
+            test_model = AutoModelForCausalLM.from_pretrained(
+                candidate_model,
+                torch_dtype=torch.float16,
+                device_map="auto",
+                trust_remote_code=True,
+                low_cpu_mem_usage=True,
+                revision="main"
+            )
+            
+            # 成功加载后赋值给全局变量
+            tokenizer = test_tokenizer
+            model = test_model
+            model_name = candidate_model
+            
+            print(f"Successfully loaded model: {candidate_model}")
+            return
+            
+        except Exception as e:
+            print(f"Failed to load {candidate_model}: {str(e)}")
+            continue
     
-    print("Loading model...")
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        torch_dtype=torch.float16,
-        device_map="auto",
-        trust_remote_code=True,
-        low_cpu_mem_usage=True
-    )
+    # 如果所有模型都失败，抛出异常
+    raise Exception("Failed to load any Qwen model. Please check your configuration.")
+
+def format_messages_simple(messages: List[Message]) -> str:
+    """简单的消息格式化（备用方案）"""
+    formatted = ""
+    for msg in messages:
+        if msg.role == "system":
+            formatted += f"System: {msg.content}\n\n"
+        elif msg.role == "user":
+            formatted += f"User: {msg.content}\n\n"
+        elif msg.role == "assistant":
+            formatted += f"Assistant: {msg.content}\n\n"
     
-    print("Model loaded successfully!")
+    formatted += "Assistant: "
+    return formatted
 
 def format_messages(messages: List[Message]) -> str:
     """将消息格式化为Qwen格式"""
-    formatted_messages = []
-    for msg in messages:
-        formatted_messages.append({
-            "role": msg.role,
-            "content": msg.content
-        })
-    
-    # 使用tokenizer的chat template
-    text = tokenizer.apply_chat_template(
-        formatted_messages,
-        tokenize=False,
-        add_generation_prompt=True
-    )
-    return text
+    try:
+        formatted_messages = []
+        for msg in messages:
+            formatted_messages.append({
+                "role": msg.role,
+                "content": msg.content
+            })
+        
+        # 尝试使用tokenizer的chat template
+        if hasattr(tokenizer, 'apply_chat_template'):
+            text = tokenizer.apply_chat_template(
+                formatted_messages,
+                tokenize=False,
+                add_generation_prompt=True
+            )
+            return text
+        else:
+            # 如果没有chat_template，使用简单格式化
+            return format_messages_simple(messages)
+            
+    except Exception as e:
+        print(f"Error in format_messages, using simple format: {str(e)}")
+        return format_messages_simple(messages)
 
 def generate_response(prompt: str, temperature: float, max_tokens: int, top_p: float) -> str:
     """生成模型响应"""
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
+    try:
+        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=4096)
+        
+        # 移动到模型设备
+        if hasattr(model, 'device'):
+            inputs = {k: v.to(model.device) for k, v in inputs.items()}
+        
+        with torch.no_grad():
+            # 设置生成参数
+            generation_config = GenerationConfig(
+                max_new_tokens=min(max_tokens, 2048),
+                temperature=temperature,
+                top_p=top_p,
+                do_sample=True,
+                pad_token_id=tokenizer.eos_token_id,
+                eos_token_id=tokenizer.eos_token_id,
+                repetition_penalty=1.1
+            )
+            
+            outputs = model.generate(
+                **inputs,
+                generation_config=generation_config
+            )
+        
+        # 只返回新生成的部分
+        response = tokenizer.decode(
+            outputs[0][inputs['input_ids'].shape[1]:], 
+            skip_special_tokens=True
         )
-    
-    # 只返回新生成的部分
-    response = tokenizer.decode(
-        outputs[0][inputs['input_ids'].shape[1]:], 
-        skip_special_tokens=True
-    )
-    
-    return response.strip()
+        
+        return response.strip()
+        
+    except Exception as e:
+        print(f"Error in generate_response: {str(e)}")
+        return f"抱歉，生成响应时出现错误: {str(e)}"
 
 @app.on_event("startup")
 async def startup_event():
     """应用启动时加载模型"""
-    load_model()
+    try:
+        load_model()
+    except Exception as e:
+        print(f"Failed to load model during startup: {str(e)}")
+        # 不要让启动失败，而是在请求时返回错误
 
 @app.get("/")
 async def root():
-    return {"message": "Qwen Coder API Server is running!"}
+    return {
+        "message": "Qwen Coder API Server is running!",
+        "model_loaded": model is not None,
+        "current_model": model_name
+    }
 
 @app.get("/v1/models", response_model=ModelListResponse)
 async def list_models():
     """列出可用模型"""
+    model_id = "qwen2.5-coder-7b-instruct" if model_name is None else model_name.split("/")[-1].lower()
+    
     return ModelListResponse(
         data=[
             Model(
-                id="qwen2.5-coder-7b-instruct",
+                id=model_id,
                 created=int(time.time()),
                 owned_by="qwen"
             )
@@ -160,7 +237,11 @@ async def chat_completions(request: ChatCompletionRequest):
     """处理聊天补全请求"""
     try:
         if model is None or tokenizer is None:
-            raise HTTPException(status_code=503, detail="Model not loaded")
+            # 尝试重新加载模型
+            try:
+                load_model()
+            except:
+                raise HTTPException(status_code=503, detail="Model not loaded and failed to load on demand")
         
         # 格式化消息
         prompt = format_messages(request.messages)
@@ -177,8 +258,13 @@ async def chat_completions(request: ChatCompletionRequest):
         completion_id = f"chatcmpl-{uuid.uuid4().hex[:8]}"
         
         # 计算token使用量（简化版本）
-        prompt_tokens = len(tokenizer.encode(prompt))
-        completion_tokens = len(tokenizer.encode(response_text))
+        try:
+            prompt_tokens = len(tokenizer.encode(prompt))
+            completion_tokens = len(tokenizer.encode(response_text))
+        except:
+            # 如果tokenizer编码失败，使用估算
+            prompt_tokens = len(prompt.split()) * 2
+            completion_tokens = len(response_text.split()) * 2
         
         response = ChatCompletionResponse(
             id=completion_id,
@@ -202,14 +288,18 @@ async def chat_completions(request: ChatCompletionRequest):
         
     except Exception as e:
         print(f"Error processing request: {str(e)}")
-        raise HTTPException(status_code=500, detail=str(e))
+        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
 
 @app.get("/health")
 async def health_check():
     """健康检查端点"""
     return {
-        "status": "healthy",
-        "model_loaded": model is not None and tokenizer is not None
+        "status": "healthy" if model is not None and tokenizer is not None else "unhealthy",
+        "model_loaded": model is not None and tokenizer is not None,
+        "current_model": model_name,
+        "torch_version": torch.__version__,
+        "cuda_available": torch.cuda.is_available(),
+        "device_count": torch.cuda.device_count() if torch.cuda.is_available() else 0
     }
 
 if __name__ == "__main__":