Chess Challenge submission by stephecw

README.md

@@ -14,7 +14,7 @@ Chess model submitted to the LLM Course Chess Challenge.
 ## Submission Info
 
 - **Submitted by**: [stephecw](https://huggingface.co/stephecw)
-- **Parameters**: 999,030
+- **Parameters**: 997,488
 - **Organization**: LLM-course
 
 ## Model Details
@@ -23,4 +23,4 @@ Chess model submitted to the LLM Course Chess Challenge.
 - **Vocab size**: 72
 - **Embedding dim**: 128
 - **Layers**: 6
-- **Heads**: 4
+- **Heads**: 8

config.json

@@ -10,8 +10,8 @@
   "model_type": "chess_transformer",
   "n_ctx": 256,
   "n_embd": 128,
-  "n_head": 4,
-  "n_inner": 361,
+  "n_head": 8,
+  "n_inner": 360,
   "n_layer": 6,
   "pad_token_id": 0,
   "tie_weights": true,

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0b43464978d01cb02981333cddb11f7a58ec88c9f2def41f49dfcf08f0d01b32
-size 4002568
+oid sha256:a686c5ae98a7ec29e009b43931e1e936feabe771ce50cc2ed5bb5db36196e10e
+size 3996400

src/tokenizer.py

@@ -10,14 +10,23 @@ The dataset format uses:
 - Source and destination squares (e.g., e2e4)
 - Special suffixes: (x)=capture, (+)=check, (+*)=checkmate, (o)/(O)=castling
 """
+
 from __future__ import annotations
 
 import json
 import os
+import re
 from typing import Dict, List, Optional
 
 from transformers import PreTrainedTokenizer
 
+SQUARE_RE = re.compile(r"[a-h][1-8]")
+UCI_PROMO_RE = re.compile(r"^[a-h][1-8][a-h][1-8]([qrbn])$", re.IGNORECASE)
+EQ_PROMO_RE = re.compile(r"=([QRBNqrbn])")
+PAREN_PROMO_RE = re.compile(r"\(([QRBNqrbn])\)")
+
+PROMOS = {"q", "r", "b", "n"}
+
 
 class ChessTokenizer(PreTrainedTokenizer):
     vocab_files_names = {"vocab_file": "vocab.json"}
@@ -34,19 +43,16 @@ class ChessTokenizer(PreTrainedTokenizer):
         vocab: Optional[Dict[str, int]] = None,
         **kwargs,
     ):
-        # Define 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
 
-        # Avoid duplicates when loading from disk
         kwargs.pop("pad_token", None)
         kwargs.pop("bos_token", None)
         kwargs.pop("eos_token", None)
         kwargs.pop("unk_token", None)
 
-        # Load vocab or create fixed vocab
         if vocab is not None:
             self._vocab = vocab
         elif vocab_file is not None and os.path.exists(vocab_file):
@@ -67,13 +73,9 @@ class ChessTokenizer(PreTrainedTokenizer):
 
     def _create_fixed_vocab(self) -> Dict[str, int]:
         specials = [self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN]
-
-        # a1..h8 (rank first in string is conventional, but we just need consistent list)
-        squares = [f"{file}{rank}" for rank in "12345678" for file in "abcdefgh"]
-
-        # Optional promotion tokens (evaluator can detect q/r/b/n after the 2nd square)
+        # IMPORTANT: deterministic ids matching a1,a2,...,a8,b1,... style
+        squares = [f"{f}{r}" for f in "abcdefgh" for r in "12345678"]
         promos = ["q", "r", "b", "n"]
-
         tokens = specials + squares + promos
         return {tok: i for i, tok in enumerate(tokens)}
 
@@ -84,8 +86,56 @@ class ChessTokenizer(PreTrainedTokenizer):
     def get_vocab(self) -> Dict[str, int]:
         return dict(self._vocab)
 
+    def _extract_promo_anywhere(self, mv: str) -> Optional[str]:
+        m = EQ_PROMO_RE.search(mv)
+        if m:
+            return m.group(1).lower()
+        m = PAREN_PROMO_RE.search(mv)
+        if m:
+            return m.group(1).lower()
+        m = UCI_PROMO_RE.match(mv)
+        if m:
+            return m.group(1).lower()
+        return None
+
     def _tokenize(self, text: str) -> List[str]:
-        return text.strip().split()
+        """
+        Robust tokenization:
+        - keeps special tokens ([BOS], etc.) as-is (HF handles them)
+        - accepts already-split squares: "e2 e4"
+        - accepts uci concat: "e2e4" -> e2,e4 (+promo)
+        - accepts verbose tokens containing squares: "WPe2e4(x+)" -> e2,e4 (+promo)
+        """
+        tokens: List[str] = []
+
+        for chunk in text.strip().split():
+            # already-split square?
+            if re.fullmatch(r"[a-h][1-8]", chunk):
+                tokens.append(chunk)
+                continue
+
+            # promo alone?
+            if chunk in PROMOS:
+                tokens.append(chunk)
+                continue
+
+            # otherwise: extract squares from inside
+            squares = SQUARE_RE.findall(chunk)
+            if len(squares) >= 2:
+                tokens.append(squares[0])
+                tokens.append(squares[1])
+
+                promo = self._extract_promo_anywhere(chunk)
+                if promo in PROMOS:
+                    tokens.append(promo)
+            else:
+                # allow special tokens to pass through if present
+                if chunk in {self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN}:
+                    tokens.append(chunk)
+                else:
+                    tokens.append(self.UNK_TOKEN)
+
+        return tokens
 
     def _convert_token_to_id(self, token: str) -> int:
         return self._vocab.get(token, self._vocab[self.UNK_TOKEN])
@@ -94,14 +144,29 @@ class ChessTokenizer(PreTrainedTokenizer):
         return self._ids_to_tokens.get(index, self.UNK_TOKEN)
 
     def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        """
+        Reconstruct "e2e4 e7e8q ..."
+        """
         special = {self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN}
-        return " ".join(t for t in tokens if t not in special)
-
-    def save_vocabulary(
-        self,
-        save_directory: str,
-        filename_prefix: Optional[str] = None,
-    ) -> tuple:
+        clean = [t for t in tokens if t not in special]
+
+        moves: List[str] = []
+        i = 0
+        while i < len(clean):
+            if re.fullmatch(r"[a-h][1-8]", clean[i]) and i + 1 < len(clean) and re.fullmatch(r"[a-h][1-8]", clean[i + 1]):
+                mv = clean[i] + clean[i + 1]
+                i += 2
+                if i < len(clean) and clean[i] in PROMOS:
+                    mv += clean[i]
+                    i += 1
+                moves.append(mv)
+            else:
+                moves.append(clean[i])
+                i += 1
+
+        return " ".join(moves)
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> tuple:
         os.makedirs(save_directory, exist_ok=True)
         vocab_file = os.path.join(
             save_directory,
@@ -110,6 +175,7 @@ class ChessTokenizer(PreTrainedTokenizer):
         with open(vocab_file, "w", encoding="utf-8") as f:
             json.dump(self._vocab, f, ensure_ascii=False, indent=2)
         return (vocab_file,)
-    
+
+
 from transformers import AutoTokenizer
 ChessTokenizer.register_for_auto_class("AutoTokenizer")

tokenizer.py

@@ -10,14 +10,23 @@ The dataset format uses:
 - Source and destination squares (e.g., e2e4)
 - Special suffixes: (x)=capture, (+)=check, (+*)=checkmate, (o)/(O)=castling
 """
+
 from __future__ import annotations
 
 import json
 import os
+import re
 from typing import Dict, List, Optional
 
 from transformers import PreTrainedTokenizer
 
+SQUARE_RE = re.compile(r"[a-h][1-8]")
+UCI_PROMO_RE = re.compile(r"^[a-h][1-8][a-h][1-8]([qrbn])$", re.IGNORECASE)
+EQ_PROMO_RE = re.compile(r"=([QRBNqrbn])")
+PAREN_PROMO_RE = re.compile(r"\(([QRBNqrbn])\)")
+
+PROMOS = {"q", "r", "b", "n"}
+
 
 class ChessTokenizer(PreTrainedTokenizer):
     vocab_files_names = {"vocab_file": "vocab.json"}
@@ -34,19 +43,16 @@ class ChessTokenizer(PreTrainedTokenizer):
         vocab: Optional[Dict[str, int]] = None,
         **kwargs,
     ):
