File size: 3,791 Bytes
3495881 79997ea 3495881 79997ea 3495881 79997ea 3495881 | 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 | """
Batching helpers: python-chess boards -> model tensors.
Contract (critical for correctness, tested in tests/test_phase1.py):
* Every board is canonicalised to side-to-move == White.
* Candidate moves are generated ON THE CANONICAL BOARD, so index i of the
logits always corresponds to `cand_moves[b][i]` in canonical space, and to
`real_moves[b][i]` on the original board.
* The value head therefore always predicts from the SIDE-TO-MOVE perspective.
"""
from __future__ import annotations
from typing import Sequence
import chess
import numpy as np
import torch
from .encoding import (encode_position, pack_move, unpack_fields, canonical_board,
flip_move)
class PositionBatch:
"""Tensors for a batch of positions plus the move bookkeeping."""
__slots__ = ("squares", "extras", "glob", "cand_fields", "cand_mask",
"real_moves", "n_legal")
def __init__(self, squares, extras, glob, cand_fields, cand_mask, real_moves,
n_legal):
self.squares = squares
self.extras = extras
self.glob = glob
self.cand_fields = cand_fields
self.cand_mask = cand_mask
self.real_moves = real_moves # list[list[chess.Move]] in ORIGINAL board space
self.n_legal = n_legal
def to(self, device):
self.squares = self.squares.to(device)
self.extras = self.extras.to(device)
self.glob = self.glob.to(device)
self.cand_fields = self.cand_fields.to(device)
self.cand_mask = self.cand_mask.to(device)
return self
def model_args(self):
return dict(squares=self.squares, extras=self.extras, glob=self.glob,
cand_fields=self.cand_fields, cand_mask=self.cand_mask)
def encode_batch(boards: Sequence[chess.Board],
last_moves: Sequence[chess.Move | None] | None = None,
max_cand: int | None = None,
reps: Sequence[int] | None = None) -> PositionBatch:
"""Encode boards + their legal candidate sets into padded tensors."""
n = len(boards)
if last_moves is None:
last_moves = [None] * n
sq_l, ex_l, g_l, packed_l, real_l = [], [], [], [], []
for bi, (b, lm) in enumerate(zip(boards, last_moves)):
if reps is not None:
rep = int(reps[bi]) # explicit: FEN-restored boards lack history
else:
rep = 0
try:
if b.is_repetition(2):
rep = 2
except Exception:
rep = 0
sq, ex, g, flipped = encode_position(b, lm, rep)
cb = b.mirror() if flipped else b
real, packed = [], []
for m in cb.legal_moves:
packed.append(pack_move(cb, m))
real.append(flip_move(m) if flipped else m)
sq_l.append(sq); ex_l.append(ex); g_l.append(g)
packed_l.append(packed); real_l.append(real)
L = max((len(p) for p in packed_l), default=1)
if max_cand is not None:
L = min(L, max_cand)
L = max(L, 1)
cand = np.zeros((n, L, 7), dtype=np.int64)
mask = np.zeros((n, L), dtype=bool)
n_legal = np.zeros(n, dtype=np.int64)
for i, p in enumerate(packed_l):
k = min(len(p), L)
n_legal[i] = k
if k:
cand[i, :k] = unpack_fields(np.asarray(p[:k], dtype=np.int64))
mask[i, :k] = True
real_l[i] = real_l[i][:k]
return PositionBatch(
squares=torch.from_numpy(np.stack(sq_l)).long(),
extras=torch.from_numpy(np.stack(ex_l)).float(),
glob=torch.from_numpy(np.stack(g_l)).float(),
cand_fields=torch.from_numpy(cand),
cand_mask=torch.from_numpy(mask),
real_moves=real_l,
n_legal=torch.from_numpy(n_legal),
)
|