modeling_chessnet.py

# =====================================================================
# 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")