Chess Challenge submission by MaximeMuhlethaler

README.md

@@ -12,15 +12,12 @@ license: mit
 Chess model submitted to the LLM Course Chess Challenge.
 
 ## Submission Info
-
 - **Submitted by**: [MaximeMuhlethaler](https://huggingface.co/MaximeMuhlethaler)
 - **Parameters**: 980,000
 - **Organization**: LLM-course
+- **Architecture**: Custom Chess Transformer (Regex Tokenizer + EOS Protection)
 
 ## Model Details
-
-- **Architecture**: Chess Transformer (GPT-style)
 - **Vocab size**: 1700
-- **Embedding dim**: 112
 - **Layers**: 6
 - **Heads**: 8

config.json

@@ -1,7 +1,4 @@
 {
-  "architectures": [
-    "ChessForCausalLM"
-  ],
   "bos_token_id": 1,
   "dropout": 0.1,
   "dtype": "float32",
@@ -16,5 +13,10 @@
   "pad_token_id": 0,
   "tie_weights": true,
   "transformers_version": "4.57.5",
-  "vocab_size": 1700
-}
+  "unk_token_id": 3,
+  "vocab_size": 1700,
+  "auto_map": {
+    "AutoModelForCausalLM": "model.ChessForCausalLM",
+    "AutoConfig": "model.ChessConfig"
+  }
+}

model.py

