src/tutor/ml/cefr/preprocessing.py

"""Label normalisation and passage chunking for the CEFR classifier (ADR 0003).

Lives in the runtime package (not training/) on purpose: inference must apply
the *exact* same preprocessing as training (no train/serve skew), and these
functions are dependency-free so they run in CI without the train group.
"""

import re
from dataclasses import dataclass

from tutor.domain.models import CEFRLevel

CANONICAL_LEVELS: tuple[str, ...] = tuple(level.value for level in CEFRLevel)

# Evidence for the drop list: docs/evals/m1_data_eda_all.md ("odd labels" detail).
# Bare macro levels (A/B/C) are ambiguous between two canonical levels -> dropped.
_DROP_LABELS = {"", "NA", "N/A", "EMPTY", "UNRATED", "UNASSESSABLE", "A", "B", "C"}

_SENTENCE_BOUNDARY = re.compile(r"(?<=[.!?])\s+")


def normalize_level(raw: str | None) -> str | None:
    """Map a raw label to a canonical level, or ``None`` if the row must be dropped.

    "X+" maps to "X": under an ordinal reading, X+ sits between X and the next
    level, so flooring is the conservative choice (keeps 114 elg_en and 1,505
    elg_nl rows).
    """
    if raw is None:
        return None
    label = raw.strip().upper()
    if label in _DROP_LABELS:
        return None
    if label.endswith("+"):
        label = label[:-1]
    return label if label in CANONICAL_LEVELS else None


@dataclass(frozen=True)
class Passage:
    """The training/inference unit: a sentence, or a chunk of a longer text."""

    text: str
    level: str
    lang: str
    corpus: str
    doc_id: str
    source_format: str


def split_sentences(text: str) -> list[str]:
    """Naive sentence split on terminal punctuation; good enough for packing."""
    return [s for s in _SENTENCE_BOUNDARY.split(text.strip()) if s]


def _hard_split(sentence: str, max_words: int) -> list[str]:
    """Last resort for a single 'sentence' longer than max_words (lists, no punctuation)."""
    words = sentence.split()
    return [" ".join(words[i : i + max_words]) for i in range(0, len(words), max_words)]


def chunk_text(
    text: str,
    *,
    target_words: int = 200,
    max_words: int = 300,
    min_tail_words: int = 50,
) -> list[str]:
    """Greedy sentence packing into ~target_words chunks (never above max_words).

    A short final chunk (< min_tail_words) is merged into the previous one so no
    chunk is uninformatively small. Texts already within max_words pass through.
    """
    stripped = text.strip()
    if not stripped:
        return []
    if len(stripped.split()) <= max_words:
        return [stripped]

    sentences: list[str] = []
    for sentence in split_sentences(stripped):
        if len(sentence.split()) > max_words:
            sentences.extend(_hard_split(sentence, max_words))
        else:
            sentences.append(sentence)

    chunks: list[list[str]] = []
    current: list[str] = []
    current_words = 0
    for sentence in sentences:
        n_words = len(sentence.split())
        if current and current_words + n_words > max_words:
            chunks.append(current)
            current, current_words = [], 0
        current.append(sentence)
        current_words += n_words
        if current_words >= target_words:
            chunks.append(current)
            current, current_words = [], 0
    if current:
        if chunks and current_words < min_tail_words:
            chunks[-1].extend(current)
        else:
            chunks.append(current)

    return [" ".join(chunk) for chunk in chunks]


def passages_from_record(
    *,
    text: str | None,
    level_raw: str | None,
    lang: str,
    corpus: str,
    doc_id: str,
    source_format: str,
    chunking: bool = True,
    target_words: int = 200,
    max_words: int = 300,
) -> list[Passage]:
    """Apply the full ADR 0003 record policy: label mapping, dropping, chunking.

    Sentence-level rows always pass through unchunked. Longer formats are
    chunked (chunks inherit the document label: weak labels, accepted and
    documented) unless ``chunking=False`` (the truncation experiment arm).
    """
    level = normalize_level(level_raw)
    if level is None or not text or not text.strip():
        return []

    def passage(chunk: str) -> Passage:
        return Passage(
            text=chunk,
            level=level,
            lang=lang,
            corpus=corpus,
            doc_id=doc_id,
            source_format=source_format,
        )

    if not chunking or source_format == "sentence-level":
        return [passage(text.strip())]
    return [
        passage(chunk) for chunk in chunk_text(text, target_words=target_words, max_words=max_words)
    ]