File size: 3,626 Bytes
1829c5b | 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 | #!/usr/bin/env python3
"""Chess encoding: board positions and moves as integer sequences.
No external tokenizer needed — pure chess-aware encoding.
"""
import chess
# Piece values (0 = empty, 1-6 = white, 7-12 = black)
PIECE_VALUES = {
None: 0,
chess.Piece(chess.PAWN, chess.WHITE): 1,
chess.Piece(chess.KNIGHT, chess.WHITE): 2,
chess.Piece(chess.BISHOP, chess.WHITE): 3,
chess.Piece(chess.ROOK, chess.WHITE): 4,
chess.Piece(chess.QUEEN, chess.WHITE): 5,
chess.Piece(chess.KING, chess.WHITE): 6,
chess.Piece(chess.PAWN, chess.BLACK): 7,
chess.Piece(chess.KNIGHT, chess.BLACK): 8,
chess.Piece(chess.BISHOP, chess.BLACK): 9,
chess.Piece(chess.ROOK, chess.BLACK): 10,
chess.Piece(chess.QUEEN, chess.BLACK): 11,
chess.Piece(chess.KING, chess.BLACK): 12,
}
# Special tokens
PAD_TOKEN = 0
BOS_TOKEN = 13
EOS_TOKEN = 14
SEP_TOKEN = 15
SIDE_WHITE_TOKEN = 16
SIDE_BLACK_TOKEN = 17
MOVE_FROM_OFFSET = 18 # 18 + square (0-63) = from-square token
MOVE_TO_OFFSET = 82 # 82 + square (0-63) = to-square token
# Total vocab: 82 + 64 = 146 + promotions
# Promotion tokens: 146 + promo_type(0-3) * 64 + to_square
PROMO_OFFSET = 146 # 146 + promo*64 + to_square
# promo: 0=queen, 1=rook, 2=bishop, 3=knight
VOCAB_SIZE = PROMO_OFFSET + 4 * 64 # = 402
def encode_board(board):
"""Encode chess board to token sequence.
Returns list of ints: [BOS, piece(64x), side, SEP]
"""
tokens = [BOS_TOKEN]
for square in chess.SQUARES:
piece = board.piece_at(square)
tokens.append(PIECE_VALUES[piece])
tokens.append(SIDE_WHITE_TOKEN if board.turn == chess.WHITE else SIDE_BLACK_TOKEN)
tokens.append(SEP_TOKEN)
return tokens
def encode_move(move):
"""Encode a move to token sequence.
Returns [from_square_token, to_square_token] or [from_square_token, promo_token]
"""
from_sq = move.from_square
to_sq = move.to_square
tokens = [MOVE_FROM_OFFSET + from_sq]
if move.promotion is not None:
promo_map = {chess.QUEEN: 0, chess.ROOK: 1, chess.BISHOP: 2, chess.KNIGHT: 3}
promo_type = promo_map[move.promotion]
tokens.append(PROMO_OFFSET + promo_type * 64 + to_sq)
else:
tokens.append(MOVE_TO_OFFSET + to_sq)
return tokens
def decode_move(tokens, board):
"""Decode token sequence back to a chess.Move.
Input: list of token IDs (just the move tokens, not board tokens).
"""
if len(tokens) < 2:
return None
from_token = tokens[0]
to_token = tokens[1]
from_sq = from_token - MOVE_FROM_OFFSET
if from_sq < 0 or from_sq > 63:
return None
if to_token >= PROMO_OFFSET:
promo_offset = to_token - PROMO_OFFSET
promo_type = promo_offset // 64
to_sq = promo_offset % 64
promo_map = {0: chess.QUEEN, 1: chess.ROOK, 2: chess.BISHOP, 3: chess.KNIGHT}
promotion = promo_map.get(promo_type)
move = chess.Move(from_sq, to_sq, promotion=promotion)
else:
to_sq = to_token - MOVE_TO_OFFSET
if to_sq < 0 or to_sq > 63:
return None
move = chess.Move(from_sq, to_sq)
if move in board.legal_moves:
return move
return None
def encode_game(board, move):
"""Encode a board+move pair for training.
Input sequence: board_tokens + move_tokens
Target: move_tokens (shifted by 1)
"""
board_tokens = encode_board(board)
move_tokens = encode_move(move)
full = board_tokens + move_tokens
return full
def decode_position_to_tokens(board):
"""Get board encoding for inference."""
return encode_board(board)
|