File size: 3,837 Bytes
8ce1de8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | """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() # one generation at a time keeps memory/CPU use predictable on small hosts
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)}
|