src/tutor/ml/cefr/inference.py

"""ONNX CPU inference service for the CEFR classifier.

Runtime-side by design: depends only on onnxruntime + tokenizers + numpy — no
torch, no transformers — so the Space image stays lean. The preprocessing is
the *same code* training used (chunker and aggregation imported from this
package), with parameters frozen by the export script in ``meta.json``: zero
train/serve skew by construction.
"""

import json
import os
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import onnxruntime as ort
from tokenizers import Tokenizer

from tutor.config import Settings
from tutor.ml.cefr.aggregation import aggregate_chunk_probs, expected_rank
from tutor.ml.cefr.preprocessing import CANONICAL_LEVELS, chunk_text

_ARTIFACT_FILES = ["model.onnx", "tokenizer.json", "meta.json"]


@dataclass(frozen=True)
class CEFRPrediction:
    level: str
    score: float  # continuous expected rank (0=A1 ... 5=C2), e.g. 2.7 = "strong B1"
    per_level: dict[str, float]  # mean probability per level across chunks
    n_chunks: int


def softmax(logits: np.ndarray) -> np.ndarray:
    """Row-wise, numerically stable softmax."""
    shifted = logits - logits.max(axis=-1, keepdims=True)
    exp = np.exp(shifted)
    return exp / exp.sum(axis=-1, keepdims=True)


def build_onnx_feed(
    input_names: list[str],
    input_ids: np.ndarray,
    attention_mask: np.ndarray,
) -> dict[str, np.ndarray]:
    """Map our two tensors onto whatever inputs the exported graph declares."""
    feed: dict[str, np.ndarray] = {}
    for name in input_names:
        if name == "input_ids":
            feed[name] = input_ids
        elif name == "attention_mask":
            feed[name] = attention_mask
        elif name == "token_type_ids":
            feed[name] = np.zeros_like(input_ids)
        else:
            msg = f"unexpected ONNX model input '{name}'"
            raise ValueError(msg)
    return feed


class CEFRClassifier:
    """Chunk -> ONNX forward -> expected-rank aggregation, end to end."""

    def __init__(
        self,
        session: "ort.InferenceSession | None",
        tokenizer: "Tokenizer | None",
        *,
        max_length: int = 512,
        target_words: int = 200,
        max_words: int = 300,
        batch_size: int = 16,
    ) -> None:
        # session/tokenizer may be None in unit tests that override _predict_probs.
        self._session = session
        self._tokenizer = tokenizer
        self.max_length = max_length
        self.target_words = target_words
        self.max_words = max_words
        self.batch_size = batch_size
        self._input_names = (
            [item.name for item in session.get_inputs()] if session is not None else []
        )

    @classmethod
    def from_dir(cls, artifact_dir: Path | str) -> "CEFRClassifier":
        artifact_dir = Path(artifact_dir)
        missing = [name for name in _ARTIFACT_FILES if not (artifact_dir / name).exists()]
        if missing:
            msg = f"CEFR artifact dir {artifact_dir} is missing {missing}"
            raise FileNotFoundError(msg)

        meta = json.loads((artifact_dir / "meta.json").read_text(encoding="utf-8"))
        if list(meta["levels"]) != list(CANONICAL_LEVELS):
            msg = f"artifact level order {meta['levels']} != canonical {CANONICAL_LEVELS}"
            raise ValueError(msg)

        tokenizer = Tokenizer.from_file(str(artifact_dir / "tokenizer.json"))
        tokenizer.enable_truncation(max_length=meta["max_length"])
        tokenizer.enable_padding(pad_id=meta["pad_id"], pad_token=meta["pad_token"])

        options = ort.SessionOptions()
        session = ort.InferenceSession(
            str(artifact_dir / "model.onnx"),
            sess_options=options,
            providers=["CPUExecutionProvider"],
        )
        return cls(
            session,
            tokenizer,
            max_length=meta["max_length"],
            target_words=meta["target_words"],
            max_words=meta["max_words"],
        )

    @classmethod
    def from_hub(cls, repo_id: str) -> "CEFRClassifier":
        """Download the artifact from a HF model repo (HF_TOKEN env honoured)."""
        from huggingface_hub import snapshot_download  # runtime dep, lazy: keeps import cheap

        local_dir = snapshot_download(repo_id, allow_patterns=_ARTIFACT_FILES)
        return cls.from_dir(local_dir)

    def _predict_probs(self, texts: list[str]) -> list[list[float]]:
        """One probability row (canonical level order) per text."""
        probs: list[list[float]] = []
        for start in range(0, len(texts), self.batch_size):
            encodings = self._tokenizer.encode_batch(texts[start : start + self.batch_size])
            input_ids = np.array([e.ids for e in encodings], dtype=np.int64)
            attention_mask = np.array([e.attention_mask for e in encodings], dtype=np.int64)
            feed = build_onnx_feed(self._input_names, input_ids, attention_mask)
            logits = self._session.run(None, feed)[0]
            probs.extend(softmax(np.asarray(logits, dtype=np.float64)).tolist())
        return probs

    def classify_text(self, text: str) -> CEFRPrediction:
        chunks = chunk_text(text, target_words=self.target_words, max_words=self.max_words)
        if not chunks:
            msg = "cannot classify an empty text"
            raise ValueError(msg)
        probs = self._predict_probs(chunks)
        level, score = aggregate_chunk_probs(probs)
        mean_probs = np.asarray(probs).mean(axis=0)
        per_level = {lvl: float(p) for lvl, p in zip(CANONICAL_LEVELS, mean_probs, strict=True)}
        return CEFRPrediction(level=level, score=score, per_level=per_level, n_chunks=len(chunks))

    def expected_rank_of(self, probs: list[float]) -> float:
        """Convenience passthrough, mostly for diagnostics."""
        return expected_rank(probs)


def create_cefr_classifier(settings: Settings) -> CEFRClassifier | None:
    """Resolve the classifier from settings; None when not configured.

    Local path takes precedence (dev); otherwise a HF model repo id
    (the Space path — HF_TOKEN is read from the environment for private repos).
    """
    if settings.cefr_model_path:
        return CEFRClassifier.from_dir(settings.cefr_model_path)
    if settings.cefr_model_id:
        if os.environ.get("HF_TOKEN") is None and settings.app_env == "prod":
            # Public repos work without a token; this is only a hint in logs.
            print("CEFR model: downloading from the Hub without HF_TOKEN (public repo assumed)")
        return CEFRClassifier.from_hub(settings.cefr_model_id)
    return None