@@ -0,0 +1,150 @@
+"""
+Chess Transformer Model - The "Nuclear Patch" Edition
+"""
+from __future__ import annotations
+import math
+from typing import Optional, Tuple, Union
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+class ChessConfig(PretrainedConfig):
+    model_type = "chess_transformer"
+    
+    def __init__(
+        self,
+        vocab_size=1200,
+        n_embd=128,
+        n_layer=6,
+        n_head=4,
+        n_ctx=256,
+        n_inner=None,
+        dropout=0.1,
+        layer_norm_epsilon=1e-5,
+        tie_weights=True,
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=2,
+        unk_token_id=3,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_ctx = n_ctx
+        self.n_inner = n_inner if n_inner is not None else 3 * n_embd
+        self.dropout = dropout
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.tie_weights = tie_weights
+        
+        # On passe les IDs vitaux à kwargs pour le parent
+        kwargs["pad_token_id"] = pad_token_id
+        kwargs["bos_token_id"] = bos_token_id
+        kwargs["eos_token_id"] = eos_token_id
+        kwargs["unk_token_id"] = unk_token_id
+        
+        super().__init__(**kwargs)
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, config: ChessConfig):
+        super().__init__()
+        self.n_head = config.n_head
+        self.n_embd = config.n_embd
+        self.head_dim = config.n_embd // config.n_head
+        self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
+        self.c_proj = nn.Linear(config.n_embd, config.n_embd)
+        self.dropout = nn.Dropout(config.dropout)
+        self.register_buffer("bias", torch.tril(torch.ones(config.n_ctx, config.n_ctx)).view(1, 1, config.n_ctx, config.n_ctx), persistent=False)
+    
+    def forward(self, x, attention_mask=None):
+        B, T, C = x.size()
+        qkv = self.c_attn(x)
+        q, k, v = qkv.split(self.n_embd, dim=2)
+        q = q.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+        k = k.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+        v = v.view(B, T, self.n_head, self.head_dim).transpose(1, 2)
+        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
+        att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
+        if attention_mask is not None:
+            att = att.masked_fill(attention_mask.view(B, 1, 1, T) == 0, float('-inf'))
+        att = F.softmax(att, dim=-1)
+        att = self.dropout(att)
+        y = att @ v
+        y = y.transpose(1, 2).contiguous().view(B, T, C)
+        return self.c_proj(y)
+
+class FeedForward(nn.Module):
+    def __init__(self, config: ChessConfig):
+        super().__init__()
+        self.c_fc = nn.Linear(config.n_embd, config.n_inner)
+        self.c_proj = nn.Linear(config.n_inner, config.n_embd)
+        self.dropout = nn.Dropout(config.dropout)
+    def forward(self, x):
+        return self.dropout(self.c_proj(F.gelu(self.c_fc(x))))
+
+class TransformerBlock(nn.Module):
+    def __init__(self, config: ChessConfig):
+        super().__init__()
+        self.ln_1 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+        self.attn = MultiHeadAttention(config)
+        self.ln_2 = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+        self.mlp = FeedForward(config)
+    def forward(self, x, attention_mask=None):
+        x = x + self.attn(self.ln_1(x), attention_mask)
+        x = x + self.mlp(self.ln_2(x))
+        return x
+
+class ChessForCausalLM(PreTrainedModel):
+    config_class = ChessConfig
+    base_model_prefix = "transformer"
+    
+    def __init__(self, config: ChessConfig):
+        super().__init__(config)
+        self.wte = nn.Embedding(config.vocab_size, config.n_embd)
+        self.wpe = nn.Embedding(config.n_ctx, config.n_embd)
+        self.drop = nn.Dropout(config.dropout)
+        self.h = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layer)])
+        self.ln_f = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        if config.tie_weights: self.post_init()
+
+    def get_input_embeddings(self): return self.wte
+    def set_input_embeddings(self, new_embeddings): self.wte = new_embeddings
+    def get_output_embeddings(self): return self.lm_head
+    def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings
+
+    def forward(self, input_ids, attention_mask=None, position_ids=None, labels=None, return_dict=None, **kwargs):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        if return_dict is None: return_dict = True
+
+        device = input_ids.device
+        b, t = input_ids.size()
+        if position_ids is None: position_ids = torch.arange(t, device=device).unsqueeze(0)
+        
+        x = self.wte(input_ids) + self.wpe(position_ids)
+        x = self.drop(x)
+        for block in self.h: x = block(x, attention_mask)
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+        if labels is None:
+
+            nuclear_bad_ids = [0, 1, 2, 3] 
+            
+            logits[:, :, nuclear_bad_ids] = float("-inf")
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id)(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
+
+        if not return_dict:
+            return ((loss,) + (logits,)) if loss is not None else (logits,)
+        
+        return CausalLMOutputWithPast(loss=loss, logits=logits)
+
+from transformers import AutoConfig, AutoModelForCausalLM
+AutoConfig.register("chess_transformer", ChessConfig)
+AutoModelForCausalLM.register(ChessConfig, ChessForCausalLM)

pytorch_model.bin

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

tokenizer.py

@@ -9,7 +9,6 @@ import re
 from typing import Dict, List, Optional
 from transformers import PreTrainedTokenizer
 
-# --- REGEX (Pour nettoyer les coups) ---
 MOVE_RE = re.compile(r"([a-h][1-8])([a-h][1-8])")
 PROMO_RE = re.compile(r"=([NBRQ])")
 
@@ -41,7 +40,6 @@ class ChessTokenizer(PreTrainedTokenizer):
 
         for t in ["pad_token", "bos_token", "eos_token", "unk_token"]: kwargs.pop(t, None)
             
-        # FIX CHEMIN
         if vocab is None:
             if vocab_file is None:
                 vocab_file = os.path.join(os.path.dirname(__file__), "vocab.json")
@@ -56,7 +54,7 @@ class ChessTokenizer(PreTrainedTokenizer):
         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)
 
-    # AUTO-COPIE (Vital pour le submit)
+
     def save_pretrained(self, save_directory: str, **kwargs):
         super().save_pretrained(save_directory, **kwargs)
         src_path = os.path.abspath(__file__)
@@ -72,7 +70,6 @@ class ChessTokenizer(PreTrainedTokenizer):
     def _create_default_vocab(self):
         return {t: i for i, t in enumerate([self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN])}
     
-    # LA FONCTION QUI GERE LA TAILLE FIXE
     @classmethod
     def build_vocab_from_dataset(cls, dataset_name, split="train", column="text", min_frequency=2, max_vocab_size=1700, max_samples=100000):
         from datasets import load_dataset
@@ -87,9 +84,7 @@ class ChessTokenizer(PreTrainedTokenizer):
             moves = [normalize_move(t) for t in ex[column].split()]
             counter.update(moves)
             
-        # ON FORCE LA TAILLE MAXIMALE ICI
         special = [cls.PAD_TOKEN, cls.BOS_TOKEN, cls.EOS_TOKEN, cls.UNK_TOKEN]
-        # On prend les N plus fréquents pour remplir jusqu'à max_vocab_size
         most_common = counter.most_common(max_vocab_size - len(special))
         
         vocab = {t: i for i, t in enumerate(special + [t for t, c in most_common])}