-        # Define 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
 
-        # Avoid duplicates when loading from disk
         kwargs.pop("pad_token", None)
         kwargs.pop("bos_token", None)
         kwargs.pop("eos_token", None)
         kwargs.pop("unk_token", None)
 
-        # Load vocab or create fixed vocab
         if vocab is not None:
             self._vocab = vocab
         elif vocab_file is not None and os.path.exists(vocab_file):
@@ -67,13 +73,9 @@ class ChessTokenizer(PreTrainedTokenizer):
 
     def _create_fixed_vocab(self) -> Dict[str, int]:
         specials = [self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN]
-
-        # a1..h8 (rank first in string is conventional, but we just need consistent list)
-        squares = [f"{file}{rank}" for rank in "12345678" for file in "abcdefgh"]
-
-        # Optional promotion tokens (evaluator can detect q/r/b/n after the 2nd square)
+        # IMPORTANT: deterministic ids matching a1,a2,...,a8,b1,... style
+        squares = [f"{f}{r}" for f in "abcdefgh" for r in "12345678"]
         promos = ["q", "r", "b", "n"]
-
         tokens = specials + squares + promos
         return {tok: i for i, tok in enumerate(tokens)}
 
@@ -84,8 +86,56 @@ class ChessTokenizer(PreTrainedTokenizer):
     def get_vocab(self) -> Dict[str, int]:
         return dict(self._vocab)
 
