Update app.py

app.py

@@ -1,63 +1,217 @@
 import os
-from fastapi import FastAPI, Request
-from fastapi.responses import JSONResponse
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
-
 # 设置缓存目录，避免 /.cache 权限问题
 os.environ["HF_HOME"] = "/tmp"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp"
 os.environ["HF_HUB_CACHE"] = "/tmp"
 
-# 初始化 FastAPI
-app = FastAPI()
-
-# 模型 ID
-MODEL_ID = "Qwen/Qwen2.5-Coder-7B-Instruct"
+import time
+import uuid
+from typing import List, Optional, Union, Dict, Any
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import json
 
-print("Loading model... (this may take a while the first time)")
-tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True, cache_dir="/tmp")
+# 初始化FastAPI应用
+app = FastAPI(title="Qwen Coder API", version="1.0.0")
 
-# 加载模型到 GPU (T4 支持 bfloat16，显存不够可换成 float16)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_ID,
-    device_map="auto",
-    torch_dtype=torch.bfloat16,
-    trust_remote_code=True,
-    cache_dir="/tmp"
+# 添加CORS中间件
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
 )
-model.eval()
-print("Model loaded.")
 
-# 生成接口 (兼容 OpenAI /v1/completions 简单版)
-@app.post("/v1/completions")
-async def completions(request: Request):
-    data = await request.json()
-    prompt = data.get("prompt") or ""
-    max_tokens = data.get("max_tokens", 512)
+# 全局变量
+model = None
+tokenizer = None
 
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+# Pydantic模型定义
+class Message(BaseModel):
+    role: str
+    content: str
+
+class ChatCompletionRequest(BaseModel):
+    model: str
+    messages: List[Message]
+    temperature: Optional[float] = 0.7
+    max_tokens: Optional[int] = 2048
+    stream: Optional[bool] = False
+    top_p: Optional[float] = 0.9
+
+class ChatCompletionChoice(BaseModel):
+    index: int
+    message: Message
+    finish_reason: str
+
+class Usage(BaseModel):
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+class ChatCompletionResponse(BaseModel):
+    id: str
+    object: str = "chat.completion"
+    created: int
+    model: str
+    choices: List[ChatCompletionChoice]
+    usage: Usage
+
+class Model(BaseModel):
+    id: str
+    object: str = "model"
+    created: int
+    owned_by: str = "qwen"
 
+class ModelListResponse(BaseModel):
+    object: str = "list"
+    data: List[Model]
+
+def load_model():
+    """加载Qwen Coder模型"""
+    global model, tokenizer
+    
+    model_name = "Qwen/Qwen2.5-Coder-7B-Instruct"
+    
+    print("Loading tokenizer...")
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name,
+        trust_remote_code=True
+    )
+    
+    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
+    )
+    
+    print("Model loaded successfully!")
+
+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
+
+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,
-            temperature=0.7,
-            top_p=0.9,
+            pad_token_id=tokenizer.eos_token_id
         )
+    
+    # 只返回新生成的部分
+    response = tokenizer.decode(
+        outputs[0][inputs['input_ids'].shape[1]:], 
+        skip_special_tokens=True
+    )
+    
+    return response.strip()
+
+@app.on_event("startup")
+async def startup_event():
+    """应用启动时加载模型"""
+    load_model()
 
-    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+@app.get("/")
+async def root():
+    return {"message": "Qwen Coder API Server is running!"}
 
-    # OpenAI API 格式返回
-    return JSONResponse({
-        "id": "cmpl-1",
-        "object": "text_completion",
-        "choices": [
-            {"index": 0, "text": text, "finish_reason": "stop"}
+@app.get("/v1/models", response_model=ModelListResponse)
+async def list_models():
+    """列出��用模型"""
+    return ModelListResponse(
+        data=[
+            Model(
+                id="qwen2.5-coder-7b-instruct",
+                created=int(time.time()),
+                owned_by="qwen"
+            )
         ]
-    })
+    )
 
-@app.get("/")
-def root():
-    return {"status": "ok", "model": MODEL_ID}
+@app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
+async def chat_completions(request: ChatCompletionRequest):
+    """处理聊天补全请求"""
+    try:
+        if model is None or tokenizer is None:
+            raise HTTPException(status_code=503, detail="Model not loaded")
+        
+        # 格式化消息
+        prompt = format_messages(request.messages)
+        
+        # 生成响应
+        response_text = generate_response(
+            prompt,
+            request.temperature,
+            request.max_tokens,
+            request.top_p
+        )
+        
+        # 构造响应
+        completion_id = f"chatcmpl-{uuid.uuid4().hex[:8]}"
+        
+        # 计算token使用量（简化版本）
+        prompt_tokens = len(tokenizer.encode(prompt))
+        completion_tokens = len(tokenizer.encode(response_text))
+        
+        response = ChatCompletionResponse(
+            id=completion_id,
+            created=int(time.time()),
+            model=request.model,
+            choices=[
+                ChatCompletionChoice(
+                    index=0,
+                    message=Message(role="assistant", content=response_text),
+                    finish_reason="stop"
+                )
+            ],
+            usage=Usage(
+                prompt_tokens=prompt_tokens,
+                completion_tokens=completion_tokens,
+                total_tokens=prompt_tokens + completion_tokens
+            )
+        )
+        
+        return response
+        
+    except Exception as e:
+        print(f"Error processing request: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/health")
+async def health_check():
+    """健康检查端点"""
+    return {
+        "status": "healthy",
+        "model_loaded": model is not None and tokenizer is not None
+    }
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)