app/classifier.py

"""Model loader abstraction: one interface, two backends.

The service never imports torch/transformers directly. It asks
:func:`load_classifier` for an object exposing ``predict(texts) -> list[dict]``,
where each dict maps every label to a probability that sums to 1.0.

Two implementations satisfy that contract:

* :class:`StubClassifier` — pure-Python, deterministic, zero downloads. Scores
  text with a small hand-built emotion lexicon so the distribution is plausible
  and *stable* (the same sentence always yields the same probabilities), which
  is what makes offline demos, tests, and load tests meaningful. Selected when
  ``OFFLINE=1`` (the default).
* :class:`TransformersClassifier` — the real fine-tuned DistilBERT loaded once
  via a Hugging Face ``pipeline``. Selected when ``OFFLINE=0``.

Both are warmed up on construction so the first real request is not slow.
"""
from __future__ import annotations

import logging
import math
import re
from typing import Dict, List, Protocol

from app import ID2LABEL, LABELS
from app.config import Settings

logger = logging.getLogger(__name__)

Distribution = Dict[str, float]


class Classifier(Protocol):
    """Anything the service can serve predictions from."""

    backend: str

    def predict(self, texts: List[str]) -> List[Distribution]:
        """Return one ``{label: probability}`` dict per input text."""
        ...


def _softmax(scores: List[float]) -> List[float]:
    hi = max(scores)
    exps = [math.exp(s - hi) for s in scores]
    total = sum(exps)
    return [e / total for e in exps]


# --- Offline stub -----------------------------------------------------------

# A compact emotion lexicon. Not meant to rival the real model — it exists so the
# offline path produces a deterministic, label-aware distribution instead of a
# constant, which keeps demos and tests honest. Weights are deliberately modest
# so the softmax stays smooth rather than collapsing to a one-hot vector.
_LEXICON: Dict[str, Dict[str, float]] = {
    "sadness": {
        "sad": 2.2, "unhappy": 2.0, "cry": 2.0, "crying": 2.0, "lonely": 2.0,
        "depressed": 2.4, "miserable": 2.2, "down": 1.4, "hopeless": 2.2,
        "lost": 1.6, "hurt": 1.6, "grief": 2.4, "disappointed": 1.8, "empty": 1.8,
    },
    "joy": {
        "happy": 2.2, "joy": 2.4, "glad": 2.0, "great": 1.6, "wonderful": 2.0,
        "delighted": 2.2, "excited": 1.8, "smiling": 2.0, "grateful": 1.6,
        "love": 0.6, "amazing": 1.8, "fantastic": 2.0, "relieved": 1.6, "thrilled": 2.2,
    },
    "love": {
        "love": 2.4, "adore": 2.4, "beloved": 2.2, "affection": 2.2, "tender": 2.0,
        "caring": 1.8, "sweetheart": 2.2, "cherish": 2.2, "fond": 1.8,
        "romantic": 2.0, "devoted": 2.0,
    },
    "anger": {
        "angry": 2.4, "mad": 2.0, "furious": 2.6, "rage": 2.6, "hate": 2.2,
        "annoyed": 1.8, "irritated": 1.8, "outraged": 2.4, "resent": 2.0,
        "disgusted": 1.8, "betrayed": 2.0, "unfair": 1.6,
    },
    "fear": {
        "afraid": 2.4, "scared": 2.4, "fear": 2.4, "terrified": 2.6, "anxious": 2.2,
        "nervous": 2.0, "worried": 2.0, "panic": 2.4, "dread": 2.2, "frightened": 2.4,
        "shaking": 1.8, "uneasy": 1.8,
    },
    "surprise": {
        "surprised": 2.4, "shocked": 2.2, "amazed": 2.0, "astonished": 2.4,
        "unexpected": 2.0, "wow": 1.8, "suddenly": 1.4, "stunned": 2.2,
        "speechless": 2.0, "unbelievable": 1.8,
    },
}

_WORD_RE = re.compile(r"[a-z']+")


class StubClassifier:
    """Deterministic, network-free classifier driven by an emotion lexicon."""

    backend = "stub"

    def __init__(self) -> None:
        # Slight positive prior on the two dominant training classes (joy,
        # sadness) so empty / unknown text falls back to a realistic shape
        # rather than a flat uniform distribution.
        self._prior = {lab: 0.0 for lab in LABELS}
        self._prior["joy"] = 0.30
        self._prior["sadness"] = 0.25
        logger.info("StubClassifier ready (offline, no weights loaded)")

    def _score_one(self, text: str) -> Distribution:
        scores = dict(self._prior)
        for token in _WORD_RE.findall(text.lower()):
            for label, lex in _LEXICON.items():
                w = lex.get(token)
                if w:
                    scores[label] += w
        ordered = [scores[ID2LABEL[i]] for i in range(len(LABELS))]
        probs = _softmax(ordered)
        return {ID2LABEL[i]: probs[i] for i in range(len(LABELS))}

    def predict(self, texts: List[str]) -> List[Distribution]:
        return [self._score_one(t) for t in texts]


# --- Real model -------------------------------------------------------------

class TransformersClassifier:
    """The fine-tuned DistilBERT, loaded once via a transformers pipeline."""

    backend = "transformers"

    def __init__(self, model_id: str) -> None:
        # Imported lazily so the dependency is only required for a real run.
        from transformers import pipeline  # type: ignore

        logger.info("Loading model %s via transformers pipeline ...", model_id)
        self._pipe = pipeline(
            "text-classification",
            model=model_id,
            top_k=None,  # return the full distribution, not just the argmax
        )
        # Warm up so the first served request does not pay graph-build cost.
        self._pipe("warmup")
        logger.info("Model %s loaded and warmed up", model_id)

    def predict(self, texts: List[str]) -> List[Distribution]:
        raw = self._pipe(list(texts))
        # pipeline returns list[list[{label, score}]] when top_k=None.
        out: List[Distribution] = []
        for row in raw:
            dist = {item["label"]: float(item["score"]) for item in row}
            # Guarantee every canonical label is present and ordered.
            out.append({lab: dist.get(lab, 0.0) for lab in LABELS})
        return out


def load_classifier(settings: Settings) -> Classifier:
    """Build the classifier the settings ask for.

    Falls back to the stub if a real load is requested but the ML stack is not
    installed, so the service still boots (degraded) instead of crash-looping.
    """
    if settings.offline:
        return StubClassifier()
    try:
        return TransformersClassifier(settings.model_id)
    except Exception:  # pragma: no cover - exercised only with the real stack
        logger.exception(
            "Failed to load real model %s; falling back to offline stub. "
            "Install the 'ml' extra and ensure the weights are reachable.",
            settings.model_id,
        )
        return StubClassifier()