Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465#!/usr/bin/env python3"""Build a subsampled test set for cheap in-training monitoring evals.
The full Sudoku-Extreme test set is 422,786 puzzles = 551 batches of 768, and eacheval batch runs the full ACT loop (16 steps x 21 network applications), so a fulleval costs ~3h on halfmind. A random subsample of ~12k puzzles gives +/-0.3%standard error at p~0.87 — ample for tracking progress — at ~5 min per eval.
The final reported number still comes from a full-test-set eval viascripts/eval_checkpoint.py.
Usage: python scripts/make_test_subset.py --src data/sudoku-extreme-1k-aug-1000 \ --dst data/sudoku-extreme-testsub-12k --n 12288 --seed 0"""
import argparseimport jsonimport os
import numpy as np
def main() -> None: ap = argparse.ArgumentParser() ap.add_argument("--src", required=True, help="dataset dir containing test/") ap.add_argument("--dst", required=True) ap.add_argument("--n", type=int, default=12288, help="puzzles to keep (multiple of batch size is tidiest)") ap.add_argument("--seed", type=int, default=0) args = ap.parse_args()
src_test = os.path.join(args.src, "test") dst_test = os.path.join(args.dst, "test") os.makedirs(dst_test, exist_ok=True)
inputs = np.load(os.path.join(src_test, "all__inputs.npy"), mmap_mode="r") labels = np.load(os.path.join(src_test, "all__labels.npy"), mmap_mode="r") total = inputs.shape[0] n = min(args.n, total)
rng = np.random.default_rng(args.seed) idx = np.sort(rng.choice(total, size=n, replace=False))
np.save(os.path.join(dst_test, "all__inputs.npy"), np.asarray(inputs[idx])) np.save(os.path.join(dst_test, "all__labels.npy"), np.asarray(labels[idx])) # one example per puzzle, one puzzle per group (matches how the builder emits Sudoku) np.save(os.path.join(dst_test, "all__puzzle_identifiers.npy"), np.zeros(n, dtype=np.int32)) np.save(os.path.join(dst_test, "all__puzzle_indices.npy"), np.arange(n + 1, dtype=np.int32)) np.save(os.path.join(dst_test, "all__group_indices.npy"), np.arange(n + 1, dtype=np.int32))
with open(os.path.join(src_test, "dataset.json")) as f: meta = json.load(f) meta.update(total_groups=n, total_puzzles=n, mean_puzzle_examples=1) with open(os.path.join(dst_test, "dataset.json"), "w") as f: json.dump(meta, f)
with open(os.path.join(args.dst, "subset_provenance.json"), "w") as f: json.dump({"src": args.src, "n": n, "of_total": int(total), "seed": args.seed}, f, indent=2)
print(f"wrote {n} of {total} test puzzles to {dst_test} (seed {args.seed})")
if __name__ == "__main__": main()