+    def _extract_promo_anywhere(self, mv: str) -> Optional[str]:
+        m = EQ_PROMO_RE.search(mv)
+        if m:
+            return m.group(1).lower()
+        m = PAREN_PROMO_RE.search(mv)
+        if m:
+            return m.group(1).lower()
+        m = UCI_PROMO_RE.match(mv)
+        if m:
+            return m.group(1).lower()
+        return None
+
     def _tokenize(self, text: str) -> List[str]:
-        return text.strip().split()
+        """
+        Robust tokenization:
+        - keeps special tokens ([BOS], etc.) as-is (HF handles them)
+        - accepts already-split squares: "e2 e4"
+        - accepts uci concat: "e2e4" -> e2,e4 (+promo)
+        - accepts verbose tokens containing squares: "WPe2e4(x+)" -> e2,e4 (+promo)
+        """
+        tokens: List[str] = []
+
+        for chunk in text.strip().split():
+            # already-split square?
+            if re.fullmatch(r"[a-h][1-8]", chunk):
+                tokens.append(chunk)
+                continue
+
+            # promo alone?
+            if chunk in PROMOS:
+                tokens.append(chunk)
+                continue
+
+            # otherwise: extract squares from inside
+            squares = SQUARE_RE.findall(chunk)
+            if len(squares) >= 2:
+                tokens.append(squares[0])
+                tokens.append(squares[1])
+
+                promo = self._extract_promo_anywhere(chunk)
+                if promo in PROMOS:
+                    tokens.append(promo)
+            else:
+                # allow special tokens to pass through if present
+                if chunk in {self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN}:
+                    tokens.append(chunk)
+                else:
+                    tokens.append(self.UNK_TOKEN)
+
+        return tokens
 
     def _convert_token_to_id(self, token: str) -> int:
         return self._vocab.get(token, self._vocab[self.UNK_TOKEN])
@@ -94,14 +144,29 @@ class ChessTokenizer(PreTrainedTokenizer):
         return self._ids_to_tokens.get(index, self.UNK_TOKEN)
 
     def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        """
+        Reconstruct "e2e4 e7e8q ..."
+        """
         special = {self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN}
