training/prep_data.py

"""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 (<level>/<Subject>_<category>_<level>.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()