File size: 3,384 Bytes
d2d7586 | 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 | """Chess-board preprocessing for the Hugging Face model.
This is deliberately a board encoder rather than a text tokenizer. Each
board becomes 64 mover-relative piece IDs, and each legal move is represented
by the tokenized board after that move. Successors remain in the original
position's perspective.
"""
from __future__ import annotations
from collections.abc import Sequence
from typing import Any
import chess
import torch
EMPTY = 0
MY_PAWN = 1
MY_KNIGHT = 2
MY_BISHOP = 3
MY_ROOK = 4
MY_QUEEN = 5
MY_KING = 6
OPPONENT_OFFSET = 6
NUM_PIECE_STATES = 13
def canonical_square(square: chess.Square, perspective: chess.Color) -> int:
"""Map a square into mover-relative canonical order."""
return square if perspective == chess.WHITE else square ^ 56
def tokenize_board(
board: chess.Board, perspective: chess.Color | None = None
) -> torch.Tensor:
"""Return a board as a ``[64]`` uint8 tensor of piece-role IDs."""
if perspective is None:
perspective = board.turn
tokens = torch.zeros(64, dtype=torch.uint8)
for square, piece in board.piece_map().items():
role = piece.piece_type
if piece.color != perspective:
role += OPPONENT_OFFSET
tokens[canonical_square(square, perspective)] = role
return tokens
def build_model_inputs(
boards: Sequence[chess.Board],
*,
include_candidate_uci: bool = False,
) -> dict[str, Any]:
"""Build the tensors expected by ``ChessTransitionPolicy``.
Candidate order is the order returned by ``python-chess``. If
``include_candidate_uci`` is true, the returned dictionary also contains
``candidate_uci`` for mapping output scores back to moves.
"""
if not boards:
raise ValueError("At least one board is required")
current_boards: list[torch.Tensor] = []
successor_boards: list[torch.Tensor] = []
candidate_uci: list[tuple[str, ...]] = []
candidate_counts: list[int] = []
for board in boards:
perspective = board.turn
legal_moves = list(board.legal_moves)
if not legal_moves:
raise ValueError("Every board must have at least one legal move")
current_boards.append(tokenize_board(board, perspective))
candidate_counts.append(len(legal_moves))
if include_candidate_uci:
candidate_uci.append(tuple(move.uci() for move in legal_moves))
for move in legal_moves:
board.push(move)
successor_boards.append(tokenize_board(board, perspective))
board.pop()
counts = torch.tensor(candidate_counts, dtype=torch.int32)
offsets = torch.zeros(len(boards) + 1, dtype=torch.int64)
offsets[1:] = torch.cumsum(counts.to(torch.int64), dim=0)
owners = torch.repeat_interleave(
torch.arange(len(boards), dtype=torch.int64), counts.to(torch.int64)
)
max_candidates = max(candidate_counts)
candidate_mask = (
torch.arange(max_candidates).unsqueeze(0)
< counts.unsqueeze(1)
)
result: dict[str, Any] = {
"current_boards": torch.stack(current_boards),
"successor_boards": torch.stack(successor_boards),
"candidate_owner": owners,
"candidate_offsets": offsets,
"candidate_mask": candidate_mask,
}
if include_candidate_uci:
result["candidate_uci"] = tuple(candidate_uci)
return result
|