src/tutor/ml/cefr/splitting.py

"""Document-level, stratified, deterministic split (ADR 0003).

The split is decided on *document ids*, never on passages: chunks inherit
their document's split, so chunk leakage between train and test is impossible
by construction (tests/test_cefr_splitting.py asserts it).
"""

import random
from collections import defaultdict

SPLITS = ("train", "val", "test")


def assign_splits(
    doc_strata: dict[str, str],
    *,
    ratios: tuple[float, float, float] = (0.8, 0.1, 0.1),
    seed: int = 13,
) -> dict[str, str]:
    """Map each doc_id to a split, stratified by its stratum (e.g. "corpus|level").

    Deterministic for a given (seed, strata) regardless of dict insertion order:
    each stratum is sorted then shuffled with its own seeded RNG. Allocation uses
    floor counts, so small strata (< ~1/ratio docs) contribute to train only —
    test/val never starve train of a rare (corpus, level) cell.
    """
    if abs(sum(ratios) - 1.0) > 1e-9:
        msg = f"ratios must sum to 1, got {ratios}"
        raise ValueError(msg)

    by_stratum: dict[str, list[str]] = defaultdict(list)
    for doc_id, stratum in doc_strata.items():
        by_stratum[stratum].append(doc_id)

    assignment: dict[str, str] = {}
    for stratum in sorted(by_stratum):
        docs = sorted(by_stratum[stratum])
        random.Random(f"{seed}:{stratum}").shuffle(docs)
        n_docs = len(docs)
        n_test = int(n_docs * ratios[2])
        n_val = int(n_docs * ratios[1])
        for index, doc_id in enumerate(docs):
            if index < n_test:
                assignment[doc_id] = "test"
            elif index < n_test + n_val:
                assignment[doc_id] = "val"
            else:
                assignment[doc_id] = "train"
    return assignment