Upload modeling_chessnet.py

modeling_chessnet.py

@@ -0,0 +1,242 @@
+# =====================================================================
+# 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")