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
|
#!/usr/bin/env python3
"""Stratified game-level train/holdout split over a games.jsonl
(build_games.py output).
Splits by GAME, not by individual position -- every position from one
game lands on the same side of the split, avoiding the leakage a flat
position-level shuffle would have (adjacent plies / recurring
middlegame structures from the same game are highly correlated, not
independent samples; see eval_tune/tune_eval_dna.py's own holdout
split, which doesn't do this).
Stratified on two axes so the holdout's composition mirrors the full
pool's rather than a flat 10% sample risking (by chance) too many
decisive games, too few draws, or skewing toward short/long games:
- result: White win / Black win / Draw
- length: quartile bin of the game's annotated ply count (quartiles
of the actual corpus, not guessed fixed cutoffs)
10% (configurable) is sampled independently within each of the
resulting (result, length-quartile) strata, not 10% of the pool
overall.
Usage:
python3 split_holdout.py --in runs/sd10_twic_2M/games.jsonl \
--train-out runs/sd10_twic_2M/games.train.jsonl \
--holdout-out runs/sd10_twic_2M/games.holdout.jsonl \
--holdout-frac 0.1 --seed 0
"""
import argparse
import json
import math
import random
from collections import defaultdict
def result_bucket(result):
if result == "1-0":
return "white_win"
if result == "0-1":
return "black_win"
if result == "1/2-1/2":
return "draw"
return "unknown"
def quartile_bin(length, cuts):
for i, c in enumerate(cuts):
if length <= c:
return i
return len(cuts)
def main():
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("--in", dest="infile", required=True)
ap.add_argument("--train-out", required=True)
ap.add_argument("--holdout-out", required=True)
ap.add_argument("--holdout-frac", type=float, default=0.1)
ap.add_argument("--seed", type=int, default=0)
args = ap.parse_args()
games = []
with open(args.infile) as f:
for line in f:
games.append(json.loads(line))
print(f"Loaded {len(games)} games from {args.infile}")
lengths = sorted(len(g["positions"]) for g in games)
def pct(p):
idx = min(int(p * len(lengths)), len(lengths) - 1)
return lengths[idx]
cuts = [pct(0.25), pct(0.50), pct(0.75)]
print(f"Length quartile cutpoints (annotated plies): {cuts}")
strata = defaultdict(list)
for g in games:
key = (result_bucket(g["result"]),
quartile_bin(len(g["positions"]), cuts))
strata[key].append(g)
rng = random.Random(args.seed)
train_games = []
holdout_games = []
print(f"\n{'stratum':<24} {'total':>7} {'holdout':>8} {'train':>8}")
for key in sorted(strata.keys()):
bucket = strata[key]
rng.shuffle(bucket)
n_holdout = math.ceil(len(bucket) * args.holdout_frac)
holdout_games.extend(bucket[:n_holdout])
train_games.extend(bucket[n_holdout:])
label = f"{key[0]}/q{key[1]}"
print(f"{label:<24} {len(bucket):>7} {n_holdout:>8} "
f"{len(bucket) - n_holdout:>8}")
print(f"\n{'TOTAL':<24} {len(games):>7} {len(holdout_games):>8} "
f"{len(train_games):>8}")
with open(args.train_out, "w") as f:
for g in train_games:
f.write(json.dumps(g) + "\n")
with open(args.holdout_out, "w") as f:
for g in holdout_games:
f.write(json.dumps(g) + "\n")
n_train_pos = sum(len(g["positions"]) for g in train_games)
n_holdout_pos = sum(len(g["positions"]) for g in holdout_games)
print(f"\nWrote {args.train_out}: {len(train_games)} games, "
f"{n_train_pos} positions")
print(f"Wrote {args.holdout_out}: {len(holdout_games)} games, "
f"{n_holdout_pos} positions")
if __name__ == "__main__":
main()
|
