Create app.py

app.py

@@ -0,0 +1,127 @@
+import os, json, time, uuid
+from threading import Thread
+
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from fastapi import FastAPI, HTTPException, Depends
+from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
+from pydantic import BaseModel
+from typing import List, Optional
+
+# ── Config ────────────────────────────────────────────────────────────────────
+MODEL_ID   = os.getenv("MODEL_ID", "Qwen/Qwen2.5-1.5B-Instruct")
+API_KEY    = os.getenv("API_KEY", "my-secret-key")   # set this in Space secrets
+DEVICE     = "cuda" if torch.cuda.is_available() else "cpu"
+DTYPE      = torch.float16 if DEVICE == "cuda" else torch.float32
+
+print(f"Loading {MODEL_ID} on {DEVICE} ...")
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
+model     = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    torch_dtype=DTYPE,
+    device_map="auto",
+    trust_remote_code=True,
+)
+model.eval()
+print("Model ready.")
+
+# ── FastAPI app ───────────────────────────────────────────────────────────────
+app     = FastAPI(title="LLM API", version="1.0")
+bearer  = HTTPBearer(auto_error=False)
+
+def verify_key(creds: Optional[HTTPAuthorizationCredentials] = Depends(bearer)):
+    if API_KEY and (creds is None or creds.credentials != API_KEY):
+        raise HTTPException(status_code=401, detail="Invalid or missing API key")
+
+# ── Schemas ───────────────────────────────────────────────────────────────────
+class Message(BaseModel):
+    role: str
+    content: str
+
+class ChatRequest(BaseModel):
+    model:       Optional[str]  = MODEL_ID
+    messages:    List[Message]
+    max_tokens:  Optional[int]  = 512
+    temperature: Optional[float]= 0.7
+    stream:      Optional[bool] = False
+
+# ── Routes ────────────────────────────────────────────────────────────────────
+@app.get("/")
+def root():
+    return {"status": "ok", "model": MODEL_ID, "device": DEVICE}
+
+@app.get("/v1/models")
+def list_models(_=Depends(verify_key)):
+    return {
+        "object": "list",
+        "data": [{"id": MODEL_ID, "object": "model", "owned_by": "user"}]
+    }
+
+@app.post("/v1/chat/completions")
+def chat_completions(req: ChatRequest, _=Depends(verify_key)):
+    msgs = [{"role": m.role, "content": m.content} for m in req.messages]
+
+    # Apply chat template
+    text   = tokenizer.apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
+    inputs = tokenizer([text], return_tensors="pt").to(model.device)
+
+    gen_kwargs = dict(
+        **inputs,
+        max_new_tokens = req.max_tokens,
+        temperature    = req.temperature,
+        do_sample      = req.temperature > 0,
+        pad_token_id   = tokenizer.eos_token_id,
+    )
+
+    cid = f"chatcmpl-{uuid.uuid4().hex[:12]}"
+
+    # ── Streaming ──────────────────────────────────────────────────────────────
+    if req.stream:
+        streamer = TextIteratorStreamer(
+            tokenizer, skip_special_tokens=True, skip_prompt=True
+        )
+        gen_kwargs["streamer"] = streamer
+        Thread(target=model.generate, kwargs=gen_kwargs, daemon=True).start()
+
+        def event_stream():
+            for token in streamer:
+                chunk = {
+                    "id": cid, "object": "chat.completion.chunk",
+                    "created": int(time.time()), "model": MODEL_ID,
+                    "choices": [{"delta": {"content": token}, "index": 0, "finish_reason": None}]
+                }
+                yield f"data: {json.dumps(chunk)}\n\n"
+            done = {
+                "id": cid, "object": "chat.completion.chunk",
+                "created": int(time.time()), "model": MODEL_ID,
+                "choices": [{"delta": {}, "index": 0, "finish_reason": "stop"}]
+            }
+            yield f"data: {json.dumps(done)}\n\n"
+            yield "data: [DONE]\n\n"
+
+        return StreamingResponse(event_stream(), media_type="text/event-stream")
+
+    # ── Non-streaming ──────────────────────────────────────────────────────────
+    with torch.no_grad():
+        output = model.generate(**gen_kwargs)
+
+    prompt_len  = inputs.input_ids.shape[1]
+    reply       = tokenizer.decode(output[0][prompt_len:], skip_special_tokens=True)
+    comp_tokens = output.shape[1] - prompt_len
+
+    return JSONResponse({
+        "id": cid, "object": "chat.completion",
+        "created": int(time.time()), "model": MODEL_ID,
+        "choices": [{
+            "index": 0,
+            "message": {"role": "assistant", "content": reply},
+            "finish_reason": "stop"
+        }],
+        "usage": {
+            "prompt_tokens":     prompt_len,
+            "completion_tokens": comp_tokens,
+            "total_tokens":      prompt_len + comp_tokens,
+        }
+    })