tests/test_cefr_inference.py

import numpy as np
import pytest

from tutor.ml.cefr.inference import CEFRClassifier, build_onnx_feed, softmax


def test_softmax_rows_sum_to_one_and_is_stable() -> None:
    logits = np.array([[1.0, 2.0, 3.0, 4.0, 5.0, 6.0], [1000.0, 1000.0, 0.0, 0.0, 0.0, 0.0]])
    probs = softmax(logits)
    assert np.allclose(probs.sum(axis=-1), 1.0)
    assert probs[0].argmax() == 5
    assert probs[1][0] == pytest.approx(0.5)  # no overflow on large logits


def test_build_onnx_feed_handles_token_type_ids_and_rejects_unknown() -> None:
    ids = np.ones((2, 4), dtype=np.int64)
    mask = np.ones((2, 4), dtype=np.int64)
    feed = build_onnx_feed(["input_ids", "attention_mask", "token_type_ids"], ids, mask)
    assert set(feed) == {"input_ids", "attention_mask", "token_type_ids"}
    assert feed["token_type_ids"].sum() == 0
    with pytest.raises(ValueError, match="unexpected ONNX model input"):
        build_onnx_feed(["pixel_values"], ids, mask)


class _FakeClassifier(CEFRClassifier):
    """Real chunking + aggregation, canned model output."""

    def __init__(self, probs_row: list[float]) -> None:
        super().__init__(session=None, tokenizer=None, target_words=50, max_words=80)
        self._probs_row = probs_row
        self.seen_chunks: list[str] = []

    def _predict_probs(self, texts: list[str]) -> list[list[float]]:
        self.seen_chunks.extend(texts)
        return [self._probs_row for _ in texts]


def test_classify_text_chunks_and_aggregates() -> None:
    b2_row = [0.02, 0.02, 0.1, 0.8, 0.04, 0.02]
    classifier = _FakeClassifier(b2_row)
    long_text = ". ".join(" ".join(f"w{i}" for i in range(11)) + " end" for _ in range(30)) + "."

    prediction = classifier.classify_text(long_text)

    assert len(classifier.seen_chunks) > 1  # the real chunker actually ran
    assert prediction.n_chunks == len(classifier.seen_chunks)
    assert prediction.level == "B2"
    assert prediction.score == pytest.approx(2.88, abs=0.01)
    assert prediction.per_level["B2"] == pytest.approx(0.8)


def test_classify_text_rejects_empty_input() -> None:
    classifier = _FakeClassifier([1 / 6] * 6)
    with pytest.raises(ValueError, match="empty text"):
        classifier.classify_text("   ")