Create app.py

app.py

@@ -0,0 +1,78 @@
+from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+app = FastAPI()
+
+MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+print(f"Loading model on {device}...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    torch_dtype=torch.float16 if device == "cuda" else torch.float32,
+    device_map="auto",
+    trust_remote_code=True
+)
+print("Model loaded!")
+
+class ChatRequest(BaseModel):
+    message: str
+    max_tokens: int = 512
+    temperature: float = 0.7
+
+class CompletionRequest(BaseModel):
+    messages: list
+    max_tokens: int = 512
+    temperature: float = 0.7
+    stream: bool = False
+
+@app.post("/chat")
+def chat(req: ChatRequest):
+    messages = [{"role": "user", "content": req.message}]
+    text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer([text], return_tensors="pt").to(device)
+    
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=req.max_tokens,
+        temperature=req.temperature,
+        do_sample=True
+    )
+    
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    response = response.split("assistant\n")[-1].strip()
+    
+    return {"response": response}
+
+@app.post("/v1/chat/completions")
+def completions(req: CompletionRequest):
+    text = tokenizer.apply_chat_template(req.messages, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer([text], return_tensors="pt").to(device)
+    
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=req.max_tokens,
+        temperature=req.temperature,
+        do_sample=True
+    )
+    
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    response = response.split("assistant\n")[-1].strip()
+    
+    return {
+        "choices": [{
+            "message": {"role": "assistant", "content": response},
+            "finish_reason": "stop"
+        }]
+    }
+
+@app.get("/health")
+def health():
+    return {"status": "ok"}
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)