# Version (Player (Color + Piece), Source_S, Destination_D, Suffix) from __future__ import annotations import json import os from pathlib import Path from typing import Dict, List, Optional from transformers import PreTrainedTokenizer class ChessTokenizer(PreTrainedTokenizer): """ Sub-move tokenizer for chess moves using extended UCI notation. This tokenizer splits each move into atomic components: - Players (color + piece): WP, WN, WB, WR, WQ, WK, etc. - Source square: e2 - Destination square: e4 - Optional suffixes: x (capture), + (check), * (checkmate), o/O (castling) Example: Move "WPe2e4(x+)" -> ["WP", "e2_S", "e4_D", "(x+)"] """ model_input_names = ["input_ids", "attention_mask"] vocab_files_names = {"vocab_file": "vocab.json"} # Special tokens PAD_TOKEN = "[PAD]" BOS_TOKEN = "[BOS]" EOS_TOKEN = "[EOS]" UNK_TOKEN = "[UNK]" # Atomic suffix tokens for default vocab SUFFIX_TOKENS = ["(x)", "(+)", "(*)", "(o)", "(O)", "(+*)", "(x+)"] def __init__( self, vocab_file: Optional[str] = None, vocab: Optional[Dict[str, int]] = None, **kwargs, ): # Special tokens self._pad_token = self.PAD_TOKEN self._bos_token = self.BOS_TOKEN self._eos_token = self.EOS_TOKEN self._unk_token = self.UNK_TOKEN # Remove duplicates from kwargs kwargs.pop("pad_token", None) kwargs.pop("bos_token", None) kwargs.pop("eos_token", None) kwargs.pop("unk_token", None) # Load or create vocab if vocab is not None: self._vocab = vocab elif vocab_file is not None and os.path.exists(vocab_file): with open(vocab_file, "r", encoding="utf-8") as f: self._vocab = json.load(f) else: self._vocab = self._create_default_vocab() # Reverse mapping self._ids_to_tokens = {v: k for k, v in self._vocab.items()} super().__init__( pad_token=self._pad_token, bos_token=self._bos_token, eos_token=self._eos_token, unk_token=self._unk_token, **kwargs, ) def _create_default_vocab(self) -> Dict[str, int]: """ Build a fixed vocab based on chess grammar for sub-moves. Useful for predefined grammar instead of dataset-based vocab. """ colors = ["W", "B"] pieces = ["P", "N", "B", "R", "Q", "K"] files = ["a", "b", "c", "d", "e", "f", "g", "h"] ranks = ["1", "2", "3", "4", "5", "6", "7", "8"] squares = [f + r for f in files for r in ranks] players = [c + p for c in colors for p in pieces] # Source and destination tokens sources = [sq + "_S" for sq in squares] dests = [sq + "_D" for sq in squares] # Build all possible sub-tokens vocab_tokens = players + sources + dests + self.SUFFIX_TOKENS # Add special tokens at the start special_tokens = [self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN] vocab = {token: idx for idx, token in enumerate(special_tokens + vocab_tokens)} return vocab def _tokenize(self, text: str) -> List[str]: """ Convert a string of moves into sub-move tokens. """ tokens: List[str] = [] moves = text.strip().split() for move in moves: if not move: continue # Color + Piece tokens.append(move[:2]) # WP, BN, etc. # Source square with _S tokens.append(move[2:4] + "_S") # Destination square with _D tokens.append(move[4:6] + "_D") if (len(move)>6): tokens.append(move[6:]) return tokens def _convert_token_to_id(self, token: str) -> int: return self._vocab.get(token, self._vocab.get(self.UNK_TOKEN, 0)) def _convert_id_to_token(self, index: int) -> str: return self._ids_to_tokens.get(index, self.UNK_TOKEN) def convert_tokens_to_string(self, tokens: List[str]) -> str: """Convert a list of tokens back to a string.""" # Filter out special tokens for cleaner output special = {self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN} clean_tokens = [] for t in tokens: if t in special: continue # Remove everything from _ onward if "_" in t: clean_tokens.append(t.split("_")[0]) else: clean_tokens.append(t) result = "" temp = "".join(token for token in clean_tokens) for i, str in enumerate(temp): if str in ["W", "B"]: if result == "": result += str elif temp[i-1].isnumeric() or temp[i-1]==")": result += " " + str else : result += str else : result += str return result.split()[0] @property def vocab_size(self) -> int: return len(self._vocab) def get_vocab(self) -> Dict[str, int]: return dict(self._vocab) def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> tuple: if not os.path.isdir(save_directory): os.makedirs(save_directory, exist_ok=True) vocab_file = os.path.join( save_directory, (filename_prefix + "-" if filename_prefix else "") + "vocab.json", ) with open(vocab_file, "w", encoding="utf-8") as f: json.dump(self._vocab, f, ensure_ascii=False, indent=2) return (vocab_file,) @classmethod def build_vocab_from_iterator(cls, iterator, min_frequency: int = 1) -> "ChessTokenizer": """ Build vocab from dataset iterator using sub-move tokens. """ from collections import Counter token_counts = Counter() for game in iterator: sub_tokens = cls()._tokenize(game) token_counts.update(sub_tokens) tokens = [token for token, count in token_counts.items() if count >= min_frequency] tokens = sorted(tokens) special_tokens = [cls.PAD_TOKEN, cls.BOS_TOKEN, cls.EOS_TOKEN, cls.UNK_TOKEN] vocab = {token: idx for idx, token in enumerate(special_tokens + tokens)} return cls(vocab=vocab) @classmethod def build_vocab_from_dataset( cls, dataset_name: str = "dlouapre/lichess_2025-01_1M", split: str = "train", column: str = "text", min_frequency: int = 500, max_samples: Optional[int] = 100000, ) -> "ChessTokenizer": from datasets import load_dataset dataset = load_dataset(dataset_name, split=split) if max_samples is not None: dataset = dataset.select(range(min(max_samples, len(dataset)))) def game_iterator(): for example in dataset: yield example[column] return cls.build_vocab_from_iterator(game_iterator(), min_frequency=min_frequency) def count_vocab_from_dataset( dataset_name: str = "dlouapre/lichess_2025-01_1M", split: str = "train", column: str = "text", max_samples: Optional[int] = 10000, ) -> Dict[str, int]: """ Count sub-move token frequencies in a dataset (useful for vocab analysis). """ from collections import Counter from datasets import load_dataset dataset = load_dataset(dataset_name, split=split) if max_samples is not None: dataset = dataset.select(range(min(max_samples, len(dataset)))) token_counts = Counter() for example in dataset: moves = example[column].strip().split() # Use sub-tokenization tokenizer = ChessTokenizer() for move in moves: token_counts.update(tokenizer._tokenize(move)) return dict(token_counts)