import torch import torch.nn as nn from typing import List, Dict, Tuple from huggingface_hub import PyTorchModelHubMixin try: from mapping import MAX_HALFMOVES, MAX_FULLMOVES, EMPTY_SQ_IDX, PIECE_TO_IDX, SQUARE_TO_IDX, IDX_TO_UCI_MOVE except: from .mapping import MAX_HALFMOVES, MAX_FULLMOVES, EMPTY_SQ_IDX, PIECE_TO_IDX, SQUARE_TO_IDX, IDX_TO_UCI_MOVE # --- Tokenizer --- # class FENTokenizer(nn.Module): """Convert FEN (and repetitions) to a sequence of tokens""" def __init__(self, hidden_size,dtype): super().__init__() self.side_embed = nn.Embedding(2,hidden_size,dtype=dtype) # black/white embedding self.castling_embed_k = nn.Parameter(torch.randn(1,1,hidden_size,dtype=dtype)) self.castling_embed_q = nn.Parameter(torch.randn(1,1,hidden_size,dtype=dtype)) self.castling_embed_K = nn.Parameter(torch.randn(1,1,hidden_size,dtype=dtype)) self.castling_embed_Q = nn.Parameter(torch.randn(1,1,hidden_size,dtype=dtype)) self.no_castling_embed = nn.Parameter(torch.randn(1,1,hidden_size,dtype=dtype)) self.piece_embed = nn.Embedding(13,hidden_size,dtype=dtype) # 6 for white pieces, 6 for black pieces, 1 for empty self.no_en_passant_embed = nn.Parameter(torch.randn(1,1,hidden_size,dtype=dtype)) # use positional embed for the target square, or a special one for '-' self.half_move_embed = nn.Embedding(MAX_HALFMOVES,hidden_size,dtype=dtype) self.full_move_embed = nn.Embedding(MAX_FULLMOVES,hidden_size,dtype=dtype) self.repetition_embed = nn.Embedding(3,hidden_size,dtype=dtype) self.pos_embed = nn.Embedding(64,hidden_size,dtype=dtype) # positional embedding def _parse_fen_string(self, fen_str: str) -> Dict: parts = fen_str.split() if len(parts) != 6: raise ValueError(f"Invalid FEN string: {fen_str}. Expected 6 fields") return { "piece_placement": parts[0], "side_to_move": parts[1], "castling": parts[2], "en_passant": parts[3], "halfmove_clock": parts[4], "fullmove_number": parts[5], } def forward(self, fen_list: List[str], repetitions: torch.Tensor) -> torch.Tensor: """ Args: fen: List of fen strings Returns: torch tensor of shape (n_fen,73,hidden_size) where 73 tokens consists of: 64 piece tokens (fen's first field) + 1 which-side-to-move token (fen's second field) + 4 casting rights tokens (fen's third field) + 1 en-passant target token (fen's fourth field) + 1 half move clock token (fen's fifth field) + 1 full move number token (fen's fifth field) + 1 repetition count token (repetitions input) """ batch_size = len(fen_list) assert batch_size == repetitions.shape[0] assert len(repetitions.size()) == 1 batch_tokens = [] device = self.side_embed.weight.device # Precompute all square indices square_indices = torch.arange(64, device=device) all_pos_embeds = self.pos_embed(square_indices) # (64,D) for fen_str in fen_list: parsed_fen = self._parse_fen_string(fen_str) tokens = [] # --- 1. Piece Placement (64 tokens) --- piece_indices = torch.full((64,), EMPTY_SQ_IDX, dtype=torch.long, device=device) current_rank = 7 # Start from rank 8 current_file = 0 # Start from file 'a' for char in parsed_fen["piece_placement"]: if char == '/': current_rank -= 1 current_file = 0 elif char.isdigit(): current_file += int(char) elif char in PIECE_TO_IDX: sq_idx = current_rank * 8 + current_file if 0 <= sq_idx < 64: piece_indices[sq_idx] = PIECE_TO_IDX[char] else: raise ValueError(f"Invalid FEN piece placement: {parsed_fen['piece_placement']}") current_file += 1 else: raise ValueError(f"Invalid character in FEN piece placement: {char}") piece_embeds = self.piece_embed(piece_indices) # (64, D) # Add positional embeddings board_tokens = piece_embeds + all_pos_embeds # (64, D) tokens.append(board_tokens) # --- 2. Side to Move (1 token) --- side_idx = 0 if parsed_fen["side_to_move"] == 'w' else 1 side_token = self.side_embed(torch.tensor(side_idx, device=device)).unsqueeze(0) # (1, D) tokens.append(side_token) # --- 3. Castling Rights (4 tokens) --- castling_str = parsed_fen["castling"] castling_tokens = torch.cat([ self.castling_embed_K if 'K' in castling_str else self.no_castling_embed.expand(1, 1, -1), self.castling_embed_Q if 'Q' in castling_str else self.no_castling_embed.expand(1, 1, -1), self.castling_embed_k if 'k' in castling_str else self.no_castling_embed.expand(1, 1, -1), self.castling_embed_q if 'q' in castling_str else self.no_castling_embed.expand(1, 1, -1) ], dim=1).squeeze(0) # (4, D) tokens.append(castling_tokens) # --- 4. En Passant Target (1 token) --- en_passant_str = parsed_fen["en_passant"] if en_passant_str == '-': en_passant_token = self.no_en_passant_embed.squeeze(0) # (1, D) else: if en_passant_str in SQUARE_TO_IDX: sq_idx = SQUARE_TO_IDX[en_passant_str] en_passant_token = self.pos_embed(torch.tensor(sq_idx, device=device)).unsqueeze(0) # (1, D) else: raise ValueError(f"Invalid en passant square: {en_passant_str}") tokens.append(en_passant_token) # --- 5. Half Move Clock (1 token) --- try: half_move_int = int(parsed_fen["halfmove_clock"]) except ValueError: raise ValueError(f"Invalid halfmove clock value: {parsed_fen['halfmove_clock']}") # Clamp value before embedding lookup half_move_clamped = torch.clamp(torch.tensor(half_move_int, device=device), 0, MAX_HALFMOVES - 1) half_move_token = self.half_move_embed(half_move_clamped).unsqueeze(0) # (1, D) tokens.append(half_move_token) # --- 6. Full Move Number (1 token) --- try: full_move_int = int(parsed_fen["fullmove_number"]) except ValueError: raise ValueError(f"Invalid fullmove number value: {parsed_fen['fullmove_number']}") # Clamp value (min 1 for full moves) before embedding lookup (adjusting for 0-based index) full_move_clamped = torch.clamp(torch.tensor(full_move_int, device=device), 1, MAX_FULLMOVES) - 1 full_move_token = self.full_move_embed(full_move_clamped).unsqueeze(0) # (1, D) tokens.append(full_move_token) # Concatenate all tokens for this FEN string # Shapes: (64, D), (1, D), (4, D), (1, D), (1, D), (1, D) -> Total 72 tokens fen_embedding = torch.cat(tokens, dim=0) # (72, D) batch_tokens.append(fen_embedding) # Stack into a batch batch_tokens = torch.stack(batch_tokens, dim=0) # (B,72,D) # ---7. Repetition Count (1 token) --- repetitions = repetitions - 1 # from 1~3 to 0~2 repetitions = torch.clamp(repetitions,0,2) # if repetition count >3 but no player claimed a draw, it will be treated as 3 repetitions repetition_tokens = self.repetition_embed(repetitions) # (B,D) repetition_tokens = repetition_tokens.unsqueeze(1) # (B,1,D) return torch.cat([batch_tokens,repetition_tokens], dim=1) # (B, 73, D) # --- Helper Modules --- # class SwiGLUFFN(nn.Module): def __init__(self, d_model, dim_feedforward, dropout: float, bias_up: bool=False, bias_gate: bool=False, bias_down: bool=True, dtype=None): super().__init__() self.up_proj = nn.Linear(d_model,dim_feedforward,bias=bias_up,dtype=dtype) self.gate_proj = nn.Linear(d_model,dim_feedforward,bias=bias_gate,dtype=dtype) self.down_proj = nn.Linear(dim_feedforward,d_model,bias=bias_down,dtype=dtype) self.dropout = nn.Dropout(dropout) def forward(self, x): x = self.up_proj(x) * self.dropout(nn.functional.silu(self.gate_proj(x))) return self.down_proj(x) class TransformerEncoderLayer(nn.Module): """Custom transformer encoder layer with RMSNorm and SwiGLUFFN""" def __init__(self, d_model: int, nhead: int, dim_feedforward: int, dropout: float, batch_first: bool=True, norm_first: bool=False, dtype=None): super().__init__() self.norm_first = norm_first self.norm1 = nn.RMSNorm(d_model,dtype=dtype) self.dropout_sa = nn.Dropout(dropout) self.self_attn = nn.MultiheadAttention( d_model, nhead, dropout=dropout, bias=False, batch_first=batch_first, dtype=dtype ) self.norm2 = nn.RMSNorm(d_model,dtype=dtype) self.dropout_ff = nn.Dropout(dropout) self.mlp = SwiGLUFFN( d_model, dim_feedforward, dropout=dropout, bias_up=False, bias_gate=False, bias_down=True, dtype=dtype ) def forward(self, x, return_attention=False): if self.norm_first: if return_attention: x_norm = self.norm1(x) attn_output, attn_weights = self._sa_block(x_norm,return_attention=True) x = x + attn_output x = x + self._ff_block(self.norm2(x)) return x, attn_weights else: x = x + self._sa_block(self.norm1(x)) x = x + self._ff_block(self.norm2(x)) return x else: if return_attention: attn_output, attn_weights = self._sa_block(x, return_attention=True) x = self.norm1(x + attn_output) x = self.norm2(x + self._ff_block(x)) return x, attn_weights else: x = self.norm1(x + self._sa_block(x)) x = self.norm2(x + self._ff_block(x)) return x def _sa_block(self, x, return_attention=False): if return_attention: attn_output, attn_weights = self.self_attn(x,x,x,need_weights=True,average_attn_weights=False) return self.dropout_sa(attn_output), attn_weights else: x = self.self_attn(x,x,x)[0] return self.dropout_sa(x) def _ff_block(self,x): x = self.mlp(x) return self.dropout_ff(x) nn.TransformerEncoderLayer # --- Model Arch --- # class ChessFormerModel(nn.Module, PyTorchModelHubMixin): def __init__(self, num_blocks, hidden_size, intermediate_size, num_heads, dropout: float=0.00, possible_moves: int=len(IDX_TO_UCI_MOVE), # 1969 structurally valid moves dtype=None): super().__init__() self.fen_tokenizer = FENTokenizer(hidden_size,dtype=dtype) self.act_token = nn.Parameter(torch.randn((1,1,hidden_size),dtype=dtype) * 0.02) self.val_token = nn.Parameter(torch.randn((1,1,hidden_size),dtype=dtype) * 0.02) self.act_proj = nn.Linear(hidden_size,possible_moves,dtype=dtype) self.val_proj = nn.Linear(hidden_size,1,dtype=dtype) self.blocks = nn.ModuleList( TransformerEncoderLayer( d_model=hidden_size, nhead=num_heads, dim_feedforward=intermediate_size, dropout=dropout, batch_first=True, norm_first=True, dtype=dtype ) for _ in range(num_blocks) ) self.dtype=dtype self.possible_moves = possible_moves self.final_norm = nn.RMSNorm(hidden_size) self._initialize_weights() def _initialize_weights(self): """Initialize weights""" for m in self.modules(): if isinstance(m,nn.Linear): nn.init.kaiming_normal_(m.weight,mode='fan_in',nonlinearity='relu') if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.Embedding): nn.init.normal_(m.weight, std=0.02) elif isinstance(m, nn.LayerNorm): if hasattr(m, 'weight'): nn.init.constant_(m.weight, 1.0) if hasattr(m, 'bias') and m.bias is not None: nn.init.constant_(m.weight, 0.0) elif isinstance(m, nn.RMSNorm): if hasattr(m, 'weight'): nn.init.constant_(m.weight, 1.0) tokenizer_params = dict(self.fen_tokenizer.named_parameters()) params_to_init = [ self.act_token, self.val_token, tokenizer_params.get('castling_embed_k'), tokenizer_params.get('castling_embed_q'), tokenizer_params.get('castling_embed_K'), tokenizer_params.get('castling_embed_Q'), tokenizer_params.get('no_castling_embed'), tokenizer_params.get('no_en_passant_embed') ] for param in params_to_init: if param is not None and param.requires_grad: nn.init.normal_(param, std=0.02) def forward(self, fen: List[str], repetitions: torch.Tensor, return_attention: bool=False) -> torch.Tensor: x = self.fen_tokenizer(fen,repetitions) # (B,73,D), pos embed are added here bs = x.shape[0] x = torch.cat([x,self.act_token.expand(bs,-1,-1),self.val_token.expand(bs,-1,-1)],dim=1) # (B,75,D) attention_maps = [] if return_attention else None for block in self.blocks: if return_attention: x, attn = block(x, return_attention=True) attention_maps.append(attn) else: x = block(x) x = self.final_norm(x) act = x[:,-2,:] val = x[:,-1,:] act_logits = self.act_proj(act) # (B,1969) val = self.val_proj(val) # (B,1) if return_attention: return act_logits, val.squeeze(1), attention_maps else: return act_logits, val.squeeze(1) def load_model(ckpt_path): checkpoint = torch.load(ckpt_path) model_config = checkpoint["model_config"] model = ChessFormerModel(**model_config) model.load_state_dict(checkpoint["model_state_dict"]) return model if __name__ == "__main__": checkpoint = torch.load("./ckpts/chessformer-sl_13.pth",map_location=torch.device("cpu")) model = ChessFormerModel(**checkpoint["config"]) model.load_state_dict(checkpoint["model_state_dict"]) model.push_to_hub("kaupane/ChessFormer-SL")