api/main.py

from __future__ import annotations

import logging
import os
import re
import time
import unicodedata
from contextlib import asynccontextmanager
from pathlib import Path

import numpy as np
import onnxruntime as ort
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, field_validator
from transformers import AutoTokenizer

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
log = logging.getLogger(__name__)

ROOT = Path(__file__).resolve().parent.parent
MODEL_DIR = ROOT / "models" / "muril_scam_v1"
QUANT_ONNX_PATH = ROOT / "models" / "muril_scam_v1_quant.onnx"
ONNX_PATH = QUANT_ONNX_PATH if QUANT_ONNX_PATH.exists() else ROOT / "models" / "muril_scam_v1.onnx"

MAX_LENGTH = 128
SCAM_THRESHOLD = 0.50

ALLOWED_CATEGORIES = {"digital_arrest", "otp_fraud", "courier_scam", "safe"}
ALLOWED_FLAGS = {"urgency", "authority_impersonation", "threat", "payment_demand"}

_DIGITAL_ARREST_RE = re.compile(
    r"\b(cbi|ed|nia|ncb|cyber\s*cell|digital\s*arrest|warrant|fir|court|money\s*laundering|hawala|customs)\b",
    re.IGNORECASE,
)
_OTP_FRAUD_RE = re.compile(
    r"\b(otp|one.?time.?password|verification\s*code|share\s*code|kyc|bank\s*update|account\s*blocked)\b",
    re.IGNORECASE,
)
_COURIER_SCAM_RE = re.compile(
    r"\b(parcel|package|courier|delivery|shipment|customs\s*hold|seized)\b",
    re.IGNORECASE,
)

FLAGS = {
    "urgency": re.compile(
        r"\b(urgent|immediately|right\s*now|turant|jaldi|deadline|last\s*chance|within\s+\d+\s+(minute|hour))\b",
        re.IGNORECASE,
    ),
    "authority_impersonation": re.compile(
        r"\b(cbi|ed|ncb|nia|police|officer|inspector|court|government|commissioner)\b",
        re.IGNORECASE,
    ),
    "threat": re.compile(
        r"\b(arrest|warrant|jail|prison|legal\s*action|fir|criminal\s*case|freeze|suspend)\b",
        re.IGNORECASE,
    ),
    "payment_demand": re.compile(
        r"\b(pay|payment|transfer|send\s*money|upi|neft|rtgs|wire|fee|penalty|amount|rupee)\b",
        re.IGNORECASE,
    ),
}

_URL_RE = re.compile(r"https?://\S+|www\.\S+", re.IGNORECASE)
_PHONE_RE = re.compile(
    r"(\+91[\s\-]?)?[6-9]\d{9}|\b\d{10}\b|\b\d{5}[\s\-]\d{5}\b|\+\d{1,3}[\s\-]\d{6,12}"
)
_EMAIL_RE = re.compile(r"\S+@\S+\.\S+")
_AMOUNT_RE = re.compile(r"₹\s?\d[\d,]*(\.\d+)?|\brs\.?\s?\d[\d,]*", re.IGNORECASE)
_OTP_RE = re.compile(r"\b\d{4,8}\b")
_AADHAAR_RE = re.compile(r"\b\d{4}\s\d{4}\s\d{4}\b")
_PAN_RE = re.compile(r"\b[A-Z]{5}[0-9]{4}[A-Z]\b")
_SPACE_RE = re.compile(r"\s{2,}")


def normalize(text: str) -> str:
    text = unicodedata.normalize("NFC", text)
    text = _AADHAAR_RE.sub("[AADHAAR]", text)
    text = _PAN_RE.sub("[PAN]", text)
    text = _URL_RE.sub("[URL]", text)
    text = _EMAIL_RE.sub("[EMAIL]", text)
    text = _PHONE_RE.sub("[PHONE]", text)
    text = _AMOUNT_RE.sub("[AMOUNT]", text)
    text = _OTP_RE.sub("[CODE]", text)
    chars = [ch.lower() if ord(ch) < 128 else ch for ch in text]
    return _SPACE_RE.sub(" ", "".join(chars)).strip()


def detect_category(text: str, label: str) -> str:
    if label == "safe":
        return "safe"
    if _DIGITAL_ARREST_RE.search(text):
        return "digital_arrest"
    if _OTP_FRAUD_RE.search(text):
        return "otp_fraud"
    if _COURIER_SCAM_RE.search(text):
        return "courier_scam"
    return "safe"


