1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
|
#!/usr/bin/env python3
"""Reassemble an annotate.py run's flat per-position output back into
game-shaped containers: one record per source game, holding its
ordered list of annotated positions.
Why this exists: extract_positions.py's positions.jsonl (fen, score,
depth, pv, exact per row) is a flat list -- the row order follows each
shard's script order, which happens to group consecutive rows by game
(write_shard_script writes one game's positions fully before moving to
the next), but there's no game_id recorded, so that grouping is
unrecoverable from positions.jsonl alone. That matters for anything
that needs to avoid splitting one game's (highly correlated,
near-duplicate) positions across a train/holdout boundary.
This script recovers it deterministically, without re-running any
search: re-load the same PGN file with the same max_games (pure
parsing, no engine involved -- cheap), re-shard it the same way
annotate.py did (same n_shards, same deterministic round-robin), and
walk each game's known position count (positions_for_game) to slice
the already-extracted flat rows back into per-game chunks.
Output shape -- deliberately the same "a game containing an ordered
list of annotated positions" shape a self-play PGN with per-move
score/PV comments already has, so a future loader for self-play games
can produce the same games.jsonl schema and both sources mix freely at
training-set-assembly time:
{"game_id": "...", "result": "1-0", "positions": [
{"fen": ..., "depth": ..., "score": ..., "score_raw": ...,
"pv": ..., "exact": ..., "exact_reason": ...},
...
]}
Usage:
python3 build_games.py --rundir runs/sd10_twic_2M \
--pgn ../pgn/twic_filtered.pgn --max-games 30000 --shards 16 \
--out runs/sd10_twic_2M/games.jsonl
"""
import argparse
import json
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parent))
from annotate import load_games, shard_games, positions_for_game
def main():
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("--rundir", required=True,
help="directory holding positions.jsonl (extract_positions.py output)")
ap.add_argument("--positions", default=None,
help="defaults to <rundir>/positions.jsonl")
ap.add_argument("--pgn", nargs="+", required=True,
help="same PGN file(s) annotate.py was originally run against")
ap.add_argument("--max-games", type=int, required=True,
help="same --max-games (or games actually used) as the original run")
ap.add_argument("--shards", type=int, default=16,
help="same --shards as the original run")
ap.add_argument("--prefilter-quiet", action="store_true",
help="must match whether the original annotate.py run "
"used --prefilter-quiet -- changes positions_for_game's "
"per-game ply count, which this script relies on to "
"re-slice the flat positions.jsonl")
ap.add_argument("--out", required=True)
args = ap.parse_args()
rundir = Path(args.rundir)
positions_path = Path(args.positions) if args.positions else rundir / "positions.jsonl"
print(f"Loading all rows from {positions_path} ...")
flat_rows = []
with open(positions_path) as f:
for line in f:
flat_rows.append(json.loads(line))
print(f" {len(flat_rows)} rows")
print(f"Re-loading {args.pgn} (max_games={args.max_games}) to recover game boundaries ...")
games = load_games(args.pgn, args.max_games)
print(f" {len(games)} games")
shards = shard_games(games, args.shards)
# Per-shard row ranges: shard_NN's rows are a contiguous slice of
# flat_rows in the same order extract_positions.py wrote them
# (shard 00 first, then 01, ...), each shard's own length equal to
# its total position count.
shard_lengths = []
for shard in shards:
shard_lengths.append(sum(
1 for g in shard
for _ in positions_for_game(g, args.prefilter_quiet)))
if sum(shard_lengths) != len(flat_rows):
sys.exit(f"row count mismatch: recomputed shards expect "
f"{sum(shard_lengths)} total positions, "
f"positions.jsonl has {len(flat_rows)} -- wrong --pgn/"
f"--max-games/--shards for this run?")
n_games_written = 0
n_positions_written = 0
with open(args.out, "w") as out:
row_cursor = 0
for shard_idx, shard in enumerate(shards):
for game_idx, game in enumerate(shard):
n_positions = sum(1 for _ in positions_for_game(game, args.prefilter_quiet))
game_rows = flat_rows[row_cursor:row_cursor + n_positions]
row_cursor += n_positions
rec = {
"game_id": f"{shard_idx:02d}_{game_idx:05d}",
"result": game.headers.get("Result", "*"),
"positions": game_rows,
}
out.write(json.dumps(rec) + "\n")
n_games_written += 1
n_positions_written += len(game_rows)
assert row_cursor == len(flat_rows)
print(f"Wrote {args.out}: {n_games_written} games, "
f"{n_positions_written} total positions")
if __name__ == "__main__":
main()
|
