examples/fastapi_app.py

"""FastAPI service with automatic PyTorch or CPU INT8 ONNX inference."""

from __future__ import annotations

import hmac
import json
import os
import threading
from contextlib import asynccontextmanager
from pathlib import Path
from typing import Annotated

import numpy as np
import torch
from fastapi import Depends, FastAPI, HTTPException
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from huggingface_hub import hf_hub_download
from pydantic import BaseModel, Field
from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer

DEFAULT_MODEL_ID = "NajahUniv/AraUni-MARBERTv2-Intent-Classifier"
MODEL_ID = os.getenv("MODEL_ID", DEFAULT_MODEL_ID)
MODEL_REVISION = os.getenv("MODEL_REVISION")
MODEL_DEVICE = os.getenv("MODEL_DEVICE", "auto")
MODEL_BACKEND = os.getenv("MODEL_BACKEND", "auto")
MODEL_PRECISION = os.getenv("MODEL_PRECISION", "auto")
MODEL_API_KEY = os.getenv("MODEL_API_KEY")
MAX_BATCH_SIZE = int(os.getenv("MAX_BATCH_SIZE", "64"))


def choose_device() -> str:
    if MODEL_BACKEND == "onnx" and MODEL_DEVICE == "auto":
        return "cpu"
    if MODEL_DEVICE != "auto":
        return MODEL_DEVICE
    if torch.cuda.is_available():
        return "cuda"
    if torch.backends.mps.is_available():
        return "mps"
    return "cpu"


def resolve_runtime(device: str) -> tuple[str, str]:
    if MODEL_BACKEND not in {"auto", "pytorch", "onnx"}:
        raise ValueError("MODEL_BACKEND must be auto, pytorch, or onnx")
    if MODEL_PRECISION not in {"auto", "fp32", "bf16", "fp16", "int8"}:
        raise ValueError("MODEL_PRECISION must be auto, fp32, bf16, fp16, or int8")
    backend = "pytorch" if MODEL_BACKEND == "auto" else MODEL_BACKEND
    precision = MODEL_PRECISION
    if precision == "auto":
        if backend == "onnx":
            precision = "int8"
        elif device == "cuda":
            precision = "bf16" if torch.cuda.is_bf16_supported() else "fp16"
        else:
            precision = "fp32"
    if backend == "onnx" and (device != "cpu" or precision != "int8"):
        raise ValueError("the published ONNX artifact supports CPU INT8 only")
    if backend == "pytorch" and precision == "int8":
        raise ValueError("MODEL_PRECISION=int8 requires MODEL_BACKEND=onnx")
    if backend == "pytorch" and precision in {"bf16", "fp16"} and device != "cuda":
        raise ValueError("this example enables bf16/fp16 only on CUDA")
    return backend, precision


def hub_or_local_file(filename: str) -> str:
    local = Path(MODEL_ID) / filename
    if local.is_file():
        return str(local)
    return hf_hub_download(MODEL_ID, filename, revision=MODEL_REVISION)


class ClassifyRequest(BaseModel):
    texts: list[str] = Field(min_length=1)
    top_k: int = Field(default=5, ge=1)
    threshold: float | None = Field(default=None, gt=0, lt=1)