def extract_flags(text: str) -> list[str]:
    matches = [name for name, pattern in FLAGS.items() if pattern.search(text)]
    return [f for f in matches if f in ALLOWED_FLAGS]


class AnalyzeRequest(BaseModel):
    text: str
    language: str = "auto"

    @field_validator("text")
    @classmethod
    def validate_text(cls, value: str) -> str:
        value = value.strip()
        if not value:
            raise ValueError("text must not be empty")
        if len(value) > 2000:
            raise ValueError("text must be under 2000 characters")
        return value


class AnalyzeResponse(BaseModel):
    score: float
    label: str
    category: str
    flags: list[str]
    latency_ms: int


class ScamDetector:
    session: ort.InferenceSession | None = None
    tokenizer: AutoTokenizer | None = None
    input_names: list[str] = []
    output_names: list[str] = []

    @classmethod
    def load(cls) -> None:
        if cls.session is not None and cls.tokenizer is not None:
            return

        if not MODEL_DIR.exists():
            raise RuntimeError(f"Model directory not found: {MODEL_DIR}")
        if not ONNX_PATH.exists():
            raise RuntimeError(f"ONNX model not found: {ONNX_PATH}. Run: python src/export_onnx.py")

        log.info("Loading tokenizer from %s", MODEL_DIR)
        cls.tokenizer = AutoTokenizer.from_pretrained(str(MODEL_DIR))

        log.info("Loading ONNX model from %s", ONNX_PATH)
        options = ort.SessionOptions()
        options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
        options.intra_op_num_threads = os.cpu_count() or 2

        providers = ["CPUExecutionProvider"]
        cls.session = ort.InferenceSession(str(ONNX_PATH), sess_options=options, providers=providers)

        cls.input_names = [i.name for i in cls.session.get_inputs()]
        cls.output_names = [o.name for o in cls.session.get_outputs()]
        cls._warmup()

    @classmethod
    def _warmup(cls) -> None:
        assert cls.session is not None and cls.tokenizer is not None
        encoded = cls.tokenizer(
            "warmup text",
            return_tensors="np",
            padding="max_length",
            truncation=True,
            max_length=MAX_LENGTH,
        )
        feed = {
            name: encoded[name].astype(np.int64, copy=False)
            for name in cls.input_names
            if name in encoded
        }
        cls.session.run(cls.output_names, feed)

    @classmethod
    def predict(cls, text: str) -> dict:
        if cls.session is None or cls.tokenizer is None:
            raise RuntimeError("Model is not loaded")

        start = time.perf_counter()
        normalized = normalize(text)
        encoded = cls.tokenizer(
            normalized,
            return_tensors="np",
            padding="max_length",
            truncation=True,
            max_length=MAX_LENGTH,
        )
        feed = {
            name: encoded[name].astype(np.int64, copy=False)
            for name in cls.input_names
            if name in encoded
        }
        outputs = cls.session.run(cls.output_names, feed)

        logits = outputs[0][0]
        logits = logits - np.max(logits)
        probs = np.exp(logits) / np.exp(logits).sum()
        scam_score = float(probs[1])
        label = "scam" if scam_score >= SCAM_THRESHOLD else "safe"

        category = detect_category(text, label)
        if category not in ALLOWED_CATEGORIES:
            category = "safe"

        latency_ms = int((time.perf_counter() - start) * 1000)
        return {
            "score": round(scam_score, 4),
            "label": label,
            "category": category,
            "flags": extract_flags(text),
            "latency_ms": latency_ms,
        }


@asynccontextmanager
async def lifespan(_app: FastAPI):
    ScamDetector.load()
    yield


app = FastAPI(
    title="Scam ONNX API",
    version="1.0.0",
    lifespan=lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:3000", "http://127.0.0.1:3000"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/health")
def health() -> dict:
    return {
        "status": "ok",
        "model_loaded": ScamDetector.session is not None,
        "onnx_path": str(ONNX_PATH),
    }


@app.post("/analyze", response_model=AnalyzeResponse)
def analyze(payload: AnalyzeRequest) -> AnalyzeResponse:
    try:
        result = ScamDetector.predict(payload.text)
        return AnalyzeResponse(**result)
    except Exception as exc:
        log.exception("Analyze failure")
        raise HTTPException(status_code=500, detail=str(exc)) from exc