Update tokenizer.py

tokenizer.py

@@ -1,13 +1,14 @@
 from __future__ import annotations
 import json
 import os
+import re
 from typing import Dict, List, Optional
 from transformers import PreTrainedTokenizer
 import torch
 
 class ChessTokenizer(PreTrainedTokenizer):
     """
-    vocab size: 149 (4 special + 12 pieces + 64 from_sq + 64 to_sq + 5 suffix)
+    vocab size: 4 special + 12 pieces + 64 from_sq + 64 to_sq + 5 suffix + 4 promotions (qrbn)
     """
 
     model_input_names = ["input_ids", "attention_mask"]
@@ -24,6 +25,7 @@ class ChessTokenizer(PreTrainedTokenizer):
         self.colors_pieces = [f'{c}{p}' for c in ['W','B'] for p in ['P','N','B','R','Q','K']]
         self.squares = [f'{f}{r}' for r in '12345678' for f in 'abcdefgh']
         self.suffixes = ["(x)", "(+)", "(+*)", "(o)", "(O)"]
+        self.promotions = ["q", "r", "b", "n"]  # promotion Token
 
         if vocab is not None:
             self._vocab = vocab
@@ -32,14 +34,11 @@ class ChessTokenizer(PreTrainedTokenizer):
                 self._vocab = json.load(f)
         else:
             self._vocab = {t: i for i, t in enumerate(special_tokens)}
-            for cp in self.colors_pieces:
-                self._vocab[cp] = len(self._vocab)
-            for sq in self.squares:
-                self._vocab[f"{sq}_f"] = len(self._vocab)
-            for sq in self.squares:
-                self._vocab[f"{sq}_t"] = len(self._vocab)
-            for suf in self.suffixes:
-                self._vocab[suf] = len(self._vocab)
+            for cp in self.colors_pieces: self._vocab[cp] = len(self._vocab)
+            for sq in self.squares: self._vocab[f"{sq}_f"] = len(self._vocab)
+            for sq in self.squares: self._vocab[f"{sq}_t"] = len(self._vocab)
+            for suf in self.suffixes: self._vocab[suf] = len(self._vocab)
+            for promo in self.promotions: self._vocab[promo] = len(self._vocab)
 
         self._ids_to_tokens = {v: k for k, v in self._vocab.items()}
 
@@ -54,8 +53,9 @@ class ChessTokenizer(PreTrainedTokenizer):
     @property
     def vocab_size(self) -> int:
         return len(self._vocab)
-
+    
     def get_vocab(self) -> Dict[str, int]:
+        """Return the vocabulary as a dictionary."""
         return dict(self._vocab)
 
     def _tokenize(self, text: str) -> List[str]:
@@ -64,43 +64,55 @@ class ChessTokenizer(PreTrainedTokenizer):
         for part in parts:
             if part in self._vocab:
                 tokens.append(part)
-            elif len(part) >= 6:
-                piece, f_sq, t_sq = part[:2], part[2:4] + "_f", part[4:6] + "_t"
+                continue
+            
+            # Deal with WPe7e8q
+            if len(part) >= 6:
+                piece = part[:2]
+                f_sq = part[2:4] + "_f"
+                t_sq = part[4:6] + "_t"
                 if piece in self._vocab: tokens.append(piece)
                 if f_sq in self._vocab: tokens.append(f_sq)
                 if t_sq in self._vocab: tokens.append(t_sq)
-                if len(part) > 6 and part[6:] in self.suffixes:
-                    tokens.append(part[6:])
+                
+                # Check if rest part include promotion or suffix
+                rest = part[6:]
+                if not rest: continue
+                
+                # Extract promotion letter (q, r, b, n)
+                promo_match = re.search(r'[qrbnQRBN]', rest)
+                if promo_match:
+                    p_char = promo_match.group(0).lower()
+                    if p_char in self._vocab: tokens.append(p_char)
+                
+                # Extract suffixes
+                for suf in self.suffixes:
+                    if suf in rest:
+                        tokens.append(suf)
         return tokens
 
     def _convert_id_to_token(self, index: int) -> str:
         token = self._ids_to_tokens.get(index, self.UNK_TOKEN)
         if token in [self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN]:
             return ""
+        # Same if promotion or suffix, delete _f or _t
+        if token in self.promotions or token in self.suffixes:
+            return token
         return token.replace("_f", "").replace("_t", "")
 
     def convert_tokens_to_string(self, tokens: List[str]) -> str:
         res = []
         for t in tokens:
             if not t: continue
-            # if piece token，new move，add space
+            # If piece (WP), new move, add space
             if len(t) == 2 and (t.startswith('W') or t.startswith('B')):
                 res.append(" " + t)
             else:
                 res.append(t)
         return "".join(res).strip()
+
     def _convert_token_to_id(self, token: str) -> int:
         return self._vocab.get(token, self._vocab.get(self.UNK_TOKEN))
-    def _convert_id_to_token(self, index: int) -> str:
-        token = self._ids_to_tokens.get(index, self.UNK_TOKEN)
-        if token in [self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN]:
-            return ""
-        if token in self.suffixes:
-            return token
-        return token.replace("_f", "").replace("_t", "")
-
-    def convert_tokens_to_string(self, tokens: List[str]) -> str:
-        return "".join([t for t in tokens if t])
 
     def decode(self, token_ids, skip_special_tokens=True, **kwargs) -> str:
         if hasattr(token_ids, "tolist"):