tokenizer.py

1"""One authoritative codec for data preparation and inference (no model dependency)."""
2import chess
3
4VERSION = "chess-square-triple-v1"
5BOS, EOS, PAD = 69, 70, 71
6PROMOTIONS = {None: 64, chess.KNIGHT: 65, chess.BISHOP: 66, chess.ROOK: 67, chess.QUEEN: 68}
7REVERSE = {v: k for k, v in PROMOTIONS.items()}
8
9
10def encode_move(move):
11    if (not move or move.drop or move.promotion not in PROMOTIONS
12            or any(type(s) is not int or not 0 <= s < 64 for s in (move.from_square, move.to_square))):
13        raise ValueError("null/drop/unsupported promotion")
14    return [move.from_square, move.to_square, PROMOTIONS[move.promotion]]
15
16
17def decode_move(tokens):
18    if (len(tokens) != 3 or any(type(t) is not int for t in tokens)
19            or not all(0 <= t < 64 for t in tokens[:2]) or tokens[2] not in REVERSE):
20        raise ValueError("invalid triple or phase")
21    move = chess.Move(tokens[0], tokens[1], REVERSE[tokens[2]])
22    if not move:
23        raise ValueError("null move")
24    return move
25
26
27class State:
28    """Authoritative board + phase. Cropping model history never mutates this state.
29
30    A board may be supplied for tests/serving snapshots; PGN v1 accepts only standard
31    starts. Board copies prevent caller mutation. Token operations are transactional.
32    EOS records an observed ending, including resignation; it is not a mate oracle.
33    """
34    def __init__(self, board=None):
35        # Variant subclasses can emit ordinary triples with different semantics.
36        # Require the exact standard board class, including for serving snapshots.
37        if board is not None and type(board) is not chess.Board:
38            raise ValueError("only exact standard chess.Board snapshots are supported")
39        self._board = (board if board is not None else chess.Board()).copy(stack=True)
40        if not self._board.is_valid() or self._board.chess960:
41            raise ValueError("invalid/unsupported board")
42        self._pending = []
43        self.started = False
44        self.ended = False
45
46    @property
47    def board(self):
48        return self._board.copy(stack=True)
49
50    @property
51    def phase(self):
52        return "ended" if self.ended else ("bos" if not self.started else ("source", "destination", "promotion")[len(self._pending)])
53
54    def mask(self, legal=False):
55        """128 bools, separate from raw model probabilities; never renormalizes.
56
57        legal=False enforces grammar/reserved IDs only. legal=True additionally
58        restricts move prefixes to legal continuations. EOS is allowed at every
59        move boundary, even nonterminal positions. PAD is a batching operation.
60        """
61        ids = set()
62        if not self.ended:
63            if not self.started:
64                ids = {BOS}
65            elif not legal:
66                ids = set(range(64) if len(self._pending) < 2 else range(64, 69))
67                if not self._pending:
68                    ids.add(EOS)
69            else:
70                for move in self._board.legal_moves:
71                    triple = encode_move(move)
72                    if triple[:len(self._pending)] == self._pending:
73                        ids.add(triple[len(self._pending)])
74                if not self._pending:
75                    ids.add(EOS)
76        return [i in ids for i in range(128)]
77
78    def feed(self, token):
79        if type(token) is not int or not 0 <= token < 128 or not self.mask(legal=True)[token]:
80            raise ValueError("illegal token, reserved ID, or incorrect phase")
81        if token == BOS:
82            self.started = True
83        elif token == EOS:
84            self.ended = True
85        else:
86            pending = self._pending + [token]
87            if len(pending) == 3:
88                move = decode_move(pending)
89                self._board.push(move)
90                self._pending = []
91                return move
92            self._pending = pending
93        return None
94
95    def append_move(self, move):
96        if not self.started or self.ended or self._pending or move not in self._board.legal_moves:
97            raise ValueError("illegal move or incorrect phase")
98        triple = encode_move(move)
99        self._board.push(move)
100        return triple
101
102    def extend_uci(self, moves):
103        """Incremental history encoding, returning only newly appended tokens."""
104        result = []
105        for uci in moves:
106            result.extend(self.append_move(chess.Move.from_uci(uci)))
107        return result
108
109    def finish(self, require_eos=True):
110        if not self.started or self._pending or (require_eos and not self.ended):
111            raise ValueError("incomplete token sequence")
112        return self.board
113
114
115def windows(tokens, length, stride=None):
116    """Unpacked, single-game next-token examples; no invented BOS/EOS at crops.
117
118    Returned lists have variable length. Collators pad x/y with PAD and loss_mask
119    with False; attention is causal within real tokens, never across examples.
120    Absolute token start identifies phase; it cannot recover board state.
121    """
122    stride = length if stride is None else stride
123    if length < 1 or not 1 <= stride <= length:
124        raise ValueError("invalid window length/stride")
125    state = State()
126    for token in tokens:
127        state.feed(token)
128    state.finish()
129    for start in range(0, len(tokens) - 1, stride):
130        x = tokens[start:start + length]
131        y = tokens[start + 1:start + length + 1]
132        x = x[:len(y)]
133        yield {"input_ids": x, "labels": y, "loss_mask": [True] * len(y),
134               "token_start": start, "attention": "single-game-causal"}
135