class ModelRuntime:
    def __init__(self) -> None:
        load_kwargs = {"revision": MODEL_REVISION} if MODEL_REVISION else {}
        self.tokenizer = AutoTokenizer.from_pretrained(
            MODEL_ID,
            trust_remote_code=True,
            **load_kwargs,
        )
        self.config = AutoConfig.from_pretrained(
            MODEL_ID,
            trust_remote_code=True,
            **load_kwargs,
        )
        self.device = choose_device()
        self.backend, self.precision = resolve_runtime(self.device)
        self.lock = threading.Lock()
        self.thresholds = self.config.thresholds
        self.model = None
        self.session = None
        if self.backend == "onnx":
            import onnxruntime as ort

            onnx_config = json.loads(
                Path(hub_or_local_file("onnx/onnx_config.json")).read_text(encoding="utf-8")
            )
            self.thresholds = onnx_config["thresholds"]
            self.session = ort.InferenceSession(
                hub_or_local_file("onnx/model_int8.onnx"),
                providers=["CPUExecutionProvider"],
            )
        else:
            dtype = {
                "fp32": torch.float32,
                "bf16": torch.bfloat16,
                "fp16": torch.float16,
            }[self.precision]
            self.model = AutoModelForSequenceClassification.from_pretrained(
                MODEL_ID,
                trust_remote_code=True,
                torch_dtype=dtype,
                **load_kwargs,
            ).to(self.device).eval()

    def probabilities(self, texts: list[str]) -> np.ndarray:
        if self.backend == "onnx":
            encoded = self.tokenizer(
                texts,
                return_tensors="np",
                padding=True,
                truncation=True,
                max_length=self.config.max_length,
            )
            with self.lock:
                logits = self.session.run(
                    ["logits"],
                    {
                        "input_ids": encoded["input_ids"].astype(np.int64),
                        "attention_mask": encoded["attention_mask"].astype(np.int64),
                    },
                )[0]
            return 1.0 / (1.0 + np.exp(-logits))
        encoded = self.tokenizer(
            texts,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=self.config.max_length,
        ).to(self.device)
        with self.lock, torch.inference_mode():
            return torch.sigmoid(self.model(**encoded).logits).cpu().float().numpy()

    def classify(self, request: ClassifyRequest) -> list[dict[str, object]]:
        if len(request.texts) > MAX_BATCH_SIZE:
            raise HTTPException(413, f"at most {MAX_BATCH_SIZE} texts are allowed per request")
        probabilities = self.probabilities(request.texts)
        labels = [self.config.id2label[index] for index in range(self.config.num_labels)]
        results = []
        for text, row in zip(request.texts, probabilities, strict=True):
            scores = []
            for index, label in enumerate(labels):
                threshold = (
                    request.threshold
                    if request.threshold is not None
                    else float(self.thresholds[label])
                )
                scores.append(
                    {
                        "label": label,
                        "probability": float(row[index]),
                        "threshold": threshold,
                        "selected": float(row[index]) >= threshold,
                    }
                )
            scores.sort(key=lambda item: item["probability"], reverse=True)
            results.append(
                {
                    "text": text,
                    "selected_labels": [item["label"] for item in scores if item["selected"]],
                    "scores": scores[: min(request.top_k, len(scores))],
                }
            )
        return results


runtime: ModelRuntime | None = None


@asynccontextmanager
async def lifespan(_: FastAPI):
    global runtime
    runtime = ModelRuntime()
    yield
    runtime = None


app = FastAPI(title="AraUni Multi-label Intent Classifier", lifespan=lifespan)
bearer_scheme = HTTPBearer(
    auto_error=False,
    scheme_name="BearerAuth",
    description="Enter the MODEL_API_KEY value. Swagger adds the 'Bearer' prefix.",
)


def authorize(
    credentials: Annotated[
        HTTPAuthorizationCredentials | None,
        Depends(bearer_scheme),
    ],
) -> None:
    if MODEL_API_KEY is None:
        return
    if (
        credentials is None
        or credentials.scheme.lower() != "bearer"
        or not hmac.compare_digest(credentials.credentials, MODEL_API_KEY)
    ):
        raise HTTPException(
            401,
            "invalid bearer token",
            headers={"WWW-Authenticate": "Bearer"},
        )


@app.get("/health")
def health() -> dict[str, object]:
    return {
        "status": "ok",
        "model_id": MODEL_ID,
        "device": runtime.device if runtime else None,
        "backend": runtime.backend if runtime else None,
        "precision": runtime.precision if runtime else None,
    }


@app.get("/labels", dependencies=[Depends(authorize)])
def labels() -> dict[int, str]:
    if runtime is None:
        raise HTTPException(503, "model is not ready")
    return dict(runtime.config.id2label)


@app.post("/classify", dependencies=[Depends(authorize)])
def classify(request: ClassifyRequest) -> list[dict[str, object]]:
    if runtime is None:
        raise HTTPException(503, "model is not ready")
    return runtime.classify(request)