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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
|
#!/usr/bin/env python3
"""Fit the WDL-sigmoid scale K on the TRAIN split only, then label
every non-exact position in both train and holdout with:
- p_search: sigmoid(score / K) -- the search-score-derived win
probability estimate.
- r_outcome: that position's own game's actual final result,
converted to side-to-move-relative (matching the score's own
convention: 1.0 = the side to move at this FEN eventually won
their game, 0.0 = eventually lost, 0.5 = drawn).
Does NOT blend p_search/r_outcome into a single training target --
that mixing weight (lambda) is left as a free parameter to sweep
later; storing both components separately means that sweep never
needs to re-derive anything from here.
K is fit by ternary search minimizing MSE(sigmoid(score/K), r_outcome)
over every non-exact TRAIN position (not holdout -- fitting K on
holdout data would leak information into what's supposed to be an
untouched generalization check). The same fitted K is then applied
(never refit) to label holdout positions too.
Usage:
python3 fit_k_and_label.py --train runs/sd10_twic_2M/games.train.jsonl \
--holdout runs/sd10_twic_2M/games.holdout.jsonl \
--train-out runs/sd10_twic_2M/games.train.labeled.jsonl \
--holdout-out runs/sd10_twic_2M/games.holdout.labeled.jsonl
"""
import argparse
import json
import math
def result_to_white_score(result):
"""1.0/0.5/0.0 from White's perspective, or None if unknown ('*')."""
return {"1-0": 1.0, "0-1": 0.0, "1/2-1/2": 0.5}.get(result)
def stm_relative_outcome(white_score, fen):
"""Flip a White-perspective game result to the side-to-move's own
perspective at this FEN, matching the score field's own
side-to-move-relative convention."""
stm_is_white = fen.split()[1] == "w"
return white_score if stm_is_white else (1.0 - white_score)
def collect_calibration_pairs(games_path):
"""(score, r_outcome) for every non-exact position across all
games in this file."""
pairs = []
n_unknown_result = 0
with open(games_path) as f:
for line in f:
g = json.loads(line)
white_score = result_to_white_score(g["result"])
if white_score is None:
n_unknown_result += 1
continue
for pos in g["positions"]:
if pos.get("exact"):
continue
r = stm_relative_outcome(white_score, pos["fen"])
pairs.append((pos["score"], r))
return pairs, n_unknown_result
def mse(pairs, k):
err = 0.0
for score, r in pairs:
p = 1.0 / (1.0 + math.exp(-score / k))
err += (r - p) ** 2
return err / max(len(pairs), 1)
def fit_k(pairs, lo=0.05, hi=20.0, iterations=40):
for _ in range(iterations):
m1 = lo + (hi - lo) / 3
m2 = hi - (hi - lo) / 3
if mse(pairs, m1) < mse(pairs, m2):
hi = m2
else:
lo = m1
return (lo + hi) / 2
def label_file(in_path, out_path, k):
n_games = 0
n_positions = 0
n_exact_skipped = 0
n_unknown_result = 0
with open(in_path) as fin, open(out_path, "w") as fout:
for line in fin:
g = json.loads(line)
white_score = result_to_white_score(g["result"])
if white_score is None:
n_unknown_result += 1
fout.write(json.dumps(g) + "\n")
continue
for pos in g["positions"]:
n_positions += 1
if pos.get("exact"):
n_exact_skipped += 1
continue
r = stm_relative_outcome(white_score, pos["fen"])
p_search = 1.0 / (1.0 + math.exp(-pos["score"] / k))
pos["p_search"] = p_search
pos["r_outcome"] = r
fout.write(json.dumps(g) + "\n")
n_games += 1
return n_games, n_positions, n_exact_skipped, n_unknown_result
def main():
ap = argparse.ArgumentParser(description=__doc__)
ap.add_argument("--train", required=True)
ap.add_argument("--holdout", required=True)
ap.add_argument("--train-out", required=True)
ap.add_argument("--holdout-out", required=True)
args = ap.parse_args()
print(f"Collecting calibration pairs from {args.train} (TRAIN only) ...")
pairs, n_unknown = collect_calibration_pairs(args.train)
print(f" {len(pairs)} (score, outcome) pairs "
f"({n_unknown} games skipped: unknown result)")
print("Fitting K (ternary search, minimizing sigmoid-MSE against "
"actual game outcomes) ...")
k = fit_k(pairs)
final_mse = mse(pairs, k)
print(f" K = {k:.2f} (MSE at fit = {final_mse:.5f})")
print(f"\nLabeling {args.train} -> {args.train_out} ...")
n_g, n_p, n_exact, n_unk = label_file(args.train, args.train_out, k)
print(f" {n_g} games, {n_p} positions, {n_exact} exact (unlabeled), "
f"{n_unk} unknown-result games (unlabeled)")
print(f"Labeling {args.holdout} -> {args.holdout_out} (same K, not refit) ...")
n_g, n_p, n_exact, n_unk = label_file(args.holdout, args.holdout_out, k)
print(f" {n_g} games, {n_p} positions, {n_exact} exact (unlabeled), "
f"{n_unk} unknown-result games (unlabeled)")
print(f"\nFitted K = {k:.4f} -- record this; applying a different K "
f"later (e.g. after labeling more data) is not the same model.")
if __name__ == "__main__":
main()
|
