| |
| """T0.3 · Fixed item-id files for RS-fast (spec §7), so every run scores the same items. |
| |
| python -I scripts/prep_rs_subsets.py # downloads ~10 MB of parquet into $MT_DATA/evalsets |
| |
| - MMLU-Redux 2.0: 1,000 of the items its annotators marked `ok` (label verified), stratified over the |
| 57 subjects (largest-remainder allocation), seed 0. |
| - Global-MMLU-Lite: all 400 test items for each RS-fast language (the set is already fixed and |
| human-verified; the file pins the ids and the parquet hash). |
| - BELEBELE: 200 of the 900 questions, seed 0, the same questions in every RS-fast language (BELEBELE |
| is fully parallel), so cross-language differences are not item differences. |
| - GSM8K-250 already exists (configs/eval/subsets/gsm8k_250.json, made for the T1.4 FP8 check). |
| |
| Parquet comes from the Hub's parquet conversion of a pinned dataset revision; each file's sha256 is |
| recorded. Downloaded files are data: read with pyarrow only. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import os |
| import random |
| import urllib.request |
| from pathlib import Path |
|
|
| REPO = Path(__file__).resolve().parent.parent |
| DATA = Path(os.environ.get("MT_DATA", "/path/to/data")) |
| OUT = REPO / "configs" / "eval" / "subsets" |
| RS_FAST = ["zh", "ar", "ko", "hu", "bn"] |
| BELEBELE_CFG = {"zh": "zho_Hans", "ar": "arb_Arab", "ko": "kor_Hang", "hu": "hun_Latn", "bn": "ben_Beng", |
| "en": "eng_Latn"} |
| REV = {"mmlu_redux": ("edinburgh-dawg/mmlu-redux-2.0", "372ea425445d51e1ba1188c56e5e893f8138621f"), |
| "global_mmlu_lite": ("CohereLabs/Global-MMLU-Lite", "36c2fd756f19ccf13a9a96c8e53ccecc02192b8b"), |
| "belebele": ("facebook/belebele", "7899cdfa4e1e0d733fd77c848e2c273cb1d32be2")} |
| PARQUET = "https://huggingface.co/api/datasets/{ds}/parquet/{cfg}/{split}/0.parquet" |
|
|
|
|
| def fetch(url: str, dest: Path) -> str: |
| dest.parent.mkdir(parents=True, exist_ok=True) |
| if not dest.exists(): |
| tmp = dest.with_suffix(".part") |
| urllib.request.urlretrieve(url, tmp) |
| tmp.rename(dest) |
| return hashlib.sha256(dest.read_bytes()).hexdigest() |
|
|
|
|
| def rows(ds: str, cfg: str, split: str, dest: Path) -> tuple[list[dict], str]: |
| import pyarrow.parquet as pq |
| sha = fetch(PARQUET.format(ds=ds, cfg=cfg, split=split), dest) |
| return pq.read_table(dest).to_pylist(), sha |
|
|
|
|
| def write(name: str, obj: dict) -> None: |
| OUT.mkdir(parents=True, exist_ok=True) |
| (OUT / name).write_text(json.dumps(obj, indent=1, ensure_ascii=False) + "\n") |
| print(f"{name}: {obj.get('n')} items") |
|
|
|
|
| def allocate(sizes: dict[str, int], total: int) -> dict[str, int]: |
| """Largest-remainder allocation of `total` proportional to `sizes`, capped at each size.""" |
| n = sum(sizes.values()) |
| quota = {k: total * v / n for k, v in sizes.items()} |
| alloc = {k: min(sizes[k], int(q)) for k, q in quota.items()} |
| for k in sorted(sizes, key=lambda k: (-(quota[k] - int(quota[k])), k)): |
| if sum(alloc.values()) >= total: |
| break |
| if alloc[k] < sizes[k]: |
| alloc[k] += 1 |
| return alloc |
|
|
|
|
| def mmlu_redux(seed: int, n: int) -> None: |
| import urllib.request as ur |
| ds, rev = REV["mmlu_redux"] |
| subjects = sorted(json.loads(ur.urlopen(f"https://huggingface.co/api/datasets/{ds}/parquet").read())) |
| ok, shas = {}, {} |
| for s in subjects: |
| rs, shas[s] = rows(ds, s, "test", DATA / f"evalsets/mmlu_redux/{s}_test.parquet") |
| ok[s] = [i for i, r in enumerate(rs) if r.get("error_type") == "ok"] |
| alloc = allocate({s: len(v) for s, v in ok.items()}, n) |
| rng = random.Random(seed) |
| items = {s: sorted(rng.sample(ok[s], alloc[s])) for s in subjects} |
| write("mmlu_redux_1000.json", {"dataset": ds, "revision": rev, "split": "test", "filter": "error_type == ok", |
| "seed": seed, "n": sum(alloc.values()), "ok_items": sum(map(len, ok.values())), |
| "parquet_sha256": shas, "ids": items}) |
|
|
|
|
| def global_mmlu_lite() -> None: |
| ds, rev = REV["global_mmlu_lite"] |
| out = {"dataset": ds, "revision": rev, "split": "test", "languages": RS_FAST, "parquet_sha256": {}, "ids": {}} |
| for lang in RS_FAST: |
| rs, out["parquet_sha256"][lang] = rows(ds, lang, "test", DATA / f"evalsets/global_mmlu_lite/{lang}_test.parquet") |
| out["ids"][lang] = [r["sample_id"] for r in rs] |
| out["n"] = sum(len(v) for v in out["ids"].values()) |
| write("global_mmlu_lite_rs5.json", out) |
|
|
|
|
| def belebele(seed: int, n: int) -> None: |
| ds, rev = REV["belebele"] |
| out = {"dataset": ds, "revision": rev, "split": "test", "languages": RS_FAST, "seed": seed, "parquet_sha256": {}} |
| keys = None |
| for lang in RS_FAST: |
| cfg = BELEBELE_CFG[lang] |
| rs, out["parquet_sha256"][lang] = rows(ds, cfg, "test", DATA / f"evalsets/belebele/{cfg}_test.parquet") |
| k = {(r["link"], r["question_number"]) for r in rs} |
| keys = k if keys is None else keys & k |
| pick = sorted(random.Random(seed).sample(sorted(keys), n)) |
| out.update(n=n, parallel_questions=len(keys), ids=[{"link": l, "question_number": q} for l, q in pick]) |
| write("belebele_200.json", out) |
|
|
|
|
| def main() -> None: |
| p = argparse.ArgumentParser() |
| p.add_argument("--seed", type=int, default=0) |
| a = p.parse_args() |
| mmlu_redux(a.seed, 1000) |
| global_mmlu_lite() |
| belebele(a.seed, 200) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
