blizzarman/polyglot-tutor

1"""EDA of UniversalCEFR subsets — produces the decision report for the M1 training mix.2 3Discovers the org's datasets on the Hub, computes per-subset statistics4(levels, granularity, production category, licenses, text lengths) and writes5a markdown report that ADR 0003 references to fix the training mix.6 7Usage (deps live in the "data" group, kept out of the runtime image):8    uv run --group data python scripts/eda_universalcefr.py                  # English only9    uv run --group data python scripts/eda_universalcefr.py --langs en fr de10    uv run --group data python scripts/eda_universalcefr.py --langs all --save-raw11 12The report answers four questions:13  1. How big is the *reference, document-level* pool per language (the M1 use case)?14  2. What does cross-lingual training add (reference rows across languages)?15  3. Which subsets are learner production (excluded from the reading classifier)?16  4. Which licenses apply (kept in sync with ADR 0003 / README)?17"""18 19import argparse20import statistics21import sys22from collections import Counter23from dataclasses import dataclass, field24from datetime import UTC, datetime25from pathlib import Path26 27from datasets import load_dataset28from huggingface_hub import HfApi29 30ORG = "UniversalCEFR"31CANONICAL_LEVELS = ("A1", "A2", "B1", "B2", "C1", "C2")32# Languages covered by the org (suffix convention: <name>_<iso639-1>)33KNOWN_LANGS = {"ar", "cs", "cy", "de", "en", "es", "et", "fr", "hi", "it", "nl", "pt", "ru"}34DEFAULT_REPORT = Path("docs/evals/m1_data_eda.md")35RAW_DIR = Path("data/raw")36 37 38@dataclass39class SubsetStats:40    dataset_id: str41    lang: str42    n_rows: int = 043    levels: Counter = field(default_factory=Counter)44    formats: Counter = field(default_factory=Counter)45    categories: Counter = field(default_factory=Counter)46    licenses: Counter = field(default_factory=Counter)47    word_counts: list[int] = field(default_factory=list)48    error: str | None = None49 50    @property51    def odd_levels(self) -> dict[str, int]:52        """Labels outside the canonical six (A1+, bare A/B, unlabeled...)."""53        return {lvl: n for lvl, n in self.levels.items() if lvl not in CANONICAL_LEVELS}54 55    def words_summary(self) -> str:56        if not self.word_counts:57            return "—"58        median = int(statistics.median(self.word_counts))59        if len(self.word_counts) >= 10:60            deciles = statistics.quantiles(self.word_counts, n=10)61            return f"{median} (p10={int(deciles[0])}, p90={int(deciles[-1])})"62        return str(median)63 64 65def lang_of(dataset_id: str) -> str | None:66    suffix = dataset_id.rsplit("_", 1)[-1].lower()67    return suffix if suffix in KNOWN_LANGS else None68 69 70def discover_datasets(langs: set[str]) -> list[tuple[str, str]]:71    """Return (dataset_id, lang) pairs from the org matching the requested languages."""72    api = HfApi()73    pairs: list[tuple[str, str]] = []74    for info in api.list_datasets(author=ORG):75        lang = lang_of(info.id)76        if lang and ("all" in langs or lang in langs):77            pairs.append((info.id, lang))78    return sorted(pairs)79 80 81def analyze(dataset_id: str, lang: str, save_raw: bool) -> SubsetStats:82    stats = SubsetStats(dataset_id=dataset_id, lang=lang)83    try:84        dataset = load_dataset(dataset_id, split="train")85    except Exception as exc:  # report and continue: the EDA must not die mid-run86        stats.error = f"{type(exc).__name__}: {exc}"87        return stats88 89    stats.n_rows = len(dataset)90    columns = dataset.column_names91 92    def count(column: str) -> Counter:93        if column not in columns:94            return Counter({"<missing column>": stats.n_rows})95        return Counter(str(value).strip() for value in dataset[column])96 97    stats.levels = count("cefr_level")98    stats.formats = count("format")99    stats.categories = count("category")100    stats.licenses = count("license")101    if "text" in columns:102        stats.word_counts = [len(str(text).split()) for text in dataset["text"]]103 104    if save_raw:105        RAW_DIR.mkdir(parents=True, exist_ok=True)106        dataset.to_parquet(RAW_DIR / f"{dataset_id.split('/')[-1]}.parquet")107    return stats108 109 110def _level_row(levels: Counter) -> str:111    cells = " | ".join(str(levels.get(lvl, 0)) for lvl in CANONICAL_LEVELS)112    odd = sum(n for lvl, n in levels.items() if lvl not in CANONICAL_LEVELS)113    return f"{cells} | {odd}"114 115 116def render_report(all_stats: list[SubsetStats], primary_lang: str, command: str) -> str:117    ok = [s for s in all_stats if s.error is None]118    failed = [s for s in all_stats if s.error is not None]119    lines: list[str] = [120        "# M1 data EDA — UniversalCEFR",121        "",122        f"Generated: {datetime.now(UTC).isoformat(timespec='seconds')}",123        f"Command: `{command}`",124        "",125        "## Subsets overview",126        "",127        "| dataset | lang | rows | categories | formats | licenses | words: median (p10, p90) |",128        "|---|---|---:|---|---|---|---|",129    ]130    for s in ok:131        lines.append(132            f"| `{s.dataset_id}` | {s.lang} | {s.n_rows} "133            f"| {dict(s.categories)} | {dict(s.formats)} | {dict(s.licenses)} "134            f"| {s.words_summary()} |"135        )136 137    lines += [138        "",139        "## Level distribution per subset",140        "",141        "| dataset | " + " | ".join(CANONICAL_LEVELS) + " | odd labels |",142        "|---|" + "---:|" * (len(CANONICAL_LEVELS) + 1),143    ]144    lines += [f"| `{s.dataset_id}` | {_level_row(s.levels)} |" for s in ok]145    odd_details = {s.dataset_id: s.odd_levels for s in ok if s.odd_levels}146    if odd_details:147        lines += ["", f"Odd labels detail: `{odd_details}`"]148 149    # The number M1 actually depends on: reference rows per (lang, format) and per level.150    ref = [s for s in ok if s.categories.get("reference", 0) > 0]151    lines += [152        "",153        "## Reference pool (the M1 reading-classifier candidates)",154        "",155        "Subsets whose `category` includes `reference`, i.e. texts written *for* "156        "learners rather than *by* them. Learner-production subsets are the M3 "157        "candidates (grading learner writing), not M1 training data.",158        "",159        "| lang | reference rows | from subsets |",160        "|---|---:|---|",161    ]162    by_lang: dict[str, list[SubsetStats]] = {}163    for s in ref:164        by_lang.setdefault(s.lang, []).append(s)165    for lang in sorted(by_lang):166        subsets = by_lang[lang]167        total = sum(s.categories.get("reference", 0) for s in subsets)168        names = ", ".join(f"`{s.dataset_id.split('/')[-1]}`" for s in subsets)169        lines.append(f"| {lang} | {total} | {names} |")170 171    primary = [s for s in ref if s.lang == primary_lang]172    if primary:173        pooled: Counter = Counter()174        for s in primary:175            pooled.update(s.levels)176        lines += [177            "",178            f"Pooled level distribution for `{primary_lang}` reference subsets "179            "(class balance check):",180            "",181            "| " + " | ".join(CANONICAL_LEVELS) + " | odd |",182            "|" + "---:|" * (len(CANONICAL_LEVELS) + 1),183            f"| {_level_row(pooled)} |",184        ]185 186    if failed:187        lines += ["", "## Failed subsets", ""]188        lines += [f"- `{s.dataset_id}` — {s.error}" for s in failed]189 190    lines += [191        "",192        "---",193        "Notes: word counts use whitespace tokenisation (approximate for ar/hi). "194        "Licenses are aggregated from the per-row `license` field; decisions and "195        "exclusions are recorded in `docs/adr/0003-datasets-and-licensing.md`.",196        "",197    ]198    return "\n".join(lines)199 200 201def main() -> None:202    parser = argparse.ArgumentParser(description=__doc__)203    parser.add_argument("--langs", nargs="+", default=["en"], help="ISO codes, or 'all'")204    parser.add_argument("--datasets", nargs="+", default=None, help="explicit ids (skip discovery)")205    parser.add_argument("--report", type=Path, default=DEFAULT_REPORT)206    parser.add_argument(207        "--save-raw", action="store_true", help="save subsets to data/raw/*.parquet"208    )209    args = parser.parse_args()210 211    if args.datasets:212        pairs = [(d, lang_of(d) or "?") for d in args.datasets]213    else:214        pairs = discover_datasets(set(args.langs))215    if not pairs:216        sys.exit("No matching datasets found.")217 218    print(f"Analysing {len(pairs)} subsets: {[d for d, _ in pairs]}\n")219    all_stats = []220    for dataset_id, lang in pairs:221        print(f"-> {dataset_id} ...", flush=True)222        all_stats.append(analyze(dataset_id, lang, save_raw=args.save_raw))223 224    report = render_report(all_stats, primary_lang=args.langs[0], command=" ".join(sys.argv))225    args.report.parent.mkdir(parents=True, exist_ok=True)226    args.report.write_text(report, encoding="utf-8")227    print(f"\n{report}")228    print(f"Report written to {args.report}")229 230 231if __name__ == "__main__":232    main()233