"""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, }