final-model

tokenizer.py

@@ -0,0 +1,321 @@
+"""
+Factorized UCI / verbose-UCI tokenizer for the Chess1MChallenge.
+
+Input move examples:
+- WPe2e4
+- BNg8f6
+- WQd1h5(x)+
+- WKe1g1(o)     # castling indicated by suffix in this dataset convention
+- WPe7e8=Q      # promotion styles vary; we support =Q or trailing q/r/b/n
+
+We tokenize each move into a small sequence of tokens:
+  [SIDE] [PIECE] [FROM_SQ] [TO_SQ] (optional: [PROMO_*]) (optional: [CAPTURE]) (optional: [CHECK/MATE]) (optional: [CASTLE])
+
+This keeps vocab small and compositional.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+import re
+from typing import Dict, List, Optional, Tuple, Any, Union, Sequence
+import torch
+from transformers import PreTrainedTokenizer
+
+
+class ChessTokenizer(PreTrainedTokenizer):
+    model_input_names = ["input_ids", "attention_mask"]
+    vocab_files_names = {"vocab_file": "vocab.json"}
+
+    PAD_TOKEN = "[PAD]"
+    BOS_TOKEN = "[BOS]"
+    EOS_TOKEN = "[EOS]"
+    UNK_TOKEN = "[UNK]"
+
+    # Token prefixes (purely cosmetic; still single tokens in vocab)
+    SIDE_W = "SIDE_W"
+    SIDE_B = "SIDE_B"
+
+    PIECES = ["P", "N", "B", "R", "Q", "K"]
+
+    # Suffix/flags
+    CAPTURE = "CAPTURE"   # (x)
+    CHECK = "CHECK"       # +
+    MATE = "MATE"         # ++ or (+*)
+    CASTLE = "CASTLE"     # (o) or (O)
+
+    PROMO_PREFIX = "PROMO_"  # PROMO_Q, PROMO_R, PROMO_B, PROMO_N
+
+    # Regex for your verbose convention:
+    #   <Side><Piece><from><to><optional_promo><optional_suffixes>
+    # side: W/B
+    # piece: P N B R Q K
+    # from/to: [a-h][1-8]
+    # promo: =Q or =q or trailing Q/q etc (we accept several)
+    MOVE_RE = re.compile(
+        r"^(?P<side>[WB])"
+        r"(?P<piece>[PNBRQK])"
+        r"(?P<from>[a-h][1-8])"
+        r"(?P<to>[a-h][1-8])"
+        r"(?P<rest>.*)$"
+    )
+
+    def __init__(
+        self,
+        vocab_file: Optional[str] = None,
+        vocab: Optional[Dict[str, int]] = None,
+        **kwargs,
+    ):
+        # Avoid duplicate special tokens passed by HF loaders
+        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 = {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 _build_fixed_vocab(self) -> Dict[str, int]:
+        special = [self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN]
+
+        sides = [self.SIDE_W, self.SIDE_B]
+        pieces = [f"PIECE_{p}" for p in self.PIECES]
+
+        squares = [f"SQ_{file}{rank}" for file in "abcdefgh" for rank in "12345678"]
+
+        promos = [f"{self.PROMO_PREFIX}{p}" for p in ["Q", "R", "B", "N"]]
+
+        flags = [self.CAPTURE, self.CHECK, self.MATE, self.CASTLE]
+
+        tokens = special + sides + pieces + squares + promos + flags
+        return {tok: i for i, tok in enumerate(tokens)}
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self._vocab)
+
+    def get_vocab(self) -> Dict[str, int]:
+        return dict(self._vocab)
+
+    # -------------------------
+    # Core tokenization methods
+    # -------------------------
+
+    def _tokenize(self, text: str) -> List[str]:
+        """
+        Tokenize a game string into factorized tokens.
+
+        Input is a space-separated move sequence.
+        Each move becomes multiple tokens.
+        """
+        out: List[str] = []
+        for move in text.strip().split():
+            out.extend(self._tokenize_move(move))
+        return out
+
+    def _tokenize_move(self, move: str) -> List[str]:
+        """
+        Convert one verbose-UCI move into tokens.
+        """
+        m = self.MOVE_RE.match(move)
+        if not m:
+            return [self.UNK_TOKEN]
+
+        side = m.group("side")
+        piece = m.group("piece")
+        frm = m.group("from")
+        to = m.group("to")
+        rest = m.group("rest") or ""
+
+        tokens: List[str] = []
+        tokens.append(self.SIDE_W if side == "W" else self.SIDE_B)
+        tokens.append(f"PIECE_{piece}")
+        tokens.append(f"SQ_{frm}")
+        tokens.append(f"SQ_{to}")
+
+        promo = self._parse_promotion(rest)
+        if promo is not None:
+            tokens.append(f"{self.PROMO_PREFIX}{promo}")
+
+        # Flags (order is fixed for determinism)
+        if "(x)" in rest:
+            tokens.append(self.CAPTURE)
+
+        # checkmate conventions vary; support "++" or "(+*)" (and also "#")
+        if "++" in rest or "(+*)" in rest or "#" in rest:
+            tokens.append(self.MATE)
+        elif "+" in rest:
+            tokens.append(self.CHECK)
+
+        # castling flag appears in this dataset as (o) or (O)
+        if "(o)" in rest or "(O)" in rest:
+            tokens.append(self.CASTLE)
+
+        return tokens
+
+    def _parse_promotion(self, rest: str) -> Optional[str]:
+        """
+        Detect promotion piece if present.
+        Accepts patterns like:
+          =Q, =q, e7e8Q, e7e8=q, etc.
+        Returns one of Q/R/B/N or None.
+        """
+        # Look for =<piece>
+        m = re.search(r"=([QRBNqrbn])", rest)
+        if m:
+            return m.group(1).upper()
+
+        # Or trailing promo letter (rare, but some formats do it)
+        m2 = re.search(r"([QRBNqrbn])", rest)
+        # Only treat as promo if it looks like a promo marker context (avoid grabbing random chars)
+        # Heuristic: promotion usually appears near end and not inside parentheses
+        if m2 and not "(" in rest:
+            # If rest is exactly "Q" or "q", accept
+            if rest.strip() in ["Q", "R", "B", "N", "q", "r", "b", "n"]:
+                return rest.strip().upper()
+
+        return None
+
+    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:
+        """
+        Convert factorized tokens back into a move string sequence.
+        This expects tokens to be aligned in move-chunks, so it’s mostly for debugging.
+        """
+        special = {self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN}
+        toks = [t for t in tokens if t not in special]
+
+        # Re-chunk by detecting SIDE tokens as move boundaries.
+        moves: List[str] = []
+        i = 0
+        while i < len(toks):
+            if toks[i] not in (self.SIDE_W, self.SIDE_B):
+                # If misaligned, skip until next SIDE
+                i += 1
+                continue
+            move_str, next_i = self._decode_one_move(toks, i)
+            moves.append(move_str)
+            i = next_i
+
+        return " ".join(moves)
+
+    def _decode_one_move(self, toks: List[str], i: int) -> Tuple[str, int]:
+        """
+        Decode a single move starting at index i (which should be SIDE_*).
+        Returns (move_string, next_index).
+        """
+        side_tok = toks[i]
+        side = "W" if side_tok == self.SIDE_W else "B"
+
+        # Need at least PIECE + FROM + TO
+        if i + 3 >= len(toks):
+            return "", i + 1
+
+        piece_tok = toks[i + 1]
+        from_tok = toks[i + 2]
+        to_tok = toks[i + 3]
+
+        if not piece_tok.startswith("PIECE_") or not from_tok.startswith("SQ_") or not to_tok.startswith("SQ_"):
+            return "", i + 1
+
+        piece = piece_tok.replace("PIECE_", "")
+        frm = from_tok.replace("SQ_", "")
+        to = to_tok.replace("SQ_", "")
+
+        j = i + 4
+        promo = None
+        flags: List[str] = []
+
+        # Read until next SIDE token or end
+        while j < len(toks) and toks[j] not in (self.SIDE_W, self.SIDE_B):
+            t = toks[j]
+            if t.startswith(self.PROMO_PREFIX):
+                promo = t.replace(self.PROMO_PREFIX, "")
+            elif t in (self.CAPTURE, self.CHECK, self.MATE, self.CASTLE):
+                flags.append(t)
+            j += 1
+
+        rest = ""
+        if promo is not None:
+            rest += f"={promo}"
+
+        # Match your dataset-style suffixes
+        if self.CAPTURE in flags:
+            rest += "(x)"
+        if self.MATE in flags:
+            rest += "++"   # or "(+*)" if your dataset prefers that
+        elif self.CHECK in flags:
+            rest += "+"
+        if self.CASTLE in flags:
+            rest += "(o)"
+
+        return f"{side}{piece}{frm}{to}{rest}", j
+
+    # ---------------
+    # Saving/loading
+    # ---------------
+
+    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 decode(
+        self,
+        token_ids: Union[int, Sequence[int], torch.Tensor],
+        skip_special_tokens: bool = False,
+        clean_up_tokenization_spaces: bool = False,
+        **kwargs: Any,
+    ) -> str:
+        # Normalize input type
+        if isinstance(token_ids, int):
+            ids = [token_ids]
+        elif "torch" in str(type(token_ids)):
+            # torch.Tensor
+            ids = token_ids.detach().cpu().flatten().tolist()
+        else:
+            ids = list(token_ids)
+
+        # Convert ids -> token strings
+        toks = [self._convert_id_to_token(i) for i in ids]
+
+        if skip_special_tokens:
+            special = {self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN}
+            toks = [t for t in toks if t not in special]
+
+        # IMPORTANT: reconstruct chess moves instead of joining token names
+        return self.convert_tokens_to_string(toks)

tokenizer_config.json

@@ -33,12 +33,6 @@
       "special": true
     }
   },
-  "auto_map": {
-    "AutoTokenizer": [
-      "tokenizer.py",
-      "ChessTokenizer"
-    ]
-  },
   "bos_token": "[BOS]",
   "clean_up_tokenization_spaces": false,
   "eos_token": "[EOS]",
@@ -46,5 +40,8 @@
   "model_max_length": 1000000000000000019884624838656,
   "pad_token": "[PAD]",
   "tokenizer_class": "ChessTokenizer",
+  "auto_map": {
+    "AutoTokenizer": ["tokenizer.py", "ChessTokenizer"]
+  },
   "unk_token": "[UNK]"
 }

training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ae24236124c061824973c9e4f4340e7186cd3564179ff2db75e0a72e99907bbe
+size 5777