Hugging Face's logo

LLM-course
/
chess_swdo_s1

Text Generation
Transformers
Safetensors
chess_transformer
chess
llm-course
chess-challenge
custom_code
Model card Files
xet
 Community
chess_swdo_s1 / tokenizer.py
swdo's picture
swdo
Chess Challenge submission by swdo
17f97be verified
Raw History Blame Contribute Delete
12.7 kB
from __future__ import annotations
import json
import os
import re
from pathlib import Path
from typing import Dict, List, Optional, Tuple
from transformers import PreTrainedTokenizer
class ChessTokenizer(PreTrainedTokenizer):
"""
A custom tokenizer
Example:
>>> tokenizer = ChessTokenizer()
>>> tokenizer.encode("WPe2e4 BPe7e5")
[1, 4, 6, 45, 47, 5, 6, 50, 48, 2] # [BOS, components..., 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]"
# Chess components
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 ["1", "2", "3", "4", "5", "6", "7", "8"]]
MODIFIERS = [
"x", # Capture
"+", # Check
"#", # Checkmate (alternative to +*)
"+*", # Checkmate (dataset format)
"=Q", # Promotion to Queen
"=R", # Promotion to Rook
"=B", # Promotion to Bishop
"=N", # Promotion to Knight
"O-O", # Kingside castling (alternative)
"O-O-O", # Queenside castling (alternative)
"o", # Kingside castling (dataset format)
"O", # Queenside castling (dataset format)
]
MOVE_PATTERN = re.compile(
r'^([WB])' # Color
r'([PNBRQK])' # Piece
r'([a-h][1-8])' # From square
r'([a-h][1-8])' # To square
r'(=[QRBN])?' # Promotion (optional)
r'(\([xoO+*]+\))?$' # Suffixes in parentheses (optional)
)
def __init__(
self,
vocab_file: Optional[str] = None,
vocab: Optional[Dict[str, int]] = None,
**kwargs,
):
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 any duplicate
kwargs.pop("pad_token", None)
kwargs.pop("bos_token", None)
kwargs.pop("eos_token", None)
kwargs.pop("unk_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 the fixed decomposed vocabulary
self._vocab = self._create_default_vocab()
# Create 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]:
"""
Create the fixed vocabulary from chess components.
Unlike the standard tokenizer, this creates a small fixed vocab
of ~88 tokens for decomposed move representation.
"""
tokens = []
# Special tokens first
tokens.extend([self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN])
# Colors
tokens.extend(self.COLORS)
# Pieces
tokens.extend(self.PIECES)
# Squares (64)
tokens.extend(self.SQUARES)
# Modifiers
tokens.extend(self.MODIFIERS)
return {token: idx for idx, token in enumerate(tokens)}
@classmethod
def build_vocab_from_iterator(
cls,
iterator,
min_frequency: int = 1,
) -> "ChessTokenizer":
"""
Build a tokenizer vocabulary from an iterator of game strings.
Note: For decomposed tokenizer, this ignores the iterator and
creates the fixed vocabulary. Provided for API compatibility.
Args:
iterator: An iterator yielding game strings (ignored).
min_frequency: Minimum frequency for a token (ignored).
Returns:
A ChessTokenizer with the fixed decomposed vocabulary.
"""
# Decomposed tokenizer uses fixed vocabulary
return cls()
@classmethod
def build_vocab_from_dataset(
cls,
dataset_name: str = "dlouapre/lichess_2025-01_1M",
split: str = "train",
column: str = "moves",
min_frequency: int = 1,
max_samples: Optional[int] = None,
) -> "ChessTokenizer":
"""
Build a tokenizer vocabulary from a Hugging Face dataset.
Note: For decomposed tokenizer, this ignores the dataset and
creates the fixed vocabulary. Provided for API compatibility.
Args:
dataset_name: Name of the dataset on Hugging Face Hub (ignored).
split: Dataset split to use (ignored).
column: Column containing move strings (ignored).
min_frequency: Minimum frequency for inclusion (ignored).
max_samples: Maximum samples to process (ignored).
Returns:
A ChessTokenizer with the fixed decomposed vocabulary.
"""
print(f"Note: Decomposed tokenizer uses fixed vocabulary (~88 tokens)")
return cls()
@property
def vocab_size(self) -> int:
return len(self._vocab)
def get_vocab(self) -> Dict[str, int]:
return dict(self._vocab)
def _parse_move(self, move: str) -> List[str]:
"""
Parse a single move into component tokens.
Args:
move: Move in extended UCI format (e.g., "WPe2e4", "BNg8f6(x+)")
Returns:
List of component tokens.
"""
match = self.MOVE_PATTERN.match(move)
if not match:
return [self.UNK_TOKEN]
tokens = []
# Color - map 'W' -> '[W]' and 'B' -> '[B]'
color = match.group(1)
tokens.append(f"[{color}]")
# Piece
tokens.append(match.group(2))
# From square
tokens.append(match.group(3))
# To square
tokens.append(match.group(4))
# Promotion (optional)
if match.group(5):
tokens.append(match.group(5))
# Parse suffixes (optional)
if match.group(6):
suffix = match.group(6)
suffix_content = suffix[1:-1]
if "x" in suffix_content:
tokens.append("x")
if "+*" in suffix_content:
tokens.append("+*")
elif "+" in suffix_content:
tokens.append("+")
if suffix_content == "o":
tokens.append("o")
elif suffix_content == "O":
tokens.append("O")
return tokens
def _tokenize(self, text: str) -> List[str]:
"""
Tokenize a string of moves into component tokens.
Args:
text: Space-separated moves in extended UCI format.
Returns:
List of component tokens.
"""
tokens = []
moves = text.strip().split()
for move in moves:
move_tokens = self._parse_move(move)
tokens.extend(move_tokens)
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 tokens back to move string.
Reconstructs moves from component tokens.
"""
special = {self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN}
result = []
current_move = []
for token in tokens:
if token in special:
if current_move:
result.append(self._reconstruct_move(current_move))
current_move = []
continue
current_move.append(token)
# Check if we have a complete move
if self._is_complete_move(current_move):
result.append(self._reconstruct_move(current_move))
current_move = []
# Handle remaining tokens
if current_move:
result.append(self._reconstruct_move(current_move))
return " ".join(result)
def _is_complete_move(self, tokens: List[str]) -> bool:
"""Check if tokens form a complete move."""
if len(tokens) < 4:
return False
# Basic move: Color + Piece + From + To
if (tokens[0] in self.COLORS and
tokens[1] in self.PIECES and
tokens[2] in self.SQUARES and
tokens[3] in self.SQUARES):
if len(tokens) == 4:
return True
# Check for modifiers
remaining = tokens[4:]
for t in remaining:
if t in self.COLORS:
return True
if t not in self.MODIFIERS and not t.startswith("="):
return True
return True
return False
def _reconstruct_move(self, tokens: List[str]) -> str:
"""Reconstruct a move string from component tokens."""
if not tokens:
return ""
if len(tokens) >= 4:
# Convert [W] -> W and [B] -> B for colors
color = tokens[0]
if color in self.COLORS:
color = color[1]
move = color + "".join(tokens[1:4])
# Add modifiers
suffixes = []
for t in tokens[4:]:
if t.startswith("="):
move += t
elif t in ["x", "+", "+*", "o", "O"]:
suffixes.append(t)
if suffixes:
move += "(" + "".join(suffixes) + ")"
return move
return "".join(tokens)
def save_vocabulary(
self,
save_directory: str,
filename_prefix: Optional[str] = None,
) -> Tuple[str]:
"""
Save the vocabulary to a file.
Args:
save_directory: Directory to save the vocabulary.
filename_prefix: Optional prefix for the vocabulary file.
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 count_vocab_from_dataset(
dataset_name: str = "dlouapre/lichess_2025-01_1M",
split: str = "train",
column: str = "moves",
max_samples: Optional[int] = None,
) -> Dict[str, int]:
"""
Count token frequencies in a dataset.
Note: For decomposed tokenizer, this counts component frequencies
rather than whole-move frequencies.
Args:
dataset_name: Name of the dataset.
split: Dataset split.
column: Column with moves.
max_samples: Max samples to process.
Returns:
Dictionary of token frequencies.
"""
from collections import Counter
from datasets import load_dataset
tokenizer = ChessTokenizer()
dataset = load_dataset(dataset_name, split=split)
if max_samples:
dataset = dataset.select(range(min(max_samples, len(dataset))))
counts = Counter()
for example in dataset:
tokens = tokenizer.tokenize(example[column])
counts.update(tokens)
return dict(counts)