| """FastAPI inference service. |
| |
| Web page : GET / (src/static/index.html) | API docs: /docs |
| Run locally : uvicorn src.api:app --reload |
| Render start: uvicorn src.api:app --host 0.0.0.0 --port $PORT |
| |
| Environment variables |
| CHECKPOINT_PATH path to the checkpoint (default checkpoints/best.pt) |
| CHECKPOINT_URL optional: if the checkpoint file is missing it is downloaded from this URL at startup |
| (e.g. a GitHub release asset or a direct Google-Drive download link) |
| """ |
| from __future__ import annotations |
|
|
| import os |
| import threading |
| import urllib.request |
| from contextlib import asynccontextmanager |
|
|
| import torch |
| from fastapi import FastAPI, HTTPException |
| from fastapi.responses import FileResponse |
| from pydantic import BaseModel, Field |
|
|
| from src.generate import generate_text |
| from src.utils import load_checkpoint, resolve |
|
|
| STATE: dict = {} |
| MAX_NEW_TOKENS_CAP = 500 |
| _lock = threading.Lock() |
|
|
|
|
| def _load(): |
| path = resolve(os.environ.get("CHECKPOINT_PATH", "checkpoints/best.pt")) |
| url = os.environ.get("CHECKPOINT_URL") |
| if not path.exists() and url: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| print(f"Downloading checkpoint from {url} ...") |
| urllib.request.urlretrieve(url, path) |
| if not path.exists(): |
| STATE["error"] = (f"Checkpoint not found at {path}. Train a model (python -m src.train) or set CHECKPOINT_URL.") |
| return |
| torch.set_num_threads(max(1, min(2, os.cpu_count() or 1))) |
| model, tok, ckpt = load_checkpoint(path, "cpu") |
| STATE.update(model=model, tok=tok, ckpt=ckpt, error=None) |
|
|
|
|
| @asynccontextmanager |
| async def lifespan(app: FastAPI): |
| _load() |
| yield |
|
|
|
|
| app = FastAPI(title="MiniGPT API", description="Character-level GPT trained from scratch on Tiny Shakespeare", lifespan=lifespan) |
|
|
|
|
| class GenerateRequest(BaseModel): |
| prompt: str = Field("ROMEO:", min_length=1, max_length=1000) |
| max_new_tokens: int = Field(100, ge=1, le=MAX_NEW_TOKENS_CAP) |
| temperature: float = Field(0.8, gt=0.0, le=2.0) |
| top_k: int | None = Field(None, ge=1, le=1000) |
| seed: int | None = None |
|
|
|
|
| def _require_model(): |
| if STATE.get("model") is None: |
| raise HTTPException(status_code=503, detail=STATE.get("error", "model not loaded")) |
| return STATE["model"], STATE["tok"] |
|
|
|
|
| @app.get("/health") |
| def health(): |
| return {"status": "ok" if STATE.get("model") is not None else "model_not_loaded", "detail": STATE.get("error")} |
|
|
|
|
| @app.get("/", include_in_schema=False) |
| def home(): |
| """The web page (src/static/index.html) that calls POST /generate.""" |
| return FileResponse(resolve("src/static/index.html")) |
|
|
|
|
| @app.get("/info") |
| def info(): |
| model, _ = _require_model() |
| return {"name": "MiniGPT", "parameters": model.count_parameters(), "context_length": model.context_length, |
| "vocab_size": model.vocab_size, "n_layers": model.n_layers, "n_heads": model.n_heads, "d_model": model.d_model, |
| "usage": "POST /generate {prompt, max_new_tokens, temperature, top_k} | docs at /docs"} |
|
|
|
|
| @app.post("/generate") |
| def generate(req: GenerateRequest): |
| model, tok = _require_model() |
| unknown = tok.unknown_chars(req.prompt) |
| with _lock: |
| if req.seed is not None: |
| torch.manual_seed(req.seed) |
| text = generate_text(model, tok, req.prompt, req.max_new_tokens, req.temperature, top_k=req.top_k) |
| return {"prompt": req.prompt, "generated": text, "completion": text[len(req.prompt):], |
| "model_parameters": model.count_parameters(), "context_length": model.context_length, |
| "temperature": req.temperature, "max_new_tokens": req.max_new_tokens, |
| "note": (f"characters not in the vocabulary were replaced by a space: {unknown}" if unknown else None)} |
|
|
