"""Build a training file from public MCQ datasets, excluding any dev/test overlap. python prep_data.py download --dataset CohereLabs/Global-MMLU --configs si --out raw/gmmlu python prep_data.py split --data dev.jsonl --prefix dev --eval-n 400 python prep_data.py split --data 'raw/gmmlu/*.jsonl' --prefix gm --eval-n 400 python prep_data.py build --train gm_train.jsonl dev_train.jsonl --exclude dev_eval.jsonl gm_eval.jsonl test.jsonl --out train.jsonl naist-nlp/SinhalaMMLU must not be used for training: the shared-task test set is built from it. Only counts are printed, never question text. """ import argparse import glob import json import re from collections import Counter from pathlib import Path from smmlu import get_gold, get_level, get_options, get_question, load_records, loose_hash MIN_STEM_CHARS = 30 # shorter stems ("Which is correct?") are too generic to match on alone def download(args): from datasets import get_dataset_config_names, load_dataset out = Path(args.out) out.mkdir(parents=True, exist_ok=True) configs = args.configs or get_dataset_config_names(args.dataset) for cfg in configs: ds = load_dataset(args.dataset, cfg) for split, d in ds.items(): path = out / f"{re.sub(r'[^A-Za-z0-9_.-]', '_', cfg)}__{split}.jsonl" with open(path, "w", encoding="utf-8", newline="\n") as f: for row in d: f.write(json.dumps(row, ensure_ascii=False, default=str) + "\n") print(f"{cfg}/{split}: {len(d)} rows fields={d.column_names}") def gold_of(rec, options, zero_based): if zero_based: for k in ("answer", "label", "gold", "correct"): v = str(rec.get(k, "")).strip() if v in ("0", "1", "2", "3"): return int(v) + 1 return get_gold(rec, options) def build(args): full_keys, stem_keys = set(), set() for pattern in args.exclude: for path in glob.glob(pattern): for r in load_records(path): try: q, opts = get_question(r), get_options(r) except KeyError: continue full_keys.add(loose_hash(q + "||" + "||".join(sorted(opts)))) if len(q) >= MIN_STEM_CHARS: stem_keys.add(loose_hash(q)) print(f"exclusion keys: {len(full_keys)} question+options, {len(stem_keys)} long stems") seen, rows, stats, per_source = set(), [], Counter(), Counter() for pattern in args.train: for path in sorted(glob.glob(pattern)): for r in load_records(path): try: q, opts = get_question(r), get_options(r) except KeyError: stats["unparseable"] += 1 continue if len(opts) != 4 or any(not o.strip() for o in opts): stats["not 4 non-empty options"] += 1 continue gold = gold_of(r, opts, args.zero_based) if gold is None: stats["no gold"] += 1 continue full = loose_hash(q + "||" + "||".join(sorted(opts))) if full in full_keys or (len(q) >= MIN_STEM_CHARS and loose_hash(q) in stem_keys): stats["removed: overlaps dev/test"] += 1 continue if full in seen: stats["removed: duplicate"] += 1 continue seen.add(full) rows.append({"question": q, "options": opts, "answer": str(gold), "subject": r.get("subject"), "level": get_level(r), "source": Path(path).parent.name}) per_source[Path(path).parent.name] += 1 with open(args.out, "w", encoding="utf-8", newline="\n") as f: for row in rows: f.write(json.dumps(row, ensure_ascii=False) + "\n") print(f"kept {len(rows)} -> {args.out}") print("per source:", dict(per_source)) print("dropped:", dict(stats)) print("gold distribution:", dict(sorted(Counter(r["answer"] for r in rows).items()))) print("levels:", dict(Counter(r["level"] for r in rows))) def merge(args): """Merge per-subject JSON files (/__.json) into one JSONL. Fields already in a record win; level, subject and category are filled from the path only when a record lacks them. """ rows, filled = [], Counter() files = sorted(f for pattern in args.files for f in glob.glob(pattern, recursive=True)) for path in files: p = Path(path) level = p.parent.name.lower() stem_parts = p.stem.rsplit("_", 2) for i, r in enumerate(load_records(path)): r = dict(r) if not (r.get("level") or r.get("difficulty")): r["level"] = level filled["level"] += 1 if not r.get("subject") and len(stem_parts) == 3: r["subject"] = stem_parts[0] filled["subject"] += 1 if not r.get("category") and len(stem_parts) == 3: r["category"] = stem_parts[1] filled["category"] += 1 r["_source_file"] = p.name rows.append(r) keys = Counter((get_level(r), str(r.get("subject")), str(r.get("q_no", r.get("id")))) for r in rows) with open(args.out, "w", encoding="utf-8", newline="\n") as f: for r in rows: f.write(json.dumps(r, ensure_ascii=False) + "\n") print(f"{len(files)} files, {len(rows)} records -> {args.out}") print("filled from path (field missing in record):", dict(filled) or "none") print("per level:", dict(Counter(get_level(r) for r in rows))) print("(level, subject, q_no) keys unique:", all(c == 1 for c in keys.values()), "| duplicate keys:", sum(c - 1 for c in keys.values() if c > 1)) def split(args): """Split a labelled file (one file or a glob) into train/eval parts, stratified by level. --eval-n fixes the eval size (spread across levels in proportion); otherwise --train-frac is used. """ import random recs_all = [r for f in sorted(glob.glob(args.data)) for r in load_records(f)] by_level = {} for r in recs_all: by_level.setdefault(get_level(r), []).append(r) rng = random.Random(args.seed) parts = {"train": [], "eval": []} for lvl, recs in sorted(by_level.items()): rng.shuffle(recs) if args.eval_n: cut = len(recs) - round(args.eval_n * len(recs) / len(recs_all)) else: cut = int(len(recs) * args.train_frac) parts["train"] += recs[:cut] parts["eval"] += recs[cut:] print(f"{lvl or '(none)'}: {cut} train / {len(recs) - cut} eval") for name, recs in parts.items(): path = f"{args.prefix}_{name}.jsonl" with open(path, "w", encoding="utf-8", newline="\n") as f: for r in recs: f.write(json.dumps(r, ensure_ascii=False) + "\n") print(f"wrote {len(recs)} -> {path}") def main(): ap = argparse.ArgumentParser() sub = ap.add_subparsers(dest="cmd", required=True) p = sub.add_parser("merge") p.add_argument("--files", nargs="+", required=True, help="glob patterns, e.g. 'raw/test/**/*.json'") p.add_argument("--out", required=True) p = sub.add_parser("split") p.add_argument("--data", required=True) p.add_argument("--prefix", required=True) p.add_argument("--train-frac", type=float, default=0.5) p.add_argument("--eval-n", type=int, default=0, help="fixed eval size instead of --train-frac") p.add_argument("--seed", type=int, default=42) p = sub.add_parser("download") p.add_argument("--dataset", required=True) p.add_argument("--configs", nargs="*") p.add_argument("--out", required=True) p = sub.add_parser("build") p.add_argument("--train", nargs="+", required=True, help="files or glob patterns") p.add_argument("--exclude", nargs="+", default=[], help="dev/test files whose questions must not be trained on") p.add_argument("--zero-based", action="store_true", help="gold answers are 0-3 instead of 1-4/A-D") p.add_argument("--out", required=True) args = ap.parse_args() {"download": download, "build": build, "split": split, "merge": merge}[args.cmd](args) if __name__ == "__main__": main()