import os from fastapi import FastAPI, HTTPException from pydantic import BaseModel from transformers import pipeline MODEL_ID = os.getenv("MODEL_ID", "killykilly/nigerian-transaction-classifier") MAX_BATCH = int(os.getenv("MAX_BATCH", "256")) classifier = pipeline( "text-classification", model=MODEL_ID, token=os.getenv("HF_TOKEN") or None, truncation=True, ) app = FastAPI(title="Nigerian Transaction Classifier API") class ClassifyRequest(BaseModel): texts: list[str] @app.get("/health") def health(): return {"status": "ok", "model": MODEL_ID} @app.post("/classify") def classify(request: ClassifyRequest): if len(request.texts) > MAX_BATCH: raise HTTPException(413, f"batch too large (max {MAX_BATCH})") if not request.texts: return {"predictions": []} results = classifier(request.texts, batch_size=32) return { "predictions": [ {"label": r["label"], "score": float(r["score"])} for r in results ] }