""" Custom Chess Tokenizer for the Chess Challenge. We build a vocabulary with: - W/B prefix for White/Black - Piece letter: P=Pawn, N=Knight, B=Bishop, R=Rook, Q=Queen, K=King - Source and rank and file: e.g e 2 - Destination and rank and file: e.g e 4 - Special suffixes: (x)=capture, (+)=check, (+*)=checkmate, (o)/(O)=castling """ from __future__ import annotations import json import os from pathlib import Path import shutil import inspect from typing import Dict, List, Optional from transformers import PreTrainedTokenizer class ChessTokenizer(PreTrainedTokenizer): """ A custom tokenizer for chess moves. Example: >>> tokenizer = ChessTokenizer() >>> tokenizer.encode("WPe2e4 BPe7e5") # [BOS, W, P, e, 2, e, 4, B, P, e, 7, e, 5, EOS] """ 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]" SEP_TOKEN = "[SEP]" def __init__( self, vocab_file: Optional[str] = None, vocab: Optional[Dict[str, int]] = None, **kwargs, ): """ Initialize the chess tokenizer. Args: vocab_file: Path to a JSON file containing the vocabulary mapping. vocab: Dictionary mapping tokens to IDs (alternative to vocab_file). **kwargs: Additional arguments passed to PreTrainedTokenizer. """ # Initialize 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 self._sep_token = self.SEP_TOKEN # Remove any duplicate special-token entries passed through kwargs # to avoid "multiple values for keyword" errors when loading from disk. kwargs.pop("pad_token", None) kwargs.pop("bos_token", None) kwargs.pop("eos_token", None) kwargs.pop("unk_token", None) kwargs.pop("sep_token", None) # Load or create vocabulary 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: # Create a minimal vocabulary with just special tokens # The full vocabulary should be built from the dataset self._vocab = self._create_default_vocab() # Create reverse mapping self._ids_to_tokens = {v: k for k, v in self._vocab.items()} # Call parent init AFTER setting up vocab super().__init__( pad_token=self._pad_token, bos_token=self._bos_token, eos_token=self._eos_token, unk_token=self._unk_token, sep_token=self._sep_token, **kwargs, ) def _create_default_vocab(self) -> Dict[str, int]: """ Create a minimal default vocabulary with just special tokens. For the full vocabulary, use `build_vocab_from_dataset()`. This minimal vocab is just a placeholder - you should build from data. """ special_tokens = [self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN, self.SEP_TOKEN] vocab = {token: idx for idx, token in enumerate(special_tokens)} return vocab @classmethod def build_vocab_from_dataset( cls, dataset_name: str = "dlouapre/lichess_2025-01_1M", split: str = "train", column: str = "text", save_path: Optional[str] = None, ) -> "ChessTokenizer": """ Build a tokenizer vocabulary from a Hugging Face dataset. Args: dataset_name: Name of the dataset on Hugging Face Hub. split: Dataset split to use. column: Column containing the game strings. Returns: A ChessTokenizer with the built vocabulary. Args: save_path: Optional path to write the generated vocab JSON. If not provided, the vocab will be saved to ``./chess_tokenizer_vocab.json``. """ from datasets import load_dataset # If a saved vocab exists at `save_path`, load it and return a tokenizer if save_path is None: cwd = os.getcwd() save_path = os.path.join(cwd, "chess_tokenizer_vocab.json") if os.path.exists(save_path): try: with open(save_path, "r", encoding="utf-8") as f: print("Loading existing tokenizer vocab from", save_path) vocab = json.load(f) return cls(vocab=vocab) except Exception: # If loading fails, fall through to rebuild the vocab. pass dataset = load_dataset(dataset_name, split=split) # Iterator over games (respect max_samples if provided) samples = dataset[column] tokens = set() for game in samples: if not isinstance(game, str): continue moves = game.strip().split() for move in moves: # Basic parsing of move token components if len(move) < 2: continue color = move[0] piece = move[1] from_square = move[2:4] if len(move) >= 4 else '' to_square = move[4:6] if len(move) >= 6 else '' suffix = move[6:] if len(move) > 6 else '' tokens.add(color) tokens.add(piece) tokens.add(from_square) tokens.add(to_square) if suffix: tokens.add(suffix) # Sort tokens tokens = sorted(tokens) # Ensure special tokens are present at fixed ids special_tokens = [cls.PAD_TOKEN, cls.BOS_TOKEN, cls.EOS_TOKEN, cls.UNK_TOKEN, cls.SEP_TOKEN] # Build vocab mapping: special tokens first, then tokens vocab: Dict[str, int] = {} idx = 0 for st in special_tokens: vocab[st] = idx idx += 1 for t in tokens: if t in vocab: continue vocab[t] = idx idx += 1 # Create tokenizer instance with this vocab tokenizer = cls(vocab=vocab) # Save vocab to disk. Use provided `save_path` or default file name. try: if save_path is None: cwd = os.getcwd() save_path = os.path.join(cwd, "chess_tokenizer_vocab.json") # Write to a temporary file first and atomically replace final file. tmp_path = save_path + ".tmp" with open(tmp_path, "w", encoding="utf-8") as f: json.dump(vocab, f, ensure_ascii=False, indent=2) os.replace(tmp_path, save_path) except Exception: # Non-fatal: ignore save errors but don't leave temp files behind. try: if 'tmp_path' in locals() and os.path.exists(tmp_path): os.remove(tmp_path) except Exception: pass return tokenizer @property def vocab_size(self) -> int: """Return the size of the vocabulary.""" return len(self._vocab) def get_vocab(self) -> Dict[str, int]: """Return the vocabulary as a dictionary.""" return dict(self._vocab) def _tokenize(self, text: str) -> List[str]: """ Tokenize a string of moves into a list of tokens. Args: text: A string of space-separated moves. Returns: List of move tokens. """ tokens: List[str] = [] for move in text.strip().split(): if len(move) < 2: continue color, piece, from_square, to_square, suffix = self._decompose_move(move) tokens.append(color) tokens.append(piece) tokens.append(from_square) tokens.append(to_square) if suffix: tokens.append(suffix) tokens.append(self._sep_token) return tokens[:-1] # Remove last SEP token @staticmethod def _decompose_move(move: str): """Decompose a move string into components: color, piece, from_square, to_square, suffix. Returns a 5-tuple of strings (empty strings for missing parts). """ color = move[0] piece = move[1] if len(move) >= 2 else '' from_square = move[2:4] if len(move) >= 4 else '' to_square = move[4:6] if len(move) >= 6 else '' suffix = move[6:] if len(move) > 6 else '' return color, piece, from_square, to_square, suffix def _convert_token_to_id(self, token: str) -> int: """Convert a token to its ID.""" return self._vocab.get(token, self._vocab.get(self.UNK_TOKEN, 0)) def _convert_id_to_token(self, index: int) -> str: """Convert an ID to its token.""" 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} return " ".join(t for t in tokens if t not in special) def decode(self, token_ids: List[int], skip_special_tokens: bool = True) -> str: """Decode a list of token IDs back to a string.""" tokens = [self._convert_id_to_token(int(tid)) for tid in token_ids] if skip_special_tokens: special = {self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN} # SEP token should be replace by space tokens = [t if t != self.SEP_TOKEN else " " for t in tokens if t not in special] return "".join(tokens) def save_vocabulary( self, save_directory: str, filename_prefix: Optional[str] = None, ) -> tuple: """ Save the vocabulary to a JSON file. Args: save_directory: Directory to save the vocabulary. filename_prefix: Optional prefix for the filename. Returns: Tuple containing the path to the saved vocabulary file. """ 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,) def save_pretrained( self, save_directory: str, filename_prefix: Optional[str] = None, save_tokenizer_code: bool = True, ) -> None: """Save tokenizer files to a directory in a HF-compatible layout. This writes the vocab JSON (via `save_vocabulary`), a small `tokenizer_config.json` describing special tokens and the vocab filename, and optionally copies the tokenizer module source file into the directory so others can import the implementation. """ if not os.path.isdir(save_directory): os.makedirs(save_directory, exist_ok=True) # Save the vocabulary file vocab_file_tuple = self.save_vocabulary(save_directory, filename_prefix) vocab_file = vocab_file_tuple[0] # Write a minimal tokenizer config config = { "tokenizer_class": self.__class__.__name__, "vocab_file": os.path.basename(vocab_file), "pad_token": self.PAD_TOKEN, "bos_token": self.BOS_TOKEN, "eos_token": self.EOS_TOKEN, "unk_token": self.UNK_TOKEN, } config_path = os.path.join(save_directory, "tokenizer_config.json") with open(config_path, "w", encoding="utf-8") as f: json.dump(config, f, ensure_ascii=False, indent=2) # Optionally copy this module file so the tokenizer class implementation # is available alongside the saved vocab/config. This helps when # transferring the saved tokenizer to another environment. if save_tokenizer_code: try: src_file = Path(inspect.getsourcefile(self.__class__)) dst_file = Path(save_directory) / src_file.name shutil.copy2(src_file, dst_file) except Exception: # Non-fatal; we still saved vocab and config pass @classmethod def from_pretrained(cls, load_directory: str) -> "ChessTokenizer": """Load tokenizer from a directory previously written with `save_pretrained`. This primarily reads the vocab file and constructs the tokenizer. If a `tokenizer_config.json` exists it will be consulted for the vocab filename and special tokens (but we still instantiate using the provided class). """ config_path = os.path.join(load_directory, "tokenizer_config.json") vocab_file = None if os.path.exists(config_path): try: with open(config_path, "r", encoding="utf-8") as f: cfg = json.load(f) vocab_file = os.path.join(load_directory, cfg.get("vocab_file", "vocab.json")) except Exception: pass if vocab_file is None: # Fallback: look for a vocab file in the directory candidates = [p for p in os.listdir(load_directory) if p.endswith("vocab.json")] if candidates: vocab_file = os.path.join(load_directory, candidates[0]) if vocab_file is None or not os.path.exists(vocab_file): raise FileNotFoundError(f"No vocab file found in {load_directory}") return cls(vocab_file=vocab_file) 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 token frequencies in a dataset (useful for vocabulary analysis). Args: dataset_name: Name of the dataset on Hugging Face Hub. split: Dataset split to use. column: Column containing the game strings. max_samples: Maximum number of samples to process. Returns: Dictionary mapping tokens to their frequencies. """ 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)))) tokenizer = ChessTokenizer() token_counts = Counter() for example in dataset: token_counts.update(tokenizer._tokenize(example[column])) return dict(token_counts)