scripts/eda_universalcefr.py

"""EDA of UniversalCEFR subsets — produces the decision report for the M1 training mix.

Discovers the org's datasets on the Hub, computes per-subset statistics
(levels, granularity, production category, licenses, text lengths) and writes
a markdown report that ADR 0003 references to fix the training mix.

Usage (deps live in the "data" group, kept out of the runtime image):
    uv run --group data python scripts/eda_universalcefr.py                  # English only
    uv run --group data python scripts/eda_universalcefr.py --langs en fr de
    uv run --group data python scripts/eda_universalcefr.py --langs all --save-raw

The report answers four questions:
  1. How big is the *reference, document-level* pool per language (the M1 use case)?
  2. What does cross-lingual training add (reference rows across languages)?
  3. Which subsets are learner production (excluded from the reading classifier)?
  4. Which licenses apply (kept in sync with ADR 0003 / README)?
"""

import argparse
import statistics
import sys
from collections import Counter
from dataclasses import dataclass, field
from datetime import UTC, datetime
from pathlib import Path

from datasets import load_dataset
from huggingface_hub import HfApi

ORG = "UniversalCEFR"
CANONICAL_LEVELS = ("A1", "A2", "B1", "B2", "C1", "C2")
# Languages covered by the org (suffix convention: <name>_<iso639-1>)
KNOWN_LANGS = {"ar", "cs", "cy", "de", "en", "es", "et", "fr", "hi", "it", "nl", "pt", "ru"}
DEFAULT_REPORT = Path("docs/evals/m1_data_eda.md")
RAW_DIR = Path("data/raw")


@dataclass
class SubsetStats:
    dataset_id: str
    lang: str
    n_rows: int = 0
    levels: Counter = field(default_factory=Counter)
    formats: Counter = field(default_factory=Counter)
    categories: Counter = field(default_factory=Counter)
    licenses: Counter = field(default_factory=Counter)
    word_counts: list[int] = field(default_factory=list)
    error: str | None = None

    @property
    def odd_levels(self) -> dict[str, int]:
        """Labels outside the canonical six (A1+, bare A/B, unlabeled...)."""
        return {lvl: n for lvl, n in self.levels.items() if lvl not in CANONICAL_LEVELS}

    def words_summary(self) -> str:
        if not self.word_counts:
            return "—"
        median = int(statistics.median(self.word_counts))
        if len(self.word_counts) >= 10:
            deciles = statistics.quantiles(self.word_counts, n=10)
            return f"{median} (p10={int(deciles[0])}, p90={int(deciles[-1])})"
        return str(median)


def lang_of(dataset_id: str) -> str | None:
    suffix = dataset_id.rsplit("_", 1)[-1].lower()
    return suffix if suffix in KNOWN_LANGS else None


def discover_datasets(langs: set[str]) -> list[tuple[str, str]]:
    """Return (dataset_id, lang) pairs from the org matching the requested languages."""
    api = HfApi()
    pairs: list[tuple[str, str]] = []
    for info in api.list_datasets(author=ORG):
        lang = lang_of(info.id)
        if lang and ("all" in langs or lang in langs):
            pairs.append((info.id, lang))
    return sorted(pairs)


def analyze(dataset_id: str, lang: str, save_raw: bool) -> SubsetStats:
    stats = SubsetStats(dataset_id=dataset_id, lang=lang)
    try:
        dataset = load_dataset(dataset_id, split="train")
    except Exception as exc:  # report and continue: the EDA must not die mid-run
        stats.error = f"{type(exc).__name__}: {exc}"
        return stats

    stats.n_rows = len(dataset)
    columns = dataset.column_names

    def count(column: str) -> Counter:
        if column not in columns:
            return Counter({"<missing column>": stats.n_rows})
        return Counter(str(value).strip() for value in dataset[column])

    stats.levels = count("cefr_level")
    stats.formats = count("format")
    stats.categories = count("category")
    stats.licenses = count("license")
    if "text" in columns:
        stats.word_counts = [len(str(text).split()) for text in dataset["text"]]

    if save_raw:
        RAW_DIR.mkdir(parents=True, exist_ok=True)
        dataset.to_parquet(RAW_DIR / f"{dataset_id.split('/')[-1]}.parquet")
    return stats


def _level_row(levels: Counter) -> str:
    cells = " | ".join(str(levels.get(lvl, 0)) for lvl in CANONICAL_LEVELS)
    odd = sum(n for lvl, n in levels.items() if lvl not in CANONICAL_LEVELS)
    return f"{cells} | {odd}"


def render_report(all_stats: list[SubsetStats], primary_lang: str, command: str) -> str:
    ok = [s for s in all_stats if s.error is None]
    failed = [s for s in all_stats if s.error is not None]
    lines: list[str] = [
        "# M1 data EDA — UniversalCEFR",
        "",
        f"Generated: {datetime.now(UTC).isoformat(timespec='seconds')}",
        f"Command: `{command}`",
        "",
        "## Subsets overview",
        "",
        "| dataset | lang | rows | categories | formats | licenses | words: median (p10, p90) |",
        "|---|---|---:|---|---|---|---|",
    ]
    for s in ok:
        lines.append(
            f"| `{s.dataset_id}` | {s.lang} | {s.n_rows} "
            f"| {dict(s.categories)} | {dict(s.formats)} | {dict(s.licenses)} "
            f"| {s.words_summary()} |"
        )

    lines += [
        "",
        "## Level distribution per subset",
        "",
        "| dataset | " + " | ".join(CANONICAL_LEVELS) + " | odd labels |",
        "|---|" + "---:|" * (len(CANONICAL_LEVELS) + 1),
    ]
    lines += [f"| `{s.dataset_id}` | {_level_row(s.levels)} |" for s in ok]
    odd_details = {s.dataset_id: s.odd_levels for s in ok if s.odd_levels}
    if odd_details:
        lines += ["", f"Odd labels detail: `{odd_details}`"]

    # The number M1 actually depends on: reference rows per (lang, format) and per level.
    ref = [s for s in ok if s.categories.get("reference", 0) > 0]
    lines += [
        "",
        "## Reference pool (the M1 reading-classifier candidates)",
        "",
        "Subsets whose `category` includes `reference`, i.e. texts written *for* "
        "learners rather than *by* them. Learner-production subsets are the M3 "
        "candidates (grading learner writing), not M1 training data.",
        "",
        "| lang | reference rows | from subsets |",
        "|---|---:|---|",
    ]
    by_lang: dict[str, list[SubsetStats]] = {}
    for s in ref:
        by_lang.setdefault(s.lang, []).append(s)
    for lang in sorted(by_lang):
        subsets = by_lang[lang]
        total = sum(s.categories.get("reference", 0) for s in subsets)
        names = ", ".join(f"`{s.dataset_id.split('/')[-1]}`" for s in subsets)
        lines.append(f"| {lang} | {total} | {names} |")

    primary = [s for s in ref if s.lang == primary_lang]
    if primary:
        pooled: Counter = Counter()
        for s in primary:
            pooled.update(s.levels)
        lines += [
            "",
            f"Pooled level distribution for `{primary_lang}` reference subsets "
            "(class balance check):",
            "",
            "| " + " | ".join(CANONICAL_LEVELS) + " | odd |",
            "|" + "---:|" * (len(CANONICAL_LEVELS) + 1),
            f"| {_level_row(pooled)} |",
        ]

    if failed:
        lines += ["", "## Failed subsets", ""]
        lines += [f"- `{s.dataset_id}` — {s.error}" for s in failed]

    lines += [
        "",
        "---",
        "Notes: word counts use whitespace tokenisation (approximate for ar/hi). "
        "Licenses are aggregated from the per-row `license` field; decisions and "
        "exclusions are recorded in `docs/adr/0003-datasets-and-licensing.md`.",
        "",
    ]
    return "\n".join(lines)


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--langs", nargs="+", default=["en"], help="ISO codes, or 'all'")
    parser.add_argument("--datasets", nargs="+", default=None, help="explicit ids (skip discovery)")
    parser.add_argument("--report", type=Path, default=DEFAULT_REPORT)
    parser.add_argument(
        "--save-raw", action="store_true", help="save subsets to data/raw/*.parquet"
    )
    args = parser.parse_args()

    if args.datasets:
        pairs = [(d, lang_of(d) or "?") for d in args.datasets]
    else:
        pairs = discover_datasets(set(args.langs))
    if not pairs:
        sys.exit("No matching datasets found.")

    print(f"Analysing {len(pairs)} subsets: {[d for d, _ in pairs]}\n")
    all_stats = []
    for dataset_id, lang in pairs:
        print(f"-> {dataset_id} ...", flush=True)
        all_stats.append(analyze(dataset_id, lang, save_raw=args.save_raw))

    report = render_report(all_stats, primary_lang=args.langs[0], command=" ".join(sys.argv))
    args.report.parent.mkdir(parents=True, exist_ok=True)
    args.report.write_text(report, encoding="utf-8")
    print(f"\n{report}")
    print(f"Report written to {args.report}")


if __name__ == "__main__":
    main()