-        return " ".join(t for t in tokens if t not in special)
-
-    def save_vocabulary(
-        self,
-        save_directory: str,
-        filename_prefix: Optional[str] = None,
-    ) -> tuple:
+        clean = [t for t in tokens if t not in special]
+
+        moves: List[str] = []
+        i = 0
+        while i < len(clean):
+            if re.fullmatch(r"[a-h][1-8]", clean[i]) and i + 1 < len(clean) and re.fullmatch(r"[a-h][1-8]", clean[i + 1]):
+                mv = clean[i] + clean[i + 1]
+                i += 2
+                if i < len(clean) and clean[i] in PROMOS:
+                    mv += clean[i]
+                    i += 1
+                moves.append(mv)
+            else:
+                moves.append(clean[i])
+                i += 1
+
+        return " ".join(moves)
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> tuple:
         os.makedirs(save_directory, exist_ok=True)
         vocab_file = os.path.join(
             save_directory,
@@ -110,6 +175,7 @@ class ChessTokenizer(PreTrainedTokenizer):
         with open(vocab_file, "w", encoding="utf-8") as f:
             json.dump(self._vocab, f, ensure_ascii=False, indent=2)
         return (vocab_file,)
-    
+
+
 from transformers import AutoTokenizer
 ChessTokenizer.register_for_auto_class("AutoTokenizer")

vocab.json

@@ -4,68 +4,68 @@
   "[EOS]": 2,
   "[UNK]": 3,
   "a1": 4,
-  "b1": 5,
-  "c1": 6,
-  "d1": 7,
-  "e1": 8,
-  "f1": 9,
-  "g1": 10,
-  "h1": 11,
-  "a2": 12,
+  "a2": 5,
+  "a3": 6,
+  "a4": 7,
+  "a5": 8,
+  "a6": 9,
+  "a7": 10,
+  "a8": 11,
+  "b1": 12,
   "b2": 13,
-  "c2": 14,
-  "d2": 15,
-  "e2": 16,
-  "f2": 17,
-  "g2": 18,
-  "h2": 19,
-  "a3": 20,
-  "b3": 21,
+  "b3": 14,
+  "b4": 15,
+  "b5": 16,
+  "b6": 17,
+  "b7": 18,
+  "b8": 19,
+  "c1": 20,
+  "c2": 21,
   "c3": 22,
-  "d3": 23,
-  "e3": 24,
-  "f3": 25,
-  "g3": 26,
-  "h3": 27,
-  "a4": 28,
-  "b4": 29,
-  "c4": 30,
+  "c4": 23,
+  "c5": 24,
+  "c6": 25,
+  "c7": 26,
+  "c8": 27,
+  "d1": 28,
+  "d2": 29,
+  "d3": 30,
   "d4": 31,
-  "e4": 32,
-  "f4": 33,
-  "g4": 34,
-  "h4": 35,
-  "a5": 36,
-  "b5": 37,
-  "c5": 38,
-  "d5": 39,
+  "d5": 32,
+  "d6": 33,
+  "d7": 34,
+  "d8": 35,
+  "e1": 36,
+  "e2": 37,
+  "e3": 38,
+  "e4": 39,
   "e5": 40,
-  "f5": 41,
-  "g5": 42,
-  "h5": 43,
-  "a6": 44,
-  "b6": 45,
-  "c6": 46,
-  "d6": 47,
-  "e6": 48,
+  "e6": 41,
+  "e7": 42,
+  "e8": 43,
+  "f1": 44,
+  "f2": 45,
+  "f3": 46,
+  "f4": 47,
+  "f5": 48,
   "f6": 49,
-  "g6": 50,
-  "h6": 51,
-  "a7": 52,
-  "b7": 53,
-  "c7": 54,
-  "d7": 55,
-  "e7": 56,
-  "f7": 57,
+  "f7": 50,
+  "f8": 51,
+  "g1": 52,
+  "g2": 53,
+  "g3": 54,
+  "g4": 55,
+  "g5": 56,
+  "g6": 57,
   "g7": 58,
-  "h7": 59,
-  "a8": 60,
-  "b8": 61,
-  "c8": 62,
-  "d8": 63,
-  "e8": 64,
-  "f8": 65,
-  "g8": 66,
+  "g8": 59,
+  "h1": 60,
+  "h2": 61,
+  "h3": 62,
+  "h4": 63,
+  "h5": 64,
+  "h6": 65,
+  "h7": 66,
   "h8": 67,
   "q": 68,
   "r": 69,