Update app.py

app.py

@@ -1,43 +1,63 @@
 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"
 
-from fastapi import FastAPI
-from vllm import LLM, SamplingParams
-from pydantic import BaseModel
-from fastapi.responses import JSONResponse
-import uvicorn
+# 初始化 FastAPI
+app = FastAPI()
 
-# 1. 加载模型（会自动从 HuggingFace Hub 拉取）
-llm = LLM(model="Qwen/Qwen2.5-Coder-7B-Instruct",download_dir="/tmp")
-sampling = SamplingParams(temperature=0.2, max_tokens=1024)
+# 模型 ID
+MODEL_ID = "Qwen/Qwen2.5-Coder-7B-Instruct"
 
-app = FastAPI()
+print("Loading model... (this may take a while the first time)")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True, cache_dir="/tmp")
+
+# 加载模型到 GPU (T4 支持 bfloat16，显存不够可换成 float16)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    device_map="auto",
+    torch_dtype=torch.bfloat16,
+    trust_remote_code=True,
+    cache_dir="/tmp"
+)
+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)
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
-class Message(BaseModel):
-    role: str
-    content: str
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=max_tokens,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9,
+        )
 
-class ChatRequest(BaseModel):
-    model: str
-    messages: list[Message]
-    max_tokens: int = 1024
+    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-@app.post("/v1/chat/completions")
-def chat(req: ChatRequest):
-    prompt = "\n".join([f"{m.role}: {m.content}" for m in req.messages])
-    outputs = llm.generate([prompt], sampling)
-    text = outputs[0].outputs[0].text
+    # OpenAI API 格式返回
     return JSONResponse({
         "id": "cmpl-1",
-        "object": "chat.completion",
+        "object": "text_completion",
         "choices": [
-            {"index": 0,
-             "message": {"role": "assistant", "content": text},
-             "finish_reason": "stop"}
+            {"index": 0, "text": text, "finish_reason": "stop"}
         ]
     })
 
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+@app.get("/")
+def root():
+    return {"status": "ok", "model": MODEL_ID}