app/classifier.py

"""Zero-shot complaint classifier.

Zero-shot (NLI) is what makes runtime labels possible: the candidate labels are
an *input* to every inference call, so adding "Street Lighting" in the UI takes
effect on the very next complaint -- no retraining, no restart.

If transformers/torch are unavailable or the model cannot be downloaded, a
keyword-overlap fallback keeps the app usable; every response reports which
engine produced it.
"""
from __future__ import annotations

import logging
import math
import re
import threading
import time
from typing import Any

from .config import (
    CONFIDENCE_THRESHOLD,
    HYPOTHESIS_TEMPLATE,
    LEXICAL_WEIGHT,
    MODEL_NAME,
)

log = logging.getLogger("classifier")

# Two locks on purpose: _load_lock is held for the whole (possibly very slow)
# model download, so inference must not wait on it -- a reload would otherwise
# stall every live request. _infer_lock only guards the pipeline call itself.
_load_lock = threading.Lock()
_infer_lock = threading.Lock()
_pipe: Any = None
_state: dict[str, Any] = {
    "status": "not_loaded",   # not_loaded | loading | ready | failed
    "engine": "none",
    "model": MODEL_NAME,
    "error": "",
    "load_seconds": 0.0,
}

# Candidate phrases longer than this are trimmed -- NLI hypotheses work best short.
MAX_CANDIDATE_WORDS = 16
_WORD_RE = re.compile(r"[a-z0-9']+")
_STOPWORDS = {
    "the", "a", "an", "is", "are", "was", "were", "be", "been", "being", "to", "of",
    "in", "on", "at", "for", "with", "and", "or", "not", "no", "my", "our", "i", "we",
    "it", "this", "that", "there", "here", "from", "by", "as", "has", "have", "had",
    "do", "does", "did", "but", "so", "very", "please", "sir", "madam", "any", "some",
}


def state() -> dict[str, Any]:
    return dict(_state)


def is_ready() -> bool:
    return _state["status"] == "ready"


def load_model(force: bool = False) -> dict[str, Any]:
    """Load the zero-shot pipeline. Safe to call repeatedly and concurrently."""
    global _pipe
    with _load_lock:
        if _pipe is not None and not force:
            return state()
        _state.update(status="loading", error="")
        started = time.perf_counter()
        try:
            from transformers import pipeline as hf_pipeline

            try:
                import torch

                torch.set_num_threads(max(1, (torch.get_num_threads() or 2) - 1))
            except Exception:
                pass

            _pipe = hf_pipeline(
                "zero-shot-classification",
                model=MODEL_NAME,
                device=-1,  # CPU; small model, fine for real-time single requests
            )
            _state.update(
                status="ready",
                engine=f"zero-shot:{MODEL_NAME}",
                load_seconds=round(time.perf_counter() - started, 2),
                error="",
            )
            log.info("Loaded %s in %.2fs", MODEL_NAME, _state["load_seconds"])
        except Exception as exc:
            _pipe = None
            _state.update(
                status="failed",
                engine="keyword-fallback",
                error=f"{type(exc).__name__}: {exc}",
                load_seconds=round(time.perf_counter() - started, 2),
            )
            log.warning("Model load failed (%s); using keyword fallback.", exc)
        return state()


def _candidate_for(label: dict[str, Any]) -> str:
    """The NLI hypothesis phrase for a label -- the label NAME, kept short.

    Feeding the description in here instead was measurably worse: on a 6-case
    benchmark, descriptions as candidates scored 0/6 against 3/6 for bare names.
    An NLI hypothesis wants a clean noun phrase ("Water Supply"), not a
    comma-separated keyword dump. Descriptions are still used -- as a lexical
    prior blended into the scores, see _blend().
    """
    name = (label.get("name") or "").strip()
    words = name.split()
    if len(words) > MAX_CANDIDATE_WORDS:
        name = " ".join(words[:MAX_CANDIDATE_WORDS])
    return name


def _blend(
    model_scores: list[tuple[str, float]],
    text: str,
    labels: list[dict[str, Any]],
    weight: float,
    multi_label: bool = False,
) -> list[tuple[str, float]]:
    """Mix the model scores with a keyword prior built from label descriptions.

    The descriptions carry real domain signal ("burst pipe", "lamp post") that a
    small NLI model does not get from the label name alone. Blending keeps the
    model in charge while letting those hints break ties.

    Skipped entirely in multi-label mode: there each score is an INDEPENDENT
    probability, while the lexical prior is a distribution summing to 1 across
    labels. Mixing the two would systematically depress every score and the
    renormalisation below would destroy the independence the caller asked for.
    """
    if weight <= 0 or multi_label:
        return model_scores

    lexical = dict(_keyword_scores(text, labels))
    # A flat prior means no keyword matched -- blending it would only dilute.
    if max(lexical.values(), default=0.0) - min(lexical.values(), default=0.0) < 1e-6:
        return model_scores

    blended = [
        (name, (1.0 - weight) * score + weight * lexical.get(name, 0.0))
        for name, score in model_scores
    ]
    total = sum(s for _, s in blended)
    if total > 0:
        blended = [(n, s / total) for n, s in blended]
    return blended


def _tokens(text: str) -> set[str]:
    return {w for w in _WORD_RE.findall(text.lower()) if w not in _STOPWORDS and len(w) > 2}


def _keyword_scores(text: str, labels: list[dict[str, Any]]) -> list[tuple[str, float]]:
    """Fallback scorer: overlap between complaint tokens and label tokens."""
    text_tokens = _tokens(text)
    raw = []
    for label in labels:
        label_tokens = _tokens(f"{label.get('name','')} {label.get('description','')}")
        if not label_tokens:
            raw.append(0.0)
            continue
        hits = sum(1 for t in label_tokens if t in text_tokens)
        # partial credit for stem-ish prefix matches ("leaking" vs "leak")
        near = sum(
            1 for t in label_tokens
            if t not in text_tokens and any(t.startswith(u[:4]) or u.startswith(t[:4])
                                            for u in text_tokens if len(u) > 3)
        )
        raw.append(hits + 0.35 * near)

    if max(raw, default=0.0) <= 0:
        even = 1.0 / len(labels) if labels else 0.0
        return [(l["name"], even) for l in labels]

    exps = [math.exp(s) for s in raw]
    total = sum(exps)
    return sorted(
        ((l["name"], e / total) for l, e in zip(labels, exps)),
        key=lambda p: p[1],
        reverse=True,
    )


def classify(
    text: str,
    labels: list[dict[str, Any]],
    multi_label: bool = False,
    threshold: float | None = None,
) -> dict[str, Any]:
    """Classify `text` against `labels` (list of {name, description}).

    Returns predicted label, per-label scores, engine used and latency.
    """
    threshold = CONFIDENCE_THRESHOLD if threshold is None else threshold
    text = (text or "").strip()
    started = time.perf_counter()

    if not text:
        raise ValueError("Complaint text is empty.")
    if not labels:
        raise ValueError("No active labels configured. Add at least one label first.")

    if _pipe is None and _state["status"] == "not_loaded":
        load_model()

    pipe = _pipe  # local ref: a concurrent reload must not swap this mid-call
    if pipe is not None:
        # Candidates are label names. Stored names are UNIQUE, but an ad-hoc
        # label list from the API may repeat, and the pipeline dislikes dupes.
        candidates = list(dict.fromkeys(_candidate_for(l) for l in labels if _candidate_for(l)))
        try:
            with _infer_lock:
                out = pipe(
                    text,
                    candidate_labels=candidates,
                    hypothesis_template=HYPOTHESIS_TEMPLATE,
                    multi_label=multi_label,
                )
            scored = [(lbl, float(score))
                      for lbl, score in zip(out["labels"], out["scores"])]
            scored = _blend(scored, text, labels, LEXICAL_WEIGHT, multi_label)
            engine = _state["engine"]
        except Exception as exc:
            log.warning("Inference failed (%s); falling back to keywords.", exc)
            scored = _keyword_scores(text, labels)
            engine = "keyword-fallback (inference error)"
    else:
        scored = _keyword_scores(text, labels)
        engine = "keyword-fallback"

    scored.sort(key=lambda p: p[1], reverse=True)
    top_label, top_score = scored[0]
    took_ms = int((time.perf_counter() - started) * 1000)

    return {
        "predicted_label": top_label,
        "confidence": round(top_score, 4),
        "confident": top_score >= threshold,
        "threshold": threshold,
        "scores": [{"label": n, "score": round(s, 4)} for n, s in scored],
        "engine": engine,
        "multi_label": multi_label,
        "took_ms": took_ms,
    }