"""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)