path: root/annotate_pgns/fit_k_and_label.py

#!/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()