scripts/download_prm800k.py

"""Download PRM800K step-label files (OpenAI, MIT) into data/raw/prm800k/.

    python scripts/download_prm800k.py                 # phase2_train (what we use)
    python scripts/download_prm800k.py --all           # all four phase files

The GitHub raw URLs serve the real JSONL (schema: question/label/steps/...),
which is exactly what mathcompose.datagen.prm800k_loader expects.
"""
import argparse
import sys
import urllib.request
from pathlib import Path

BASE = "https://github.com/openai/prm800k/raw/main/prm800k/data"
FILES = ["phase2_train.jsonl", "phase1_train.jsonl", "phase2_test.jsonl", "phase1_test.jsonl"]


def download(name: str, out_dir: Path) -> None:
    url = f"{BASE}/{name}"
    dest = out_dir / name
    dest.parent.mkdir(parents=True, exist_ok=True)
    print(f"downloading {name} -> {dest}")

    def _hook(block, bsize, total):
        if total > 0:
            pct = min(100, block * bsize * 100 // total)
            sys.stdout.write(f"\r  {pct:3d}%  ({block * bsize / 1e6:.0f} MB)")
            sys.stdout.flush()

    urllib.request.urlretrieve(url, dest, reporthook=_hook)
    print(f"\r  done: {dest.stat().st_size / 1e6:.1f} MB{' ' * 10}")


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--out-dir", default="data/raw/prm800k")
    ap.add_argument("--all", action="store_true", help="download all four phase files")
    args = ap.parse_args()
    names = FILES if args.all else ["phase2_train.jsonl"]
    for n in names:
        download(n, Path(args.out_dir))
    print("\nNext: python -m mathcompose.data.build_verifier_dataset "
          f"--prm800k {args.out_dir}/phase2_train.jsonl --teacher promptlens --limit 15000")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())