# ===================================================================== # modeling_chessnet.py — "Monkey" chess policy/value network # Hugging Face custom-code modeling file. # # State-dict keys match the Monkey Trainer v5 checkpoints exactly # (stem.* / blocks.* / policy_head.* / value_head.*), so trainer # weights load with zero key mismatch. # ===================================================================== from __future__ import annotations import json import os from dataclasses import dataclass from typing import List, Optional, Tuple import numpy as np import torch import torch.nn as nn import torch.nn.functional as F import chess from transformers import PreTrainedModel, PretrainedConfig from transformers.utils import ModelOutput # --------------------------------------------------------------------- # Constants (MUST stay identical to the trainer) # --------------------------------------------------------------------- INPUT_PLANES = 19 BOARD_SIZE = 8 POLICY_SIZE = 20480 # 64 * 64 * 5 PROMO_TO_IDX = {None: 0, chess.KNIGHT: 1, chess.BISHOP: 2, chess.ROOK: 3, chess.QUEEN: 4} IDX_TO_PROMO = {v: k for k, v in PROMO_TO_IDX.items()} # --------------------------------------------------------------------- # Config # --------------------------------------------------------------------- class ChessNetConfig(PretrainedConfig): model_type = "chessnet" def __init__( self, input_planes: int = INPUT_PLANES, board_size: int = BOARD_SIZE, num_blocks: int = 10, num_channels: int = 128, policy_channels: int = 32, value_channels: int = 8, policy_size: int = POLICY_SIZE, **kwargs, ): super().__init__(**kwargs) self.input_planes = input_planes self.board_size = board_size self.num_blocks = num_blocks self.num_channels = num_channels self.policy_channels = policy_channels self.value_channels = value_channels self.policy_size = policy_size # --------------------------------------------------------------------- # Board <-> Tensor <-> Move utilities # --------------------------------------------------------------------- def encode_board(board: chess.Board) -> np.ndarray: """Encode a chess.Board into a (19, 8, 8) float32 array.""" planes = np.zeros((INPUT_PLANES, 64), dtype=np.float32) for sq, piece in board.piece_map().items(): planes[piece.piece_type - 1 + (0 if piece.color == chess.WHITE else 6), sq] = 1.0 planes[12].fill(1.0 if board.turn == chess.WHITE else 0.0) planes[13].fill(1.0 if board.has_kingside_castling_rights(chess.WHITE) else 0.0) planes[14].fill(1.0 if board.has_queenside_castling_rights(chess.WHITE) else 0.0) planes[15].fill(1.0 if board.has_kingside_castling_rights(chess.BLACK) else 0.0) planes[16].fill(1.0 if board.has_queenside_castling_rights(chess.BLACK) else 0.0) if board.ep_square is not None: planes[17, board.ep_square] = 1.0 planes[18].fill(min(board.halfmove_clock / 100.0, 1.0)) return planes.reshape(INPUT_PLANES, BOARD_SIZE, BOARD_SIZE) def tensor_to_board(tensor) -> chess.Board: """Rebuild a chess.Board from a (19, 64) or (19, 8, 8) tensor/array.""" t = np.asarray(tensor).reshape(INPUT_PLANES, 64) board = chess.Board.empty() for sq in range(64): for p in range(6): if t[p, sq] == 1.0: board.set_piece_at(sq, chess.Piece(p + 1, chess.WHITE)) if t[p + 6, sq] == 1.0: board.set_piece_at(sq, chess.Piece(p + 1, chess.BLACK)) board.turn = chess.WHITE if t[12, 0] == 1.0 else chess.BLACK rights = 0 if t[13, 0] == 1.0: rights |= chess.BB_H1 if t[14, 0] == 1.0: rights |= chess.BB_A1 if t[15, 0] == 1.0: rights |= chess.BB_H8 if t[16, 0] == 1.0: rights |= chess.BB_A8 board.castling_rights = rights ep = int(np.argmax(t[17])) board.ep_square = ep if t[17, ep] == 1.0 else None return board def get_move_index(move: chess.Move) -> int: """(from_square * 64 + to_square) * 5 + promotion_type -> 0..20479""" return (move.from_square * 64 + move.to_square) * 5 + PROMO_TO_IDX[move.promotion] def index_to_move(idx: int) -> chess.Move: """Flat index -> chess.Move (legality NOT checked).""" from_sq = (idx // 5) // 64 to_sq = (idx // 5) % 64 return chess.Move(from_sq, to_sq, promotion=IDX_TO_PROMO[idx % 5]) # --------------------------------------------------------------------- # Model # --------------------------------------------------------------------- @dataclass class ChessNetOutput(ModelOutput): policy_logits: torch.FloatTensor = None value: torch.FloatTensor = None class ResidualBlock(nn.Module): def __init__(self, c: int): super().__init__() self.conv1 = nn.Conv2d(c, c, 3, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(c) self.conv2 = nn.Conv2d(c, c, 3, padding=1, bias=False) self.bn2 = nn.BatchNorm2d(c) def forward(self, x): return F.relu(self.bn2(self.conv2(F.relu(self.bn1(self.conv1(x))))) + x) class ChessNet(PreTrainedModel): config_class = ChessNetConfig base_model_prefix = "chessnet" supports_gradient_checkpointing = False _no_split_modules = ["ResidualBlock"] def __init__(self, config: ChessNetConfig): super().__init__(config) inc, nc = config.input_planes, config.num_channels pc, vc, bs = config.policy_channels, config.value_channels, config.board_size self.stem = nn.Sequential( nn.Conv2d(inc, nc, 3, padding=1, bias=False), nn.BatchNorm2d(nc), nn.ReLU(), ) self.blocks = nn.Sequential(*[ResidualBlock(nc) for _ in range(config.num_blocks)]) self.policy_head = nn.Sequential( nn.Conv2d(nc, pc, 1, bias=False), nn.BatchNorm2d(pc), nn.ReLU(), nn.Flatten(), nn.Linear(pc * bs * bs, config.policy_size), ) self.value_head = nn.Sequential( nn.Conv2d(nc, vc, 1, bias=False), nn.BatchNorm2d(vc), nn.ReLU(), nn.Flatten(), nn.Linear(vc * bs * bs, 64), nn.ReLU(), nn.Linear(64, 1), nn.Tanh(), ) self.post_init() def _init_weights(self, module): pass # keep PyTorch default init (identical to trainer) def forward(self, x, **kwargs) -> ChessNetOutput: h = self.blocks(self.stem(x)) return ChessNetOutput(policy_logits=self.policy_head(h), value=self.value_head(h)) # --------------------------------------------------------------------- # Player wrapper (inference / arena / Colab matches) # --------------------------------------------------------------------- class MonkeyPlayer: """Convenience wrapper: legal-move masking + argmax/temperature play.""" def __init__(self, model: ChessNet, device: Optional[str] = None): self.device = device or ("cuda" if torch.cuda.is_available() else "cpu") self.model = model.to(self.device).eval() @classmethod def from_pretrained(cls, repo_or_dir: str, token=None, device=None, **kw): try: model = ChessNet.from_pretrained(repo_or_dir, token=token, trust_remote_code=True, **kw) except Exception: # Fallback: manual load (works even if config.json lacks auto_map) from huggingface_hub import hf_hub_download from safetensors.torch import load_file if os.path.isdir(repo_or_dir): cfg_path = os.path.join(repo_or_dir, "config.json") w_path = os.path.join(repo_or_dir, "model.safetensors") else: cfg_path = hf_hub_download(repo_or_dir, "config.json", token=token) w_path = hf_hub_download(repo_or_dir, "model.safetensors", token=token) with open(cfg_path) as f: cfg = json.load(f) model = ChessNet(ChessNetConfig(**cfg)) model.load_state_dict(load_file(w_path)) return cls(model, device=device) @torch.no_grad() def choose_move(self, board: chess.Board, temperature: float = 0.0): """Returns (move, confidence, value).""" legal = list(board.legal_moves) if not legal: return None, 0.0, 0.0 x = torch.tensor(encode_board(board), dtype=torch.float32, device=self.device).unsqueeze(0) out = self.model(x) logits = out.policy_logits[0].float() idx = torch.tensor([get_move_index(m) for m in legal], device=logits.device) masked = logits[idx] if temperature > 0: probs = torch.softmax(masked / temperature, dim=-1) k = int(torch.multinomial(probs, 1).item()) else: probs = torch.softmax(masked, dim=-1) k = int(torch.argmax(probs).item()) return legal[k], float(probs[k]), float(out.value[0][0]) @torch.no_grad() def evaluate(self, board: chess.Board) -> float: x = torch.tensor(encode_board(board), dtype=torch.float32, device=self.device).unsqueeze(0) return float(self.model(x).value[0][0]) if __name__ == "__main__": cfg = ChessNetConfig() m = ChessNet(cfg) b = chess.Board() mv, conf, val = MonkeyPlayer(m, device="cpu").choose_move(b) print(f"params={sum(p.numel() for p in m.parameters()):,} | smoke move={mv.uci()} | ok")