File size: 5,413 Bytes
e5235fb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | #!/usr/bin/env python3
"""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"] # T0.2 (D-12)
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()
|