pipeline/fr_docs/classifier.py

"""FR-Docs classifier.

Two interchangeable engines behind one interface:

- EmbeddingClassifier (v0, ships today): embeds the document and the 30
  category anchors with a small sentence-transformer and picks the nearest
  anchor by cosine similarity. No training required.

- FineTunedClassifier (v1): loads the fine-tuned ModernBERT checkpoint
  produced by training/train.py. Drop-in replacement — same predict() API —
  so the service code never changes when you swap engines.
"""

from __future__ import annotations

from dataclasses import dataclass, asdict

import numpy as np

from .taxonomy import Category, load_taxonomy

# Confidence below which we return the parent group instead of the leaf
LEAF_CONFIDENCE_FLOOR = 0.35
# Margin over runner-up below which we flag the prediction as ambiguous
AMBIGUITY_MARGIN = 0.03

EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"


@dataclass
class Prediction:
    category_id: str
    category_label: str
    group_id: str
    group_label: str
    confidence: float
    ambiguous: bool
    top3: list[dict]
    engine: str

    def to_dict(self) -> dict:
        return asdict(self)


class EmbeddingClassifier:
    """Zero-training anchor-similarity classifier."""

    def __init__(self, model_name: str = EMBED_MODEL):
        from sentence_transformers import SentenceTransformer

        self.model = SentenceTransformer(model_name)
        self.categories: list[Category] = load_taxonomy()
        anchors = [f"{c.label}. {c.anchor}" for c in self.categories]
        self.anchor_vecs = self.model.encode(
            anchors, normalize_embeddings=True, show_progress_bar=False
        )

    def predict(self, text: str) -> Prediction:
        if not text.strip():
            other = next(c for c in self.categories if c.id == "other")
            return Prediction(
                category_id=other.id, category_label=other.label,
                group_id=other.group_id, group_label=other.group_label,
                confidence=0.0, ambiguous=True, top3=[], engine="embedding-v0",
            )
        doc_vec = self.model.encode(
            [text], normalize_embeddings=True, show_progress_bar=False
        )[0]
        sims = self.anchor_vecs @ doc_vec
        # Softmax over similarities (temperature sharpens the raw cosine range)
        probs = np.exp(sims / 0.05)
        probs /= probs.sum()

        order = np.argsort(probs)[::-1]
        best, second = order[0], order[1]
        cat = self.categories[best]
        confidence = float(probs[best])
        ambiguous = (confidence - float(probs[second])) < AMBIGUITY_MARGIN

        top3 = [
            {
                "category_id": self.categories[i].id,
                "label": self.categories[i].label,
                "confidence": round(float(probs[i]), 4),
            }
            for i in order[:3]
        ]
        return Prediction(
            category_id=cat.id,
            category_label=cat.label,
            group_id=cat.group_id,
            group_label=cat.group_label,
            confidence=round(confidence, 4),
            ambiguous=ambiguous or confidence < LEAF_CONFIDENCE_FLOOR,
            top3=top3,
            engine="embedding-v0",
        )


class LiteClassifier:
    """Dependency-free fallback: TF-IDF cosine similarity against anchors.

    Used automatically when sentence-transformers is not installed (e.g.
    restricted or edge environments). Lower accuracy than the embedding
    engine; same interface.
    """

    _token_re = None

    def __init__(self):
        import math
        import re

        self.math = math
        self._token_re = re.compile(r"[a-z]{2,}")
        self.categories: list[Category] = load_taxonomy()
        docs = [self._tokens(f"{c.label} {c.anchor}") for c in self.categories]
        # idf over anchor corpus
        df: dict[str, int] = {}
        for toks in docs:
            for t in set(toks):
                df[t] = df.get(t, 0) + 1
        n = len(docs)
        self.idf = {t: math.log((n + 1) / (d + 1)) + 1 for t, d in df.items()}
        self.anchor_vecs = [self._vec(toks) for toks in docs]

    def _tokens(self, text: str) -> list[str]:
        return self._token_re.findall(text.lower())

    def _vec(self, tokens: list[str]) -> dict[str, float]:
        tf: dict[str, float] = {}
        for t in tokens:
            tf[t] = tf.get(t, 0) + 1
        vec = {t: f * self.idf.get(t, 1.0) for t, f in tf.items()}
        norm = self.math.sqrt(sum(v * v for v in vec.values())) or 1.0
        return {t: v / norm for t, v in vec.items()}

    def predict(self, text: str) -> Prediction:
        if not text.strip():
            other = next(c for c in self.categories if c.id == "other")
            return Prediction(
                category_id=other.id, category_label=other.label,
                group_id=other.group_id, group_label=other.group_label,
                confidence=0.0, ambiguous=True, top3=[], engine="lite-tfidf",
            )
        doc_vec = self._vec(self._tokens(text))
        sims = np.array([
            sum(doc_vec.get(t, 0.0) * v for t, v in anchor.items())
            for anchor in self.anchor_vecs
        ])
        probs = np.exp(sims / 0.05)
        probs /= probs.sum()
        order = np.argsort(probs)[::-1]
        best, second = order[0], order[1]
        cat = self.categories[best]
        confidence = float(probs[best])
        ambiguous = (confidence - float(probs[second])) < AMBIGUITY_MARGIN
        top3 = [
            {
                "category_id": self.categories[i].id,
                "label": self.categories[i].label,
                "confidence": round(float(probs[i]), 4),
            }
            for i in order[:3]
        ]
        return Prediction(
            category_id=cat.id,
            category_label=cat.label,
            group_id=cat.group_id,
            group_label=cat.group_label,
            confidence=round(confidence, 4),
            ambiguous=ambiguous or confidence < LEAF_CONFIDENCE_FLOOR,
            top3=top3,
            engine="lite-tfidf",
        )


