#!/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. Both the calibration-pair collection (--train) and the labeling pass (--train/--holdout) are sharded across --workers processes via multiprocessing.Pool, each worker vectorizing its own chunk of games with numpy -- this file used to run ~80 full Python-loop passes over every pair during the ternary search (2 mse() evals x 40 iterations), which was the slowest step in the post-annotation pipeline relative to how little actual compute it does. Per-game-chunk numpy vectorization plus Pool.imap (order-preserving, like build_training_tensors.py's existing pattern) turns that into array ops. 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 from multiprocessing import Pool import numpy as np 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 sigmoid(scores, k): return 1.0 / (1.0 + np.exp(-scores / k)) def mse(scores, outcomes, k): p = sigmoid(scores, k) return float(np.mean((outcomes - p) ** 2)) if len(scores) else 0.0 def fit_k(scores, outcomes, lo=0.05, hi=20.0, iterations=40): for _ in range(iterations): m1 = lo + (hi - lo) / 3 m2 = hi - (hi - lo) / 3 if mse(scores, outcomes, m1) < mse(scores, outcomes, m2): hi = m2 else: lo = m1 return (lo + hi) / 2 def iter_chunks(path, chunk_size): """Batches raw lines (one game per line) of an open jsonl file into lists of up to chunk_size lines -- the unit of work handed to one worker process at a time. Text, not parsed JSON, so the parse itself happens in the worker, same division of labor as build_training_tensors.py's iter_chunks.""" chunk = [] with open(path) as f: for line in f: chunk.append(line) if len(chunk) >= chunk_size: yield chunk chunk = [] if chunk: yield chunk def _non_exact_scores_and_outcomes(game): """(scores, outcomes) numpy arrays for one game's non-exact positions, or None if the game's result is unknown or it has no non-exact positions.""" white_score = result_to_white_score(game["result"]) if white_score is None: return None non_exact = [p for p in game["positions"] if not p.get("exact")] if not non_exact: return None scores = np.fromiter((p["score"] for p in non_exact), dtype=np.float64, count=len(non_exact)) is_white = np.fromiter( (p["fen"].split()[1] == "w" for p in non_exact), dtype=bool, count=len(non_exact)) outcomes = np.where(is_white, white_score, 1.0 - white_score) return scores, outcomes def collect_chunk(lines): """Runs in a worker process: parses a chunk of game lines and returns (scores, outcomes, n_unknown_result) for every non-exact position across those games.""" score_pieces = [] outcome_pieces = [] n_unknown_result = 0 for line in lines: g = json.loads(line) result = _non_exact_scores_and_outcomes(g) if result is None: if result_to_white_score(g["result"]) is None: n_unknown_result += 1 continue scores, outcomes = result score_pieces.append(scores) outcome_pieces.append(outcomes) if not score_pieces: return np.empty(0), np.empty(0), n_unknown_result return (np.concatenate(score_pieces), np.concatenate(outcome_pieces), n_unknown_result) def collect_calibration_pairs(games_path, workers, chunk_size): """(scores, outcomes) numpy arrays for every non-exact position across all games in this file, built by sharding games across --workers processes.""" score_pieces = [] outcome_pieces = [] n_unknown_result = 0 chunks = iter_chunks(games_path, chunk_size) with Pool(processes=workers) as pool: for scores, outcomes, n_unk in pool.imap(collect_chunk, chunks): if len(scores): score_pieces.append(scores) outcome_pieces.append(outcomes) n_unknown_result += n_unk if not score_pieces: return np.empty(0), np.empty(0), n_unknown_result return (np.concatenate(score_pieces), np.concatenate(outcome_pieces), n_unknown_result) def label_chunk(args): """Runs in a worker process: labels every non-exact position in a chunk of game lines (vectorized per game with numpy) and returns the re-serialized lines plus counts, in the same order they came in.""" lines, k = args out_lines = [] n_games = 0 n_positions = 0 n_exact_skipped = 0 n_unknown_result = 0 for line in lines: g = json.loads(line) white_score = result_to_white_score(g["result"]) positions = g["positions"] n_positions += len(positions) if white_score is None: n_unknown_result += 1 out_lines.append(json.dumps(g)) continue idxs = [i for i, p in enumerate(positions) if not p.get("exact")] n_exact_skipped += len(positions) - len(idxs) if idxs: scores = np.fromiter((positions[i]["score"] for i in idxs), dtype=np.float64, count=len(idxs)) is_white = np.fromiter( (positions[i]["fen"].split()[1] == "w" for i in idxs), dtype=bool, count=len(idxs)) outcomes = np.where(is_white, white_score, 1.0 - white_score) p_search = sigmoid(scores, k) for j, i in enumerate(idxs): positions[i]["p_search"] = float(p_search[j]) positions[i]["r_outcome"] = float(outcomes[j]) out_lines.append(json.dumps(g)) n_games += 1 return out_lines, n_games, n_positions, n_exact_skipped, n_unknown_result def label_file(in_path, out_path, k, workers, chunk_size): n_games = 0 n_positions = 0 n_exact_skipped = 0 n_unknown_result = 0 chunks = ((chunk, k) for chunk in iter_chunks(in_path, chunk_size)) with Pool(processes=workers) as pool, open(out_path, "w") as fout: for (out_lines, chunk_games, chunk_positions, chunk_exact, chunk_unknown) in pool.imap(label_chunk, chunks): for line in out_lines: fout.write(line + "\n") n_games += chunk_games n_positions += chunk_positions n_exact_skipped += chunk_exact n_unknown_result += chunk_unknown 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) ap.add_argument("--workers", type=int, default=16, help="worker processes for both the calibration-pair " "collection and the labeling pass (default 16, " "kept below this box's 24 cores to leave other " "jobs some breathing room)") ap.add_argument("--chunk-size", type=int, default=2000, help="games per unit of work handed to a worker") ap.add_argument("--fixed-k", type=float, default=None, help="use this K directly instead of re-fitting it " "from --train. 2026-10-16: needed when labeling " "a small supplementary slice (e.g. a self-play " "augmentation batch) to continue training an " "EXISTING checkpoint -- that checkpoint's " "weights were fit against labels built with its " "own K (nnue_model.py's DEFAULT_K unless noted " "otherwise), and re-fitting K from a small, " "distributionally skewed slice (e.g. one " "deliberately overweighted toward decisive-" "material positions) would silently relabel " "everything on a different, inconsistent scale " "-- the exact mismatch this module's own closing " "warning describes. Omit for a normal from-" "scratch fit on a large, representative pool.") args = ap.parse_args() if args.fixed_k is not None: k = args.fixed_k print(f"Using fixed K = {k:.4f} (not re-fit from {args.train!r})") else: print(f"Collecting calibration pairs from {args.train} (TRAIN only) ...") scores, outcomes, n_unknown = collect_calibration_pairs( args.train, args.workers, args.chunk_size) print(f" {len(scores)} (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(scores, outcomes) final_mse = mse(scores, outcomes, 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, args.workers, args.chunk_size) 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, args.workers, args.chunk_size) 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()