Update tokenizer.py

tokenizer.py

@@ -0,0 +1,215 @@
+# src/tokenizer.py
+from __future__ import annotations
+
+import json
+import os
+from typing import Dict, List, Optional, Tuple
+
+from transformers import PreTrainedTokenizer
+
+# --- Fixed vocab pieces ---
+_SQUARES = [f"{file}{rank}" for rank in "12345678" for file in "abcdefgh"]
+_PROMOS = ["=Q", "=R", "=B", "=N"]
+
+
+class SquaresOnlyChessTokenizer(PreTrainedTokenizer):
+    """
+    Tokenizer designed to MINIMIZE illegal-move formatting issues under the provided evaluate.py,
+    WITHOUT modifying evaluate.py.
+
+    Key idea:
+      - evaluate.py extracts UCI using move_token[2:4] + move_token[4:6]
+      - so decoded move strings must look like: "W" + <any char> + from_sq + to_sq [+ "=Q/R/B/N"]
+        e.g. "WPe2e4", "WNg8f6", "WPe7e8=Q"
+      - evaluate.py stops generation on whitespace; we therefore include a SPACE token as a move separator.
+
+    Encoding (per move):
+      from_sq, to_sq, promo? , " "   (space is a separator token)
+
+    Decoding (per move):
+      "WP" + from_sq + to_sq + promo?   (constant prefix)
+
+    We strip all suffixes like (x), (+), (+*), (o)/(O) since evaluator doesn't use them.
+    """
+
+    vocab_files_names = {"vocab_file": "vocab.json"}
+    model_input_names = ["input_ids", "attention_mask"]
+
+    PAD_TOKEN = "[PAD]"
+    BOS_TOKEN = "[BOS]"
+    EOS_TOKEN = "[EOS]"
+    UNK_TOKEN = "[UNK]"
+
+    MOVE_SEP = " "  # IMPORTANT: whitespace => evaluator stops on separator
+
+    def __init__(
+        self,
+        vocab: Optional[Dict[str, int]] = None,
+        vocab_file: Optional[str] = None,
+        **kwargs,
+    ):
+        # Avoid duplicates when loading/saving
+        kwargs.pop("pad_token", None)
+        kwargs.pop("bos_token", None)
+        kwargs.pop("eos_token", None)
+        kwargs.pop("unk_token", None)
+
+        self._pad_token = self.PAD_TOKEN
+        self._bos_token = self.BOS_TOKEN
+        self._eos_token = self.EOS_TOKEN
+        self._unk_token = self.UNK_TOKEN
+
+        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._build_fixed_vocab()
+
+        self._ids_to_tokens = {i: t for t, i 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,
+        )
+
+    # -------------------------
+    # Vocab
+    # -------------------------
+    @classmethod
+    def _build_fixed_vocab(cls) -> Dict[str, int]:
+        toks = [cls.PAD_TOKEN, cls.BOS_TOKEN, cls.EOS_TOKEN, cls.UNK_TOKEN]
+        toks += [cls.MOVE_SEP]
+        toks += _SQUARES
+        toks += _PROMOS
+        return {t: i for i, t in enumerate(toks)}
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self._vocab)
+
+    def get_vocab(self) -> Dict[str, int]:
+        return dict(self._vocab)
+
+    # -------------------------
+    # Helpers: parse / normalize
+    # -------------------------
+    @staticmethod
+    def _strip_suffixes(token: str) -> str:
+        # Remove "(x)" "(+)" "(+*)" "(o)" "(O)" etc.
+        return token.split("(", 1)[0]
+
+    @staticmethod
+    def _extract_squares_and_promo(base: str) -> Tuple[Optional[str], Optional[str], Optional[str]]:
+        """
+        base expected like:
+          WPe2e4
+          BNg8f6
+          WPe7e8=Q
+        Return: (from_sq, to_sq, promo_token like '=Q' or None)
+        """
+        if len(base) < 6:
+            return None, None, None
+        from_sq = base[2:4].lower()
+        to_sq = base[4:6].lower()
+        if from_sq not in _SQUARES or to_sq not in _SQUARES:
+            return None, None, None
+
+        promo = None
+        if "=" in base:
+            promo = base[base.index("="):].upper()  # "=Q"
+            if promo not in _PROMOS:
+                promo = None
+        return from_sq, to_sq, promo
+
+    # -------------------------
+    # Tokenization API
+    # -------------------------
+    def _tokenize(self, text: str) -> List[str]:
+        """
+        Tokenize a string of moves (space-separated).
+        Special tokens are preserved if present.
+        Each move becomes: from, to, promo?, " "
+        """
+        raw = text.strip().split()
+        out: List[str] = []
+
+        for tok in raw:
+            if tok in (self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN):
+                out.append(tok)
+                continue
+
+            base = self._strip_suffixes(tok)
+            from_sq, to_sq, promo = self._extract_squares_and_promo(base)
+
+            if from_sq is None or to_sq is None:
+                out.append(self.UNK_TOKEN)
+                out.append(self.MOVE_SEP)
+                continue
+
+            out.append(from_sq)
+            out.append(to_sq)
+            if promo is not None:
+                out.append(promo)
+            out.append(self.MOVE_SEP)
+
+        return out
+
+    def _convert_token_to_id(self, token: str) -> int:
+        return self._vocab.get(token, self._vocab[self.UNK_TOKEN])
+
+    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 a text compatible with evaluate.py.
+        Each move is rendered as: "WP" + from + to + promo?
+        Moves are separated by actual spaces (MOVE_SEP token).
+        """
+        s: List[str] = []
+        at_move_start = True
+
+        for tok in tokens:
+            if tok in (self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN):
+                continue
+
+            if tok == self.MOVE_SEP:
+                s.append(" ")
+                at_move_start = True
+                continue
+
+            if tok in _PROMOS:
+                s.append(tok)
+                continue
+
+            if tok in _SQUARES:
+                if at_move_start:
+                    s.append("WP")  # constant prefix, starts with 'W'
+                    at_move_start = False
+                s.append(tok)
+                continue
+
+            # Fallback (should be rare)
+            if at_move_start:
+                s.append("WP")
+                at_move_start = False
+            s.append(tok)
+
+        return "".join(s)
+
+    # -------------------------
+    # Saving / loading
+    # -------------------------
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        os.makedirs(save_directory, exist_ok=True)
+        path = os.path.join(save_directory, (filename_prefix + "-" if filename_prefix else "") + "vocab.json")
+        with open(path, "w", encoding="utf-8") as f:
+            json.dump(self._vocab, f, ensure_ascii=False, indent=2)
+        return (path,)
+
+ChessTokenizer = SquaresOnlyChessTokenizer