Hugging Face's logo

LLM-course
/
chess-model-qtb-v2

Text Generation
Transformers
Safetensors
chess_transformer
chess
llm-course
chess-challenge
Model card Files
xet
 Community
chess-model-qtb-v2 / tokenizer.py
matheoqtb's picture
matheoqtb
Chess Challenge submission by matheoqtb
6a85e7b verified
Raw History Blame Contribute Delete
8.07 kB
"""
Custom Chess Tokenizer for the Chess Challenge.
This tokenizer treats chess moves using a 'Square-Aware' Character strategy.
Instead of full moves (e.g., WPe2e4), it splits them into meaningful atomic parts:
- Pieces/Colors: W, B, P, N, B, R, Q, K
- Full Squares: e2, e4, h8 (keeps coordinates together for geometric understanding)
- Separators: Space " "
Example: "WPe2e4" -> ["W", "P", "e2", "e4"]
"""
from __future__ import annotations
import re
import json
import os
from typing import Dict, List, Optional
from transformers import PreTrainedTokenizer
class ChessTokenizer(PreTrainedTokenizer):
"""
A custom tokenizer for chess moves using 'Square-Aware' tokenization.
It maps atomic chess components (squares like 'e4', pieces like 'P') to IDs.
This creates a small, dense vocabulary (~80 tokens) allowing deeper models.
"""
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]"
def __init__(
self,
vocab_file: Optional[str] = None,
vocab: Optional[Dict[str, int]] = None,
**kwargs,
):
# 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
# Clean kwargs to avoid conflicts
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:
# Minimal default vocab (placeholder)
self._vocab = self._create_default_vocab()
# Create reverse mapping (ID -> Token)
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]:
special_tokens = [self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN]
return {token: idx for idx, token in enumerate(special_tokens)}
@classmethod
def build_vocab_from_iterator(
cls,
iterator,
min_frequency: int = 1,
) -> "ChessTokenizer":
"""
Build vocabulary by scanning the dataset.
Splits text into pieces (W, P) and full squares (e2, e4).
"""
from collections import Counter
token_counts = Counter()
for game in iterator:
# 1. Nettoyage : on enlève les suffixes (x), (+)
game = re.sub(r'\(.*?\)', '', game)
# 2. Découpage par coups pour gérer les espaces correctement
moves = game.strip().split()
for i, move in enumerate(moves):
# Regex : Capture soit une case [a-h][1-8], soit n'importe quel autre char (.)
tokens = re.findall(r'[a-h][1-8]|.', move)
token_counts.update(tokens)
# Ajout explicite de l'espace entre les coups (sauf après le dernier)
if i < len(moves) - 1:
token_counts.update([" "])
# Filter and sort tokens
tokens = [t for t, count in token_counts.items() if count >= min_frequency]
tokens = sorted(tokens)
# Build final vocabulary dict
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 = 1, # Keep at 1 to catch all squares/pieces
max_samples: Optional[int] = 50000,
) -> "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)
@property
def vocab_size(self) -> int:
return len(self._vocab)
def get_vocab(self) -> Dict[str, int]:
return dict(self._vocab)
def _tokenize(self, text: str) -> List[str]:
"""
Tokenize input text using the Square-Aware logic.
"WPe2e4" -> ["W", "P", "e2", "e4"]
"""
# 1. Remove suffixes
text = re.sub(r'\(.*?\)', '', text)
# 2. Split into moves to manage spaces
moves = text.strip().split()
all_tokens = []
for i, move in enumerate(moves):
# Regex match: squares OR single chars
tokens = re.findall(r'[a-h][1-8]|.', move)
all_tokens.extend(tokens)
# Re-insert space token between moves
if i < len(moves) - 1:
all_tokens.append(" ")
return all_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 string.
IMPORTANT: Join with empty string "" because space " " is already a token.
"""
special = {self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN}
filtered_tokens = [t for t in tokens if t not in special]
# Join with "" because the space character is treated as a token in our vocab
return "".join(filtered_tokens)
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,)
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]:
# Utility function remains similar but should use the new regex logic if needed for analysis
# For simple counting, split() is often enough approximation, but let's be precise:
from collections import Counter
from datasets import load_dataset
import re
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:
text = re.sub(r'\(.*?\)', '', example[column])
moves = text.strip().split()
for move in moves:
tokens = re.findall(r'[a-h][1-8]|.', move)
token_counts.update(tokens)
token_counts.update([" "])
return dict(token_counts)