"""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: _) 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({"": 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()