hatsingue commited on
Commit
617e8d2
·
verified ·
1 Parent(s): 32deb84

Chess Challenge submission by hatsingue

Browse files
Files changed (9) hide show
  1. README.md +22 -0
  2. config.json +24 -0
  3. model.py +398 -0
  4. model.safetensors +3 -0
  5. special_tokens_map.json +6 -0
  6. tokenizer.py +259 -0
  7. tokenizer_config.json +47 -0
  8. training_args.bin +3 -0
  9. vocab.json +74 -0
README.md ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ library_name: transformers
3
+ tags:
4
+ - chess
5
+ - llm-course
6
+ - chess-challenge
7
+ license: mit
8
+ ---
9
+ # chess-pop-v2
10
+ Chess model submitted to the LLM Course Chess Challenge.
11
+ ## Submission Info
12
+ - **Submitted by**: [hatsingue](https://huggingface.co/hatsingue)
13
+ - **Parameters**: 862,336
14
+ - **Organization**: LLM-course
15
+ ## Usage
16
+ ```python
17
+ from transformers import AutoModelForCausalLM, AutoTokenizer
18
+ model = AutoModelForCausalLM.from_pretrained("LLM-course/chess-pop-v2", trust_remote_code=True)
19
+ tokenizer = AutoTokenizer.from_pretrained("LLM-course/chess-pop-v2", trust_remote_code=True)
20
+ ```
21
+ ## Evaluation
22
+ This model is evaluated at the [Chess Challenge Arena](https://huggingface.co/spaces/LLM-course/Chess1MChallenge).
config.json ADDED
@@ -0,0 +1,24 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "ChessForCausalLM"
4
+ ],
5
+ "auto_map": {
6
+ "AutoConfig": "model.ChessConfig",
7
+ "AutoModelForCausalLM": "model.ChessForCausalLM"
8
+ },
9
+ "bos_token_id": 1,
10
+ "dropout": 0.1,
11
+ "dtype": "float32",
12
+ "eos_token_id": 2,
13
+ "layer_norm_epsilon": 1e-05,
14
+ "model_type": "chess_transformer",
15
+ "n_ctx": 256,
16
+ "n_embd": 128,
17
+ "n_head": 4,
18
+ "n_inner": 384,
19
+ "n_layer": 4,
20
+ "pad_token_id": 0,
21
+ "tie_weights": true,
22
+ "transformers_version": "4.57.6",
23
+ "vocab_size": 72
24
+ }
model.py ADDED
@@ -0,0 +1,398 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Chess Transformer Model for the Chess Challenge.
3
+
4
+ This module provides a simple GPT-style transformer architecture
5
+ designed to fit within the 1M parameter constraint.
6
+
7
+ Key components:
8
+ - ChessConfig: Configuration class for model hyperparameters
9
+ - ChessForCausalLM: The main model class for next-move prediction
10
+ """
11
+
12
+ from __future__ import annotations
13
+
14
+ import math
15
+ from dataclasses import dataclass
16
+ from typing import Optional, Tuple, Union
17
+
18
+ import torch
19
+ import torch.nn as nn
20
+ import torch.nn.functional as F
21
+ from transformers import PretrainedConfig, PreTrainedModel
22
+ from transformers.modeling_outputs import CausalLMOutputWithPast
23
+
24
+
25
+ class RMSNorm(nn.Module):
26
+ def __init__(self, dim: int, eps: float = 1e-6):
27
+ super().__init__()
28
+ self.eps = eps
29
+ self.weight = nn.Parameter(torch.ones(dim))
30
+
31
+ def forward(self, x):
32
+ return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
33
+
34
+ class RotaryEmbedding(nn.Module):
35
+ def __init__(self, dim, max_seq_len=256):
36
+ super().__init__()
37
+ inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
38
+ self.register_buffer("inv_freq", inv_freq)
39
+
40
+ def forward(self, x, seq_len):
41
+ t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
42
+ freqs = torch.einsum("i,j->ij", t, self.inv_freq)
43
+ emb = torch.cat((freqs, freqs), dim=-1)
44
+ return emb[None, :, None, :]
45
+
46
+ def apply_rotary_emb(q, k, freqs):
47
+ def rotate_half(x):
48
+ x1, x2 = x.chunk(2, dim=-1)
49
+ return torch.cat((-x2, x1), dim=-1)
50
+
51
+ q_rot = (q * freqs.cos()) + (rotate_half(q) * freqs.sin())
52
+ k_rot = (k * freqs.cos()) + (rotate_half(k) * freqs.sin())
53
+ return q_rot, k_rot
54
+
55
+ class SwiGLU(nn.Module):
56
+ def __init__(self, dim: int, inner_dim: int, dropout: float):
57
+ super().__init__()
58
+ self.w1 = nn.Linear(dim, inner_dim, bias=False)
59
+ self.w2 = nn.Linear(inner_dim, dim, bias=False)
60
+ self.w3 = nn.Linear(dim, inner_dim, bias=False)
61
+ self.dropout = nn.Dropout(dropout)
62
+
63
+ def forward(self, x):
64
+ # L'essence de SwiGLU : (SiLU(W1x) * W3x) * W2
65
+ return self.dropout(self.w2(F.silu(self.w1(x)) * self.w3(x)))
66
+
67
+ class ModernAttention(nn.Module):
68
+ def __init__(self, config):
69
+ super().__init__()
70
+ self.n_head = config.n_head
71
+ self.head_dim = config.n_embd // config.n_head
72
+
73
+ self.wq = nn.Linear(config.n_embd, config.n_embd, bias=False)
74
+ self.wk = nn.Linear(config.n_embd, config.n_embd, bias=False)
75
+ self.wv = nn.Linear(config.n_embd, config.n_embd, bias=False)
76
+ self.wo = nn.Linear(config.n_embd, config.n_embd, bias=False)
77
+ self.dropout = nn.Dropout(config.dropout)
78
+
79
+ def forward(self, x, freqs, mask=None):
80
+ bsz, seqlen, _ = x.shape
81
+ q, k, v = self.wq(x), self.wk(x), self.wv(x)
82
+
83
+ q = q.view(bsz, seqlen, self.n_head, self.head_dim)
84
+ k = k.view(bsz, seqlen, self.n_head, self.head_dim)
85
+ v = v.view(bsz, seqlen, self.n_head, self.head_dim)
86
+
87
+ q, k = apply_rotary_emb(q, k, freqs)
88
+
89
+ scores = torch.matmul(q.transpose(1, 2), k.transpose(1, 2).transpose(-2, -1)) / math.sqrt(self.head_dim)
90
+
91
+ if mask is not None:
92
+ scores = scores + mask[:, :, :seqlen, :seqlen]
93
+
94
+ scores = F.softmax(scores.float(), dim=-1).type_as(q)
95
+ output = torch.matmul(scores, v.transpose(1, 2))
96
+ output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
97
+ return self.dropout(self.wo(output))
98
+
99
+ class ModernBlock(nn.Module):
100
+ def __init__(self, config):
101
+ super().__init__()
102
+ self.attention = ModernAttention(config)
103
+ self.feed_forward = SwiGLU(config.n_embd, config.n_inner, config.dropout)
104
+ self.attention_norm = RMSNorm(config.n_embd)
105
+ self.ffn_norm = RMSNorm(config.n_embd)
106
+
107
+ def forward(self, x, freqs, mask):
108
+ x = x + self.attention(self.attention_norm(x), freqs, mask)
109
+ x = x + self.feed_forward(self.ffn_norm(x))
110
+ return x
111
+
112
+
113
+
114
+
115
+
116
+ class ChessConfig(PretrainedConfig):
117
+ """
118
+ Configuration class for the Chess Transformer model.
119
+
120
+ This configuration is designed for a ~1M parameter model.
121
+ Students can adjust these values to explore different architectures.
122
+
123
+ Parameter budget breakdown (with default values):
124
+ - Embeddings (vocab): 1200 x 128 = 153,600
125
+ - Position Embeddings: 256 x 128 = 32,768
126
+ - Transformer Layers: 6 x ~120,000 = ~720,000
127
+ - LM Head (with weight tying): 0 (shared with embeddings)
128
+ - Total: ~906,000 parameters
129
+
130
+ Attributes:
131
+ vocab_size: Size of the vocabulary (number of unique moves).
132
+ n_embd: Embedding dimension (d_model).
133
+ n_layer: Number of transformer layers.
134
+ n_head: Number of attention heads.
135
+ n_ctx: Maximum sequence length (context window).
136
+ n_inner: Feed-forward inner dimension (default: 3 * n_embd).
137
+ dropout: Dropout probability.
138
+ layer_norm_epsilon: Epsilon for layer normalization.
139
+ tie_weights: Whether to tie embedding and output weights.
140
+ """
141
+
142
+ model_type = "chess_transformer"
143
+
144
+ def __init__(
145
+ self,
146
+ vocab_size: int = 1200,
147
+ n_embd: int = 128,
148
+ n_layer: int = 6,
149
+ n_head: int = 8,
150
+ n_ctx: int = 256,
151
+ n_inner: Optional[int] = None,
152
+ dropout: float = 0.1,
153
+ layer_norm_epsilon: float = 1e-5,
154
+ tie_weights: bool = True,
155
+ pad_token_id: int = 0,
156
+ bos_token_id: int = 1,
157
+ eos_token_id: int = 2,
158
+ **kwargs,
159
+ ):
160
+ super().__init__(
161
+ pad_token_id=pad_token_id,
162
+ bos_token_id=bos_token_id,
163
+ eos_token_id=eos_token_id,
164
+ **kwargs,
165
+ )
166
+
167
+ self.vocab_size = vocab_size
168
+ self.n_embd = n_embd
169
+ self.n_layer = n_layer
170
+ self.n_head = n_head
171
+ self.n_ctx = n_ctx
172
+ self.n_inner = n_inner if n_inner is not None else 3 * n_embd # Reduced from 4x to 3x
173
+ self.dropout = dropout
174
+ self.layer_norm_epsilon = layer_norm_epsilon
175
+ self.tie_weights = tie_weights
176
+ # Inform HF base class about tying behavior
177
+ self.tie_word_embeddings = bool(tie_weights)
178
+
179
+
180
+ class MultiHeadAttention(nn.Module):
181
+ """
182
+ Multi-head self-attention module.
183
+
184
+ This is a standard scaled dot-product attention implementation
185
+ with causal masking for autoregressive generation.
186
+ """
187
+
188
+ def __init__(self, config: ChessConfig):
189
+ super().__init__()
190
+
191
+ assert config.n_embd % config.n_head == 0, \
192
+ f"n_embd ({config.n_embd}) must be divisible by n_head ({config.n_head})"
193
+
194
+ self.n_head = config.n_head
195
+ self.n_embd = config.n_embd
196
+ self.head_dim = config.n_embd // config.n_head
197
+
198
+ # Combined QKV projection for efficiency
199
+ self.c_attn = nn.Linear(config.n_embd, 3 * config.n_embd)
200
+ self.c_proj = nn.Linear(config.n_embd, config.n_embd)
201
+
202
+ self.dropout = nn.Dropout(config.dropout)
203
+
204
+ # Causal mask (will be created on first forward pass)
205
+ self.register_buffer(
206
+ "bias",
207
+ torch.tril(torch.ones(config.n_ctx, config.n_ctx)).view(
208
+ 1, 1, config.n_ctx, config.n_ctx
209
+ ),
210
+ persistent=False,
211
+ )
212
+
213
+ def forward(
214
+ self,
215
+ x: torch.Tensor,
216
+ attention_mask: Optional[torch.Tensor] = None,
217
+ ) -> torch.Tensor:
218
+ batch_size, seq_len, _ = x.size()
219
+
220
+ # Compute Q, K, V
221
+ qkv = self.c_attn(x)
222
+ q, k, v = qkv.split(self.n_embd, dim=2)
223
+
224
+ # Reshape for multi-head attention
225
+ q = q.view(batch_size, seq_len, self.n_head, self.head_dim).transpose(1, 2)
226
+ k = k.view(batch_size, seq_len, self.n_head, self.head_dim).transpose(1, 2)
227
+ v = v.view(batch_size, seq_len, self.n_head, self.head_dim).transpose(1, 2)
228
+
229
+ # Scaled dot-product attention
230
+ attn_weights = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
231
+
232
+ # Apply causal mask
233
+ causal_mask = self.bias[:, :, :seq_len, :seq_len]
234
+ attn_weights = attn_weights.masked_fill(causal_mask == 0, float("-inf"))
235
+
236
+ # Apply attention mask (for padding)
237
+ if attention_mask is not None:
238
+ # attention_mask shape: (batch_size, seq_len) -> (batch_size, 1, 1, seq_len)
239
+ attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
240
+ attn_weights = attn_weights.masked_fill(attention_mask == 0, float("-inf"))
241
+
242
+ attn_weights = F.softmax(attn_weights, dim=-1)
243
+ attn_weights = self.dropout(attn_weights)
244
+
245
+ # Apply attention to values
246
+ attn_output = torch.matmul(attn_weights, v)
247
+
248
+ # Reshape back
249
+ attn_output = attn_output.transpose(1, 2).contiguous().view(
250
+ batch_size, seq_len, self.n_embd
251
+ )
252
+
253
+ # Output projection
254
+ attn_output = self.c_proj(attn_output)
255
+
256
+ return attn_output
257
+
258
+
259
+ class FeedForward(nn.Module):
260
+ """
261
+ Feed-forward network (MLP) module.
262
+
263
+ Standard two-layer MLP with GELU activation.
264
+ """
265
+
266
+ def __init__(self, config: ChessConfig):
267
+ super().__init__()
268
+
269
+ self.c_fc = nn.Linear(config.n_embd, config.n_inner)
270
+ self.c_proj = nn.Linear(config.n_inner, config.n_embd)
271
+ self.dropout = nn.Dropout(config.dropout)
272
+
273
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
274
+ x = self.c_fc(x)
275
+ x = F.gelu(x)
276
+ x = self.c_proj(x)
277
+ x = self.dropout(x)
278
+ return x
279
+
280
+
281
+ class TransformerBlock(nn.Module):
282
+ """
283
+ A single transformer block with attention and feed-forward layers.
284
+
285
+ Uses pre-normalization (LayerNorm before attention/FFN) for better
286
+ training stability.
287
+ """
288
+
289
+ def __init__(self, config: ChessConfig):
290
+ super().__init__()
291
+
292
+ self.ln_1 = RMSNorm(config.n_embd, eps=config.layer_norm_epsilon)
293
+ self.attn = MultiHeadAttention(config)
294
+ self.ln_2 = RMSNorm(config.n_embd, eps=config.layer_norm_epsilon)
295
+ self.mlp = FeedForward(config)
296
+
297
+ def forward(
298
+ self,
299
+ x: torch.Tensor,
300
+ attention_mask: Optional[torch.Tensor] = None,
301
+ ) -> torch.Tensor:
302
+ # Pre-norm attention
303
+ x = x + self.attn(self.ln_1(x), attention_mask=attention_mask)
304
+ # Pre-norm FFN
305
+ x = x + self.mlp(self.ln_2(x))
306
+ return x
307
+
308
+
309
+ class ChessForCausalLM(PreTrainedModel):
310
+ config_class = ChessConfig
311
+ _tied_weights_keys = ["lm_head.weight"]
312
+
313
+ def __init__(self, config: ChessConfig):
314
+ super().__init__(config)
315
+
316
+ self.wte = nn.Embedding(config.vocab_size, config.n_embd)
317
+
318
+ self.rope = RotaryEmbedding(config.n_embd // config.n_head)
319
+
320
+ self.drop = nn.Dropout(config.dropout)
321
+ self.h = nn.ModuleList([ModernBlock(config) for _ in range(config.n_layer)])
322
+ self.ln_f = RMSNorm(config.n_embd, eps=config.layer_norm_epsilon)
323
+ self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
324
+
325
+ self.post_init()
326
+ if config.tie_weights:
327
+ self.tie_weights()
328
+
329
+
330
+ def forward(
331
+ self,
332
+ input_ids: torch.LongTensor,
333
+ attention_mask: Optional[torch.Tensor] = None,
334
+ labels: Optional[torch.LongTensor] = None,
335
+ return_dict: Optional[bool] = None,
336
+ **kwargs,
337
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
338
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
339
+ batch_size, seq_len = input_ids.size()
340
+ device = input_ids.device
341
+
342
+ freqs = self.rope(input_ids, seq_len)
343
+
344
+ mask = torch.full((seq_len, seq_len), float("-inf"), device=device)
345
+ mask = torch.triu(mask, diagonal=1)
346
+ mask = mask.view(1, 1, seq_len, seq_len)
347
+
348
+ hidden_states = self.drop(self.wte(input_ids))
349
+
350
+ for block in self.h:
351
+ hidden_states = block(hidden_states, freqs, mask)
352
+
353
+ hidden_states = self.ln_f(hidden_states)
354
+ logits = self.lm_head(hidden_states)
355
+
356
+ loss = None
357
+ if labels is not None:
358
+ shift_logits = logits[..., :-1, :].contiguous()
359
+ shift_labels = labels[..., 1:].contiguous()
360
+ loss = F.cross_entropy(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
361
+
362
+ if not return_dict:
363
+ output = (logits,)
364
+ return ((loss,) + output) if loss is not None else output
365
+
366
+ return CausalLMOutputWithPast(
367
+ loss=loss,
368
+ logits=logits,
369
+ past_key_values=None,
370
+ hidden_states=None,
371
+ attentions=None,
372
+ )
373
+
374
+ def get_input_embeddings(self):
375
+ return self.wte
376
+
377
+ def set_input_embeddings(self, value):
378
+ self.wte = value
379
+
380
+ def get_output_embeddings(self):
381
+ return self.lm_head
382
+
383
+ def set_output_embeddings(self, new_embeddings):
384
+ self.lm_head = new_embeddings
385
+
386
+ def tie_weights(self):
387
+ """
388
+ C'est cette méthode que HF appelle automatiquement si
389
+ config.tie_word_embeddings est True.
390
+ """
391
+ self._tie_or_clone_weights(self.lm_head, self.wte)
392
+
393
+
394
+ # Register the model with Auto classes for easy loading
395
+ from transformers import AutoConfig, AutoModelForCausalLM
396
+
397
+ AutoConfig.register("chess_transformer", ChessConfig)
398
+ AutoModelForCausalLM.register(ChessConfig, ChessForCausalLM)
model.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:41f52b7be1bfbb8b62384dc57f7c41eb08409457995a328cb135e8d44ac9e22a
3
+ size 3453000
special_tokens_map.json ADDED
@@ -0,0 +1,6 @@
 
 
 
 
 
 
 
1
+ {
2
+ "bos_token": "[BOS]",
3
+ "eos_token": "[EOS]",
4
+ "pad_token": "[PAD]",
5
+ "unk_token": "[UNK]"
6
+ }
tokenizer.py ADDED
@@ -0,0 +1,259 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Custom Chess Tokenizer for the Chess Challenge.
3
+
4
+ This tokenizer treats each move as a single token using the extended UCI notation
5
+ from the Lichess dataset (e.g., WPe2e4, BNg8f6).
6
+
7
+ The dataset format uses:
8
+ - W/B prefix for White/Black
9
+ - Piece letter: P=Pawn, N=Knight, B=Bishop, R=Rook, Q=Queen, K=King
10
+ - Source and destination squares (e.g., e2e4)
11
+ - Special suffixes: (x)=capture, (+)=check, (+*)=checkmate, (o)/(O)=castling
12
+ """
13
+
14
+ from __future__ import annotations
15
+
16
+ import json
17
+ import os
18
+ from pathlib import Path
19
+ from typing import Dict, List, Optional
20
+ import re
21
+
22
+ from transformers import PreTrainedTokenizer
23
+
24
+
25
+ class ChessTokenizer(PreTrainedTokenizer):
26
+ """
27
+ A custom tokenizer for chess moves using extended UCI notation.
28
+
29
+ This tokenizer maps each possible chess move to a unique token ID.
30
+ The vocabulary is built from the training dataset to ensure all moves
31
+ encountered during training have a corresponding token.
32
+
33
+ Example:
34
+ >>> tokenizer = ChessTokenizer()
35
+ >>> tokenizer.encode("WPe2e4 BPe7e5")
36
+ [1, 42, 87, 2] # [BOS, e2e4, e7e5, EOS]
37
+ """
38
+
39
+ model_input_names = ["input_ids", "attention_mask"]
40
+ vocab_files_names = {"vocab_file": "vocab.json"}
41
+
42
+ # Special tokens
43
+ PAD_TOKEN = "[PAD]"
44
+ BOS_TOKEN = "[BOS]"
45
+ EOS_TOKEN = "[EOS]"
46
+ UNK_TOKEN = "[UNK]"
47
+
48
+ def __init__(
49
+ self,
50
+ vocab_file: Optional[str] = None,
51
+ vocab: Optional[Dict[str, int]] = None,
52
+ **kwargs,
53
+ ):
54
+ """
55
+ Initialize the chess tokenizer.
56
+
57
+ Args:
58
+ vocab_file: Path to a JSON file containing the vocabulary mapping.
59
+ vocab: Dictionary mapping tokens to IDs (alternative to vocab_file).
60
+ **kwargs: Additional arguments passed to PreTrainedTokenizer.
61
+ """
62
+ # Initialize special tokens
63
+ self._pad_token = self.PAD_TOKEN
64
+ self._bos_token = self.BOS_TOKEN
65
+ self._eos_token = self.EOS_TOKEN
66
+ self._unk_token = self.UNK_TOKEN
67
+
68
+ # Remove any duplicate special-token entries passed through kwargs
69
+ # to avoid "multiple values for keyword" errors when loading from disk.
70
+ kwargs.pop("pad_token", None)
71
+ kwargs.pop("bos_token", None)
72
+ kwargs.pop("eos_token", None)
73
+ kwargs.pop("unk_token", None)
74
+
75
+ self.token_pattern = re.compile(r'[a-h][1-8]|[qrbn]')
76
+
77
+ # Load or create vocabulary
78
+ if vocab is not None:
79
+ self._vocab = vocab
80
+ elif vocab_file is not None and os.path.exists(vocab_file):
81
+ with open(vocab_file, "r", encoding="utf-8") as f:
82
+ self._vocab = json.load(f)
83
+ else:
84
+ # Create a minimal vocabulary with just special tokens
85
+ # The full vocabulary should be built from the dataset
86
+ self._vocab = self._create_default_vocab()
87
+
88
+ # Create reverse mapping
89
+ self._ids_to_tokens = {v: k for k, v in self._vocab.items()}
90
+
91
+ # Call parent init AFTER setting up vocab
92
+ super().__init__(
93
+ pad_token=self._pad_token,
94
+ bos_token=self._bos_token,
95
+ eos_token=self._eos_token,
96
+ unk_token=self._unk_token,
97
+ **kwargs,
98
+ )
99
+
100
+ def _create_default_vocab(self) -> Dict[str, int]:
101
+ """
102
+ Create a minimal default vocabulary with just special tokens.
103
+
104
+ For the full vocabulary, use `build_vocab_from_dataset()`.
105
+ This minimal vocab is just a placeholder - you should build from data.
106
+ """
107
+ special_tokens = [self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN]
108
+ vocab = {token: idx for idx, token in enumerate(special_tokens)}
109
+ n = len(vocab)
110
+ for f in 'abcdefgh':
111
+ for r in '12345678':
112
+ vocab[f"{f}{r}"] = n
113
+ n += 1
114
+
115
+ for p in ['q', 'r', 'b', 'n']:
116
+ vocab[p] = n
117
+ n += 1
118
+ return vocab
119
+
120
+ @classmethod
121
+ def build_vocab_from_iterator(
122
+ cls,
123
+ iterator,
124
+ min_frequency: int = 1,
125
+ ) -> "ChessTokenizer":
126
+ """
127
+ Build a tokenizer vocabulary from an iterator of game strings.
128
+
129
+ Args:
130
+ iterator: An iterator yielding game strings (space-separated moves).
131
+ min_frequency: Minimum frequency for a token to be included.
132
+
133
+ Returns:
134
+ A ChessTokenizer with the built vocabulary.
135
+ """
136
+ return cls()
137
+
138
+ @classmethod
139
+ def build_vocab_from_dataset(
140
+ cls,
141
+ dataset_name: str = "dlouapre/lichess_2025-01_1M",
142
+ split: str = "train",
143
+ column: str = "text",
144
+ min_frequency: int = 500,
145
+ max_samples: Optional[int] = 100000,
146
+ ) -> "ChessTokenizer":
147
+ return cls()
148
+
149
+ @property
150
+ def vocab_size(self) -> int:
151
+ """Return the size of the vocabulary."""
152
+ return len(self._vocab)
153
+
154
+ def get_vocab(self) -> Dict[str, int]:
155
+ """Return the vocabulary as a dictionary."""
156
+ return dict(self._vocab)
157
+
158
+ def _tokenize(self, text: str) -> List[str]:
159
+ """
160
+ Tokenize a string of moves into a list of tokens.
161
+
162
+ Args:
163
+ text: A string of space-separated moves.
164
+
165
+ Returns:
166
+ List of move tokens.
167
+ """
168
+ text = (text.replace("(Q)", "q")
169
+ .replace("(R)", "r")
170
+ .replace("(B)", "b")
171
+ .replace("(N)", "n"))
172
+ return self.token_pattern.findall(text)
173
+
174
+ def _convert_token_to_id(self, token: str) -> int:
175
+ """Convert a token to its ID."""
176
+ return self._vocab.get(token, self._vocab.get(self.UNK_TOKEN, 0))
177
+
178
+ def _convert_id_to_token(self, index: int) -> str:
179
+ """Convert an ID to its token."""
180
+ return self._ids_to_tokens.get(index, self.UNK_TOKEN)
181
+
182
+ def convert_tokens_to_string(self, tokens: List[str]) -> str:
183
+ """Convert a list of tokens back to a string."""
184
+ special_tokens = {self.PAD_TOKEN, self.BOS_TOKEN, self.EOS_TOKEN, self.UNK_TOKEN}
185
+ clean_tokens = [t for t in tokens if t not in special_tokens]
186
+
187
+ output = []
188
+ for token in clean_tokens:
189
+ if token in ['q', 'r', 'b', 'n'] and output:
190
+ output[-1] += token
191
+ elif output and len(output[-1]) == 2 and output[-1][0] in 'abcdefgh':
192
+ output[-1] += token
193
+ else:
194
+ output.append(token)
195
+
196
+ return " ".join(output)
197
+
198
+ def save_vocabulary(
199
+ self,
200
+ save_directory: str,
201
+ filename_prefix: Optional[str] = None,
202
+ ) -> tuple:
203
+ """
204
+ Save the vocabulary to a JSON file.
205
+
206
+ Args:
207
+ save_directory: Directory to save the vocabulary.
208
+ filename_prefix: Optional prefix for the filename.
209
+
210
+ Returns:
211
+ Tuple containing the path to the saved vocabulary file.
212
+ """
213
+ if not os.path.isdir(save_directory):
214
+ os.makedirs(save_directory, exist_ok=True)
215
+
216
+ vocab_file = os.path.join(
217
+ save_directory,
218
+ (filename_prefix + "-" if filename_prefix else "") + "vocab.json",
219
+ )
220
+
221
+ with open(vocab_file, "w", encoding="utf-8") as f:
222
+ json.dump(self._vocab, f, ensure_ascii=False, indent=2)
223
+
224
+ return (vocab_file,)
225
+
226
+
227
+ def count_vocab_from_dataset(
228
+ dataset_name: str = "dlouapre/lichess_2025-01_1M",
229
+ split: str = "train",
230
+ column: str = "text",
231
+ max_samples: Optional[int] = 10000,
232
+ ) -> Dict[str, int]:
233
+ """
234
+ Count token frequencies in a dataset (useful for vocabulary analysis).
235
+
236
+ Args:
237
+ dataset_name: Name of the dataset on Hugging Face Hub.
238
+ split: Dataset split to use.
239
+ column: Column containing the game strings.
240
+ max_samples: Maximum number of samples to process.
241
+
242
+ Returns:
243
+ Dictionary mapping tokens to their frequencies.
244
+ """
245
+ from collections import Counter
246
+ from datasets import load_dataset
247
+
248
+ dataset = load_dataset(dataset_name, split=split)
249
+
250
+ if max_samples is not None:
251
+ dataset = dataset.select(range(min(max_samples, len(dataset))))
252
+
253
+ token_counts = Counter()
254
+
255
+ for example in dataset:
256
+ moves = example[column].strip().split()
257
+ token_counts.update(moves)
258
+
259
+ return dict(token_counts)
tokenizer_config.json ADDED
@@ -0,0 +1,47 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "added_tokens_decoder": {
3
+ "0": {
4
+ "content": "[PAD]",
5
+ "lstrip": false,
6
+ "normalized": false,
7
+ "rstrip": false,
8
+ "single_word": false,
9
+ "special": true
10
+ },
11
+ "1": {
12
+ "content": "[BOS]",
13
+ "lstrip": false,
14
+ "normalized": false,
15
+ "rstrip": false,
16
+ "single_word": false,
17
+ "special": true
18
+ },
19
+ "2": {
20
+ "content": "[EOS]",
21
+ "lstrip": false,
22
+ "normalized": false,
23
+ "rstrip": false,
24
+ "single_word": false,
25
+ "special": true
26
+ },
27
+ "3": {
28
+ "content": "[UNK]",
29
+ "lstrip": false,
30
+ "normalized": false,
31
+ "rstrip": false,
32
+ "single_word": false,
33
+ "special": true
34
+ }
35
+ },
36
+ "auto_map": {
37
+ "AutoTokenizer": ["tokenizer.ChessTokenizer", null]
38
+ },
39
+ "bos_token": "[BOS]",
40
+ "clean_up_tokenization_spaces": false,
41
+ "eos_token": "[EOS]",
42
+ "extra_special_tokens": {},
43
+ "model_max_length": 1000000000000000019884624838656,
44
+ "pad_token": "[PAD]",
45
+ "tokenizer_class": "ChessTokenizer",
46
+ "unk_token": "[UNK]"
47
+ }
training_args.bin ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:38a385687f3a5fcb30b7768fa4b9fe6a2526b5ce5c855b55284fd0ee550c6a8a
3
+ size 5777
vocab.json ADDED
@@ -0,0 +1,74 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "[PAD]": 0,
3
+ "[BOS]": 1,
4
+ "[EOS]": 2,
5
+ "[UNK]": 3,
6
+ "a1": 4,
7
+ "a2": 5,
8
+ "a3": 6,
9
+ "a4": 7,
10
+ "a5": 8,
11
+ "a6": 9,
12
+ "a7": 10,
13
+ "a8": 11,
14
+ "b1": 12,
15
+ "b2": 13,
16
+ "b3": 14,
17
+ "b4": 15,
18
+ "b5": 16,
19
+ "b6": 17,
20
+ "b7": 18,
21
+ "b8": 19,
22
+ "c1": 20,
23
+ "c2": 21,
24
+ "c3": 22,
25
+ "c4": 23,
26
+ "c5": 24,
27
+ "c6": 25,
28
+ "c7": 26,
29
+ "c8": 27,
30
+ "d1": 28,
31
+ "d2": 29,
32
+ "d3": 30,
33
+ "d4": 31,
34
+ "d5": 32,
35
+ "d6": 33,
36
+ "d7": 34,
37
+ "d8": 35,
38
+ "e1": 36,
39
+ "e2": 37,
40
+ "e3": 38,
41
+ "e4": 39,
42
+ "e5": 40,
43
+ "e6": 41,
44
+ "e7": 42,
45
+ "e8": 43,
46
+ "f1": 44,
47
+ "f2": 45,
48
+ "f3": 46,
49
+ "f4": 47,
50
+ "f5": 48,
51
+ "f6": 49,
52
+ "f7": 50,
53
+ "f8": 51,
54
+ "g1": 52,
55
+ "g2": 53,
56
+ "g3": 54,
57
+ "g4": 55,
58
+ "g5": 56,
59
+ "g6": 57,
60
+ "g7": 58,
61
+ "g8": 59,
62
+ "h1": 60,
63
+ "h2": 61,
64
+ "h3": 62,
65
+ "h4": 63,
66
+ "h5": 64,
67
+ "h6": 65,
68
+ "h7": 66,
69
+ "h8": 67,
70
+ "q": 68,
71
+ "r": 69,
72
+ "b": 70,
73
+ "n": 71
74
+ }