app.py

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
        ]
    }