| """Hugging Face configuration for the chess transition policy.""" |
|
|
| from __future__ import annotations |
|
|
| from transformers import PretrainedConfig |
|
|
|
|
| class ChessPolicyConfig(PretrainedConfig): |
| """Serializable configuration for ``ChessTransitionPolicy``.""" |
|
|
| model_type = "chess_transition_policy" |
|
|
| def __init__( |
| self, |
| *, |
| board_squares: int = 64, |
| piece_states: int = 13, |
| d_model: int = 384, |
| n_heads: int = 6, |
| board_layers: int = 6, |
| candidate_layers: int = 12, |
| swiglu_hidden: int = 1040, |
| dropout: float = 0.0, |
| activation_checkpointing: bool = False, |
| **kwargs, |
| ) -> None: |
| super().__init__(**kwargs) |
| if board_squares != 64: |
| raise ValueError("ChessPolicyConfig requires exactly 64 board squares") |
| if d_model % n_heads != 0: |
| raise ValueError("d_model must be divisible by n_heads") |
| if board_layers < 1 or candidate_layers < 1: |
| raise ValueError("Both transformer stacks need at least one layer") |
| self.board_squares = board_squares |
| self.piece_states = piece_states |
| self.d_model = d_model |
| self.n_heads = n_heads |
| self.board_layers = board_layers |
| self.candidate_layers = candidate_layers |
| self.swiglu_hidden = swiglu_hidden |
| self.dropout = dropout |
| self.activation_checkpointing = activation_checkpointing |
|
|
| @property |
| def head_dim(self) -> int: |
| return self.d_model // self.n_heads |
|
|
