summaryrefslogtreecommitdiff
path: root/annotate_pgns/extract_positions.py
blob: 507f95696832e7eb8414b19dbc16912f26bef80d (plain)
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
#!/usr/bin/env python3
"""Extract (FEN, score, depth, PV, exact-flag) tuples directly from a
completed annotate.py run's raw shard logs -- shard_NN.script (ordered
FENs) zipped positionally against shard_NN.out (ordered tellothers
results), bypassing PGN reconstruction entirely.

Why this exists: annotate.py's TELLOTHERS_RE only matched plain-float
scores, silently dropping +MATE<N>/-MATE<N> lines from the parsed
results list while leaving them in the position index -- a positional
zip() between the two then misaligned every later position in that
shard. This script fixes the parse (captures MATE and <EGTB>/<REC>
lines as real, index-aligned results, tagged exact=True) and reads
straight from the logs already on disk, so the 12-hour run does not
need to be repeated.

Exact positions (mate-distance scores, or any <EGTB>/<REC> PV --
EGTB/REC lines are game-theoretic/exact regardless of what their score
field looks like, e.g. the engine prints +0.00 or +-MATE100 for EGTB
hits, never a real centipawn value) are written out but flagged
exact=True rather than silently dropped, per NNUE.md's existing
decision to exclude exact-value positions from the sigmoid-fit WDL
loss at training-set-assembly time -- not at extraction time, so the
raw corpus keeps standalone analysis value.

Usage:
    python3 extract_positions.py --rundir runs/sd10_twic_2M \
        --out runs/sd10_twic_2M/positions.jsonl
"""
import argparse
import json
import re
import sys
from pathlib import Path

TELLOTHERS_RE = re.compile(
    r"^tellothers d(\d+), ([+-]?[\w.]+), [\d.]+s,\s*\d+knps, PV=(.*)$"
)
FLOAT_RE = re.compile(r"^[+-]?\d+\.\d+$")
MATE_RE = re.compile(r"^([+-])MATE(\d+)$")


def parse_score(token):
    """Returns (numeric_score_or_None, exact, reason)."""
    if FLOAT_RE.match(token):
        return float(token), False, None
    m = MATE_RE.match(token)
    if m:
        sign, n = m.groups()
        return (1 if sign == "+" else -1) * float(n), True, f"mate{n}"
    return None, True, f"unrecognized:{token}"


def extract_shard(script_path, out_path):
    fens = []
    with open(script_path, encoding="utf-8") as f:
        for line in f:
            if line.startswith("setboard "):
                fens.append(line[len("setboard "):].strip())

    results = []
    with open(out_path, encoding="utf-8", errors="replace") as f:
        for line in f:
            m = TELLOTHERS_RE.match(line.rstrip("\n"))
            if m:
                depth, score_tok, pv = m.groups()
                pv = pv.strip()
                score, exact, reason = parse_score(score_tok)
                # <EGTB>/<REC>/<=>/<#> are terminal sentinel moves appended
                # to the PV (UpdatePV(), search.c) -- the PV may carry real
                # moves before the sentinel (e.g.
                # "Kf6 Bc3+ Kxg6 Kf8 Kxh5 <EGTB>"), and in that case the
                # score field often looks like an ordinary centipawn number,
                # but the position is still exact/game-theoretic, not
                # sigmoid-fit training signal -- check the PV's last token,
                # not the whole PV string. <=> (DRAWMOVE) specifically is a
                # claimed repetition/50-move draw -- a fact about the game's
                # move history, not the position itself (the same position
                # reached a different way isn't a draw at all), so it's
                # excluded for the same reason <EGTB>/<REC> are: not
                # position-dependent training signal.
                last_move = pv.rsplit(" ", 1)[-1] if pv else ""
                if last_move in ("<EGTB>", "<REC>", "<=>", "<#>"):
                    exact = True
                    reason = {"<=>": "draw_recognizer",
                              "<#>": "mate_recognizer"}.get(
                                  last_move, last_move.strip("<>").lower())
                elif pv == "only":
                    # root.c: exactly one legal move at the root skips the
                    # search entirely (iRootScore/uRootDepth hardcoded to
                    # 0), so depth/score here are placeholders, not a real
                    # evaluation -- not indeterminate, definitely not
                    # trainable.
                    exact = True
                    reason = "unsearched_only_move"
                results.append((int(depth), score, score_tok, pv, exact, reason))

    if len(fens) != len(results):
        print(f"  MISALIGNED: {script_path.name} has {len(fens)} positions "
              f"but {out_path.name} produced {len(results)} tellothers lines "
              f"-- skipping this shard, needs manual inspection",
              file=sys.stderr)
        return []

    return list(zip(fens, results))


def main():
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--rundir", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--shards", type=int, default=16)
    args = ap.parse_args()

    rundir = Path(args.rundir)
    total = 0
    total_exact = 0
    with open(args.out, "w") as out:
        for i in range(args.shards):
            script_path = rundir / f"shard_{i:02d}.script"
            out_path = rundir / f"shard_{i:02d}.out"
            if not script_path.exists() or not out_path.exists():
                continue
            rows = extract_shard(script_path, out_path)
            for fen, (depth, score, score_tok, pv, exact, reason) in rows:
                rec = {
                    "fen": fen,
                    "depth": depth,
                    "score": score,
                    "score_raw": score_tok,
                    "pv": pv,
                    "exact": exact,
                }
                if reason:
                    rec["exact_reason"] = reason
                # Mover-relative exact outcome, straight from the sign
                # of the already-mover-relative score (unlike
                # extract_selfplay_positions.py's foreign-engine path,
                # typhoon's own "score" field here is never nulled out
                # for mate/EGTB/REC rows -- only "only" (unsearched) and
                # genuinely unparseable tokens carry no usable score).
                # draw_recognizer ("<=>") is 0.5 regardless of sign.
                if exact and reason == "draw_recognizer":
                    rec["r_outcome"] = 0.5
                elif exact and score is not None and reason not in (
                        "unsearched_only_move",) and not (
                        reason and reason.startswith("unrecognized")):
                    rec["r_outcome"] = 1.0 if score > 0 else 0.0
                out.write(json.dumps(rec) + "\n")
                total += 1
                if exact:
                    total_exact += 1
            print(f"  shard {i:02d}: {len(rows)} positions "
                  f"({sum(1 for _, r in rows if r[4])} exact)")

    print(f"Wrote {args.out}: {total} positions total, "
          f"{total_exact} exact (mate/EGTB/REC), "
          f"{total - total_exact} usable for sigmoid-fit training")


if __name__ == "__main__":
    main()