src/tutor/ml/cefr/aggregation.py

"""Chunk -> document aggregation by expected-rank mean (ADR 0003).

Shared by training evaluation and Space inference: a document's level is the
rounded mean of the expected CEFR rank of each of its chunks. The continuous
score is kept — it is useful both for metrics and for the UI ("strong B1").
"""

from collections.abc import Sequence

from tutor.ml.cefr.preprocessing import CANONICAL_LEVELS


def expected_rank(probs: Sequence[float]) -> float:
    """E[rank] of a probability distribution over the six canonical levels."""
    if len(probs) != len(CANONICAL_LEVELS):
        msg = f"expected {len(CANONICAL_LEVELS)} probabilities, got {len(probs)}"
        raise ValueError(msg)
    total = sum(probs)
    if total <= 0:
        msg = "probabilities must sum to a positive value"
        raise ValueError(msg)
    return sum(index * p for index, p in enumerate(probs)) / total


def aggregate_chunk_probs(chunk_probs: Sequence[Sequence[float]]) -> tuple[str, float]:
    """Aggregate per-chunk probability rows into (document level, continuous score)."""
    if not chunk_probs:
        msg = "cannot aggregate an empty list of chunk probabilities"
        raise ValueError(msg)
    score = sum(expected_rank(probs) for probs in chunk_probs) / len(chunk_probs)
    index = min(len(CANONICAL_LEVELS) - 1, max(0, int(score + 0.5)))  # round half up, clamp
    return CANONICAL_LEVELS[index], score