|
|
|
|
|
|
|
|
""" |
|
|
Custom Chess Tokenizer - Decomposed Move Representation. |
|
|
|
|
|
Decomposes chess moves into components: Color, Piece, From square, To square, Modifiers. |
|
|
Uses a fixed vocabulary of ~88 tokens. |
|
|
""" |
|
|
|
|
|
from __future__ import annotations |
|
|
|
|
|
import json |
|
|
import os |
|
|
import re |
|
|
from typing import Dict, List, Optional |
|
|
|
|
|
from transformers import PreTrainedTokenizer |
|
|
|
|
|
|
|
|
class ChessTokenizer(PreTrainedTokenizer): |
|
|
model_input_names = ["input_ids", "attention_mask"] |
|
|
|
|
|
PAD_TOKEN = "[PAD]" |
|
|
BOS_TOKEN = "[BOS]" |
|
|
EOS_TOKEN = "[EOS]" |
|
|
UNK_TOKEN = "[UNK]" |
|
|
|
|
|
COLORS = ["[W]", "[B]"] |
|
|
PIECES = ["P", "N", "B", "R", "Q", "K"] |
|
|
SQUARES = [f + r for f in "abcdefgh" for r in "12345678"] |
|
|
MODIFIERS = ["x", "+", "#", "+*", "=Q", "=R", "=B", "=N", "O-O", "O-O-O", "o", "O"] |
|
|
|
|
|
MOVE_PATTERN = re.compile( |
|
|
r'^([WB])([PNBRQK])([a-h][1-8])([a-h][1-8])(=[QRBN])?(\([xoO+*]+\))?$' |
|
|
) |
|
|
|
|
|
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 |
|
|
|
|
|
kwargs.pop("pad_token", None) |
|
|
kwargs.pop("bos_token", None) |
|
|
kwargs.pop("eos_token", None) |
|
|
kwargs.pop("unk_token", None) |
|
|
|
|
|
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() |
|
|
|
|
|
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]: |
|
|
tokens = [self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN] |
|
|
tokens.extend(self.COLORS) |
|
|
tokens.extend(self.PIECES) |
|
|
tokens.extend(self.SQUARES) |
|
|
tokens.extend(self.MODIFIERS) |
|
|
return {token: idx for idx, token in enumerate(tokens)} |
|
|
|
|
|
def _parse_move(self, move: str) -> List[str]: |
|
|
"""Parse a single move into component tokens.""" |
|
|
match = self.MOVE_PATTERN.match(move) |
|
|
if not match: |
|
|
return [self.UNK_TOKEN] |
|
|
|
|
|
tokens = [f"[{match.group(1)}]", match.group(2), match.group(3), match.group(4)] |
|
|
|
|
|
if match.group(5): |
|
|
tokens.append(match.group(5)) |
|
|
|
|
|
if match.group(6): |
|
|
suffix_content = match.group(6)[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.""" |
|
|
tokens = [] |
|
|
for move in text.strip().split(): |
|
|
tokens.extend(self._parse_move(move)) |
|
|
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: |
|
|
"""Reconstruct 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) |
|
|
if self._is_complete_move(current_move): |
|
|
result.append(self._reconstruct_move(current_move)) |
|
|
current_move = [] |
|
|
|
|
|
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 |
|
|
|
|
|
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 |
|
|
for t in tokens[4:]: |
|
|
if t in self.COLORS or (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 or len(tokens) < 4: |
|
|
return "".join(tokens) |
|
|
|
|
|
color = tokens[0][1] if tokens[0] in self.COLORS else tokens[0] |
|
|
move = color + "".join(tokens[1:4]) |
|
|
|
|
|
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 |
|
|
|
|
|
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=1): |
|
|
return cls() |
|
|
|
|
|
@classmethod |
|
|
def build_vocab_from_dataset(cls, **kwargs): |
|
|
return cls() |
|
|
|
|
|
@property |
|
|
def vocab_size(self) -> int: |
|
|
return len(self._vocab) |
|
|
|
|
|
def get_vocab(self) -> Dict[str, int]: |
|
|
return dict(self._vocab) |
|
|
|
