training/data_loading.py

"""Load UniversalCEFR subsets into Passages (network + `datasets` dependency).

Sibling module of the training scripts (run them from the repo root). The
record policy it applies — labels, dropping, chunking — lives in
`src/tutor/ml/cefr/preprocessing.py`, shared byte-for-byte with inference.
"""

from datasets import load_dataset

from tutor.ml.cefr.preprocessing import Passage, passages_from_record


def load_passages(
    subsets: list[str],
    *,
    chunking: bool,
    target_words: int,
    max_words: int,
) -> list[Passage]:
    passages: list[Passage] = []
    for subset in subsets:
        dataset = load_dataset(subset, split="train")
        corpus = subset.split("/")[-1]
        default_lang = corpus.rsplit("_", 1)[-1]
        before = len(passages)
        for index, row in enumerate(dataset):
            passages.extend(
                passages_from_record(
                    text=str(row.get("text") or ""),
                    level_raw=row.get("cefr_level"),
                    lang=str(row.get("lang") or default_lang),
                    corpus=corpus,
                    doc_id=f"{corpus}:{index}",
                    source_format=str(row.get("format") or "document-level"),
                    chunking=chunking,
                    target_words=target_words,
                    max_words=max_words,
                )
            )
        print(f"  {subset}: {len(dataset)} rows -> {len(passages) - before} passages")
    return passages