Create app.py

app.py

@@ -0,0 +1,213 @@
+import os
+import time
+import uuid
+from typing import Any, Dict, List, Optional
+
+import torch
+from fastapi import FastAPI, Header, HTTPException
+from fastapi.responses import PlainTextResponse
+from pydantic import BaseModel
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+
+MODEL_ID = os.environ.get(
+    "MODEL_ID",
+    "Polygl0t/Tucano2-qwen-0.5B-Instruct"
+)
+
+API_KEY = os.environ.get("API_KEY", "")
+
+MAX_INPUT_TOKENS = int(os.environ.get("MAX_INPUT_TOKENS", "2048"))
+DEFAULT_MAX_NEW_TOKENS = int(os.environ.get("DEFAULT_MAX_NEW_TOKENS", "80"))
+
+tokenizer = None
+model = None
+
+app = FastAPI(title="Tucano OpenAI-Compatible API")
+
+
+class Message(BaseModel):
+    role: str
+    content: str
+
+
+class ChatCompletionRequest(BaseModel):
+    model: Optional[str] = None
+    messages: List[Message]
+    temperature: Optional[float] = 0.1
+    max_tokens: Optional[int] = DEFAULT_MAX_NEW_TOKENS
+    top_p: Optional[float] = 0.95
+    stream: Optional[bool] = False
+
+
+def check_auth(authorization: Optional[str]) -> None:
+    if not API_KEY:
+        return
+
+    if not authorization:
+        raise HTTPException(status_code=401, detail="Missing Authorization header")
+
+    expected = f"Bearer {API_KEY}"
+
+    if authorization != expected:
+        raise HTTPException(status_code=403, detail="Invalid API key")
+
+
+@app.on_event("startup")
+def load_model() -> None:
+    global tokenizer, model
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_ID,
+        trust_remote_code=True
+    )
+
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_ID,
+        torch_dtype=torch.float32,
+        low_cpu_mem_usage=True,
+        trust_remote_code=True
+    )
+
+    model.eval()
+
+
+@app.get("/", response_class=PlainTextResponse)
+def root() -> str:
+    return "Tucano OpenAI-compatible API is running."
+
+
+@app.get("/health")
+def health() -> Dict[str, str]:
+    return {
+        "status": "ok",
+        "model": MODEL_ID
+    }
+
+
+@app.get("/v1/models")
+def list_models(
+    authorization: Optional[str] = Header(default=None)
+) -> Dict[str, Any]:
+    check_auth(authorization)
+
+    return {
+        "object": "list",
+        "data": [
+            {
+                "id": MODEL_ID,
+                "object": "model",
+                "created": 0,
+                "owned_by": "huggingface-space"
+            }
+        ]
+    }
+
+
+@app.post("/v1/chat/completions")
+def chat_completions(
+    request: ChatCompletionRequest,
+    authorization: Optional[str] = Header(default=None)
+) -> Dict[str, Any]:
+    check_auth(authorization)
+
+    if request.stream:
+        raise HTTPException(
+            status_code=400,
+            detail="Streaming is not supported by this pilot endpoint"
+        )
+
+    if not request.messages:
+        raise HTTPException(
+            status_code=400,
+            detail="At least one message is required"
+        )
+
+    messages = [
+        {
+            "role": m.role,
+            "content": m.content
+        }
+        for m in request.messages
+    ]
+
+    try:
+        prompt = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+    except Exception:
+        prompt = ""
+
+        for m in messages:
+            prompt += f"{m['role']}: {m['content']}\n"
+
+        prompt += "assistant:"
+
+    inputs = tokenizer(
+        prompt,
+        return_tensors="pt",
+        truncation=True,
+        max_length=MAX_INPUT_TOKENS
+    )
+
+    max_new_tokens = request.max_tokens or DEFAULT_MAX_NEW_TOKENS
+
+    generation_args = {
+        "input_ids": inputs["input_ids"],
+        "attention_mask": inputs.get("attention_mask"),
+        "max_new_tokens": max_new_tokens,
+        "pad_token_id": tokenizer.eos_token_id,
+    }
+
+    if request.temperature is not None and request.temperature > 0:
+        generation_args["do_sample"] = True
+        generation_args["temperature"] = request.temperature
+        generation_args["top_p"] = request.top_p or 0.95
+    else:
+        generation_args["do_sample"] = False
+
+    start_time = time.time()
+
+    with torch.no_grad():
+        output_ids = model.generate(**generation_args)
+
+    elapsed = time.time() - start_time
+
+    prompt_length = inputs["input_ids"].shape[-1]
+    generated_ids = output_ids[0][prompt_length:]
+
+    content = tokenizer.decode(
+        generated_ids,
+        skip_special_tokens=True
+    ).strip()
+
+    prompt_tokens = int(prompt_length)
+    completion_tokens = int(generated_ids.shape[-1])
+
+    return {
+        "id": f"chatcmpl-{uuid.uuid4().hex}",
+        "object": "chat.completion",
+        "created": int(time.time()),
+        "model": request.model or MODEL_ID,
+        "choices": [
+            {
+                "index": 0,
+                "message": {
+                    "role": "assistant",
+                    "content": content
+                },
+                "finish_reason": "stop"
+            }
+        ],
+        "usage": {
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": prompt_tokens + completion_tokens
+        },
+        "pilot_metadata": {
+            "served_model": MODEL_ID,
+            "elapsed_seconds": round(elapsed, 3)
+        }
+    }