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
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
|
#!/usr/bin/env python3
"""Feature extraction driver: takes the mined FENs in mate_training_data.csv
and runs each through the engine's `evalkingdump` command (KING_DEFECT_DUMP
build only) to pull _EvalKing's raw per-defect breakdown, joining it back
onto each row's label/metadata to produce the final regression-ready CSV.
Sharded across N engine subprocesses (one `script` file each, run in
parallel) since 86k FENs through a single serial engine invocation would
be needlessly slow on this 24-core box -- same reasoning as
shard_mine.py's PGN sharding.
Each engine invocation processes its shard via `script <file>` containing
alternating `setboard <fen>` / `evalkingdump` lines; evalkingdump always
emits exactly two `KINGDEFECTDUMP color=...` lines (black then white, in
that fixed loop order) per position, unconditionally -- so the Nth pair
of KINGDEFECTDUMP lines in the shard's output corresponds to the Nth row
of that shard's input, and we pick out whichever color each row's
"mated_color_or_sampled_color" column names.
"""
import argparse
import csv
import os
import re
import subprocess
import sys
import time
KV_RE = re.compile(r"(\S+)=(\S+)")
# Fixed schema, since an off-board ring square (king near an edge/corner)
# emits only "sqN_offboard=1" while an on-board square emits 9 separate
# keys -- different king positions otherwise produce entirely different
# column sets, which would silently misalign rows if we inferred columns
# from whichever row happened to be parsed first (caught by a per-shard
# "feature key mismatch" warning during testing).
PIECE_NAMES = ("pawn", "minor", "rook", "queen")
PER_SQUARE_ON_BOARD_KEYS = [f"atk_{p}" for p in PIECE_NAMES] + \
["undef", "def", "xray", "flight"]
GLOBAL_KEYS = ["filestorm", "filestormopenfile", "filestormpawn",
"queenproximity", "initloc", "countertotal",
"safetyvalue", "attackerarmyscaler", "castled",
"minorsathome", "misscastleshort", "misscastlelong"]
FEATURE_KEYS = []
for _sq in range(8):
FEATURE_KEYS.append(f"sq{_sq}_offboard")
for _k in PER_SQUARE_ON_BOARD_KEYS:
FEATURE_KEYS.append(f"sq{_sq}_{_k}")
FEATURE_KEYS += GLOBAL_KEYS
def parse_kingdefectdump_line(line):
"""Returns (color, {key: value_str}) from a 'KINGDEFECTDUMP color=... ...'
line, normalized to FEATURE_KEYS' fixed schema -- off-board squares get
offboard=1 and all per-square on-board keys defaulted to 0 rather than
left absent."""
parts = line.strip().split(None, 1)
assert parts[0] == "KINGDEFECTDUMP"
raw_kv = dict(KV_RE.findall(parts[1]))
color = raw_kv.pop("color")
kv = {}
for sq in range(8):
offboard_key = f"sq{sq}_offboard"
if offboard_key in raw_kv:
kv[offboard_key] = raw_kv[offboard_key]
for k in PER_SQUARE_ON_BOARD_KEYS:
kv[f"sq{sq}_{k}"] = "0"
else:
kv[offboard_key] = "0"
for k in PER_SQUARE_ON_BOARD_KEYS:
full_key = f"sq{sq}_{k}"
kv[full_key] = raw_kv[full_key]
for k in GLOBAL_KEYS:
kv[k] = raw_kv[k]
assert set(kv.keys()) == set(FEATURE_KEYS), \
f"parsed keys don't match fixed schema: {set(kv) ^ set(FEATURE_KEYS)}"
return color, kv
def run_shard(engine_path, rows, shard_dir, shard_idx):
script_path = os.path.join(shard_dir, f"extract_{shard_idx}.script")
out_log = os.path.join(shard_dir, f"extract_{shard_idx}.out")
with open(script_path, "w") as f:
f.write("force\n")
f.write("nopost\n")
for row in rows:
f.write(f"setboard {row['fen']}\n")
f.write("evalkingdump\n")
cmd = [engine_path, "--batch", "--command",
f"script {script_path}"]
with open(out_log, "w") as logf:
rc = subprocess.run(cmd, stdout=logf, stderr=subprocess.STDOUT).returncode
if rc != 0:
raise RuntimeError(f"shard {shard_idx} engine exited {rc}, see {out_log}")
dumps = []
with open(out_log) as f:
for line in f:
if line.startswith("KINGDEFECTDUMP"):
dumps.append(parse_kingdefectdump_line(line))
if len(dumps) != 2 * len(rows):
raise RuntimeError(
f"shard {shard_idx}: expected {2 * len(rows)} KINGDEFECTDUMP "
f"lines (2 per row), got {len(dumps)} -- a setboard likely "
f"failed silently and evalkingdump re-used a stale position. "
f"See {out_log}.")
out_rows = []
for i, row in enumerate(rows):
black_color, black_kv = dumps[2 * i]
white_color, white_kv = dumps[2 * i + 1]
assert black_color == "black" and white_color == "white", \
f"row {i}: unexpected color order {black_color}/{white_color}"
target = row["mated_color_or_sampled_color"]
kv = black_kv if target == "black" else white_kv
out_row = dict(row)
out_row.update(kv)
out_rows.append(out_row)
os.remove(script_path)
os.remove(out_log)
return out_rows
def main():
ap = argparse.ArgumentParser()
ap.add_argument("training_csv")
ap.add_argument("out_csv")
ap.add_argument("--engine", default="../typhoon_kingdump")
ap.add_argument("--num-shards", type=int, default=8)
ap.add_argument("--work-dir", default=None)
args = ap.parse_args()
engine_path = os.path.abspath(args.engine)
if not os.path.isfile(engine_path):
print(f"Engine not found: {engine_path}", file=sys.stderr)
sys.exit(1)
with open(args.training_csv, newline="") as f:
rows = list(csv.DictReader(f))
print(f"Loaded {len(rows)} rows from {args.training_csv}", file=sys.stderr)
work_dir = args.work_dir or os.path.dirname(os.path.abspath(args.out_csv))
os.makedirs(work_dir, exist_ok=True)
n = args.num_shards
shard_size = (len(rows) + n - 1) // n
shards = [rows[i:i + shard_size] for i in range(0, len(rows), shard_size)]
t0 = time.time()
all_out_rows = [[] for _ in shards]
import concurrent.futures
with concurrent.futures.ThreadPoolExecutor(max_workers=n) as pool:
futures = {}
for idx, shard_rows in enumerate(shards):
fut = pool.submit(run_shard, engine_path, shard_rows, work_dir, idx)
futures[fut] = idx
for fut in concurrent.futures.as_completed(futures):
idx = futures[fut]
out_rows = fut.result()
all_out_rows[idx] = out_rows
print(f"Shard {idx} done ({len(out_rows)} rows)", file=sys.stderr)
print(f"All shards done in {time.time() - t0:.1f}s", file=sys.stderr)
base_cols = ["fen", "label", "mated_color_or_sampled_color", "n_back",
"game_index", "control_source"]
fieldnames = base_cols + FEATURE_KEYS
with open(args.out_csv, "w", newline="") as f:
w = csv.DictWriter(f, fieldnames=fieldnames)
w.writeheader()
for out_rows in all_out_rows:
for row in out_rows:
w.writerow({k: row.get(k) for k in fieldnames})
total = sum(len(r) for r in all_out_rows)
print(f"Wrote {total} feature rows to {args.out_csv}", file=sys.stderr)
if __name__ == "__main__":
main()
|