class SklearnClassifier:
    """v0.5: trained TF-IDF + LogisticRegression checkpoint (joblib)."""

    def __init__(self, model_path: str):
        import joblib

        self.pipeline = joblib.load(model_path)
        self.categories = {c.id: c for c in load_taxonomy()}
        self.class_order = list(self.pipeline.classes_)

    def predict(self, text: str) -> Prediction:
        if not text.strip():
            other = self.categories["other"]
            return Prediction(
                category_id=other.id, category_label=other.label,
                group_id=other.group_id, group_label=other.group_label,
                confidence=0.0, ambiguous=True, top3=[], engine="sklearn-v0.5",
            )
        probs = self.pipeline.predict_proba([text])[0]
        order = np.argsort(probs)[::-1]
        best, second = order[0], order[1]
        cat = self.categories[self.class_order[best]]
        confidence = float(probs[best])
        ambiguous = (confidence - float(probs[second])) < AMBIGUITY_MARGIN
        top3 = [
            {
                "category_id": self.class_order[i],
                "label": self.categories[self.class_order[i]].label,
                "confidence": round(float(probs[i]), 4),
            }
            for i in order[:3]
        ]
        return Prediction(
            category_id=cat.id,
            category_label=cat.label,
            group_id=cat.group_id,
            group_label=cat.group_label,
            confidence=round(confidence, 4),
            ambiguous=ambiguous or confidence < LEAF_CONFIDENCE_FLOOR,
            top3=top3,
            engine="sklearn-v0.5",
        )


class FineTunedClassifier:
    """Loads the ModernBERT checkpoint trained with training/train.py."""

    def __init__(self, checkpoint_dir: str):
        import torch
        from transformers import AutoModelForSequenceClassification, AutoTokenizer

        self.torch = torch
        self.tokenizer = AutoTokenizer.from_pretrained(checkpoint_dir)
        self.model = AutoModelForSequenceClassification.from_pretrained(checkpoint_dir)
        self.model.eval()
        self.categories = {c.id: c for c in load_taxonomy()}
        # Respect the architecture's real context limit (512 for DistilBERT,
        # 8192 for ModernBERT); never exceed what the checkpoint supports.
        model_limit = getattr(self.model.config, "max_position_embeddings", 512)
        tok_limit = self.tokenizer.model_max_length
        if tok_limit is None or tok_limit > 100_000:  # some tokenizers report a sentinel
            tok_limit = model_limit
        self.max_length = min(model_limit, tok_limit, 4096)

    def predict(self, text: str) -> Prediction:
        inputs = self.tokenizer(
            text, truncation=True, max_length=self.max_length, return_tensors="pt"
        )
        with self.torch.no_grad():
            logits = self.model(**inputs).logits[0]
        probs = logits.softmax(-1)
        order = probs.argsort(descending=True)
        best = order[0].item()
        cat_id = self.model.config.id2label[best]
        cat = self.categories[cat_id]
        confidence = float(probs[best])
        ambiguous = (confidence - float(probs[order[1]])) < AMBIGUITY_MARGIN
        top3 = [
            {
                "category_id": self.model.config.id2label[i.item()],
                "label": self.categories[self.model.config.id2label[i.item()]].label,
                "confidence": round(float(probs[i]), 4),
            }
            for i in order[:3]
        ]
        return Prediction(
            category_id=cat.id,
            category_label=cat.label,
            group_id=cat.group_id,
            group_label=cat.group_label,
            confidence=round(confidence, 4),
            ambiguous=ambiguous or confidence < LEAF_CONFIDENCE_FLOOR,
            top3=top3,
            engine="modernbert-v1",
        )


def get_classifier(checkpoint_dir: str | None = None, sklearn_path: str | None = None):
    """Factory, best available first:
    ModernBERT checkpoint > sklearn checkpoint > embedding engine > lite fallback.
    """
    import os
    from pathlib import Path

    if checkpoint_dir:
        return FineTunedClassifier(checkpoint_dir)
    sklearn_path = sklearn_path or os.environ.get("FR_DOCS_SKLEARN")
    if not sklearn_path:
        default = Path(__file__).parent.parent / "checkpoints" / "fr-docs-sklearn.joblib"
        if default.exists():
            sklearn_path = str(default)
    if sklearn_path:
        return SklearnClassifier(sklearn_path)
    try:
        return EmbeddingClassifier()
    except ImportError:
        return LiteClassifier()