| |
| """Freeze evaluation splits so every run is scored on identical states.""" |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import sys |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(ROOT)) |
|
|
| from jevon.data import SamplerConfig, freeze_split |
|
|
|
|
| def main() -> None: |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--out-dir", default="data/frozen") |
| p.add_argument("--val-states", type=int, default=180) |
| p.add_argument("--test-states", type=int, default=240) |
| p.add_argument("--ood-states", type=int, default=160) |
| p.add_argument("--seed", type=int, default=20260919) |
| args = p.parse_args() |
|
|
| out = Path(args.out_dir) |
| out.mkdir(parents=True, exist_ok=True) |
| base = SamplerConfig(augment=False) |
| |
| ood = SamplerConfig(maze_sizes=(41, 51), snake_sizes=(20, 24), augment=False) |
|
|
| manifest = {} |
| manifest["val"] = freeze_split(out / "val.jsonl", base, "val", args.val_states, |
| args.seed) |
| manifest["test"] = freeze_split(out / "test.jsonl", base, "test", |
| args.test_states, args.seed + 1) |
| manifest["ood"] = freeze_split(out / "ood.jsonl", ood, "test", args.ood_states, |
| args.seed + 2) |
| manifest["config"] = {"train_maze_sizes": list(base.maze_sizes), |
| "train_snake_sizes": list(base.snake_sizes), |
| "ood_maze_sizes": list(ood.maze_sizes), |
| "ood_snake_sizes": list(ood.snake_sizes)} |
| (out / "manifest.json").write_text(json.dumps(manifest, indent=2)) |
| print(json.dumps(manifest, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
