Spaces:
Sleeping
Sleeping
File size: 3,131 Bytes
a5f2e46 | 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 | import chess
_PIECE_TOKEN = {
(chess.PAWN, chess.WHITE): 1, (chess.KNIGHT, chess.WHITE): 2,
(chess.BISHOP, chess.WHITE): 3, (chess.ROOK, chess.WHITE): 4,
(chess.QUEEN, chess.WHITE): 5, (chess.KING, chess.WHITE): 6,
(chess.PAWN, chess.BLACK): 7, (chess.KNIGHT, chess.BLACK): 8,
(chess.BISHOP, chess.BLACK): 9, (chess.ROOK, chess.BLACK): 10,
(chess.QUEEN, chess.BLACK): 11, (chess.KING, chess.BLACK): 12,
}
class PostionTokenizer:
"""A simple tokenizer for chess board positions.
Each square on the board is represented by a token ID based on the piece occupying it.
Empty squares are represented by the token ID 0.
The board is represented as an 8x8 grid, flattened into a list of 64 tokens.
"""
def __init__(self):
self.vocab = {
"P": 1,
"N": 2,
"B": 3,
"R": 4,
"Q": 5,
"K": 6,
"p": 7,
"n": 8,
"b": 9,
"r": 10,
"q": 11,
"k": 12,
".": 0,
}
self.inv_vocab = {v: k for k, v in self.vocab.items()}
def encode(self, board: chess.Board) -> list[int]:
"""Encode a chess.Board object into a list of token IDs.
The board is represented as an 8x8 grid, flattened into a list of 64 tokens.
Args:
board: A chess.Board object
Returns:
List of token IDs representing the board position
"""
result = [0] * 64
for sq, piece in board.piece_map().items():
result[sq] = _PIECE_TOKEN[(piece.piece_type, piece.color)]
return result
def decode(self, token_ids: list[int]) -> chess.Board:
"""Decode a list of token IDs back into a chess board string.
Args:
token_ids: List of token IDs representing the board position
Returns:
A chess.Board object
"""
chars = [self.inv_vocab[token_id] for token_id in token_ids]
rows = ["".join(chars[i * 8 : (i + 1) * 8]) for i in range(8)][::-1] # Reverse to get the correct order
board_str = "\n".join(rows)
board = self._ascii2board(board_str)
return board
def _ascii2board(self, ascii_board: str) -> chess.Board:
"""Convert an ASCII board representation back to a chess.Board object.
Args:
ascii_board: String representation of the board (8 lines of 8 characters)
Returns:
A chess.Board object
"""
board = chess.Board.empty()
rows = ascii_board.split("\n")
for rank in range(8):
file = 0
for char in rows[7 - rank]: # Reverse the order of ranks
if char in self.vocab and char != ".":
square = chess.square(file, rank)
piece = chess.Piece.from_symbol(char)
board.set_piece_at(square, piece)
file += 1
return board
@property
def vocab_size(self) -> int:
"""Return the size of the vocabulary."""
return len(self.vocab)
|