| import math |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
| from transformers import PreTrainedModel |
|
|
| |
| try: |
| from .configuration_dumbc import DumbcConfig |
| except ImportError: |
| from configuration_dumbc import DumbcConfig |
|
|
| class DumbChessRetina(nn.Module): |
| def __init__(self, dim=192): |
| super().__init__() |
| self.piece_embed = nn.Embedding(14, 32) |
| self.conv_rank_file = nn.Conv1d(32, 32, kernel_size=3, padding=1, groups=32) |
| self.conv_diag = nn.Conv1d(32, 32, kernel_size=3, padding=1, groups=32) |
| self.conv_antidiag = nn.Conv1d(32, 32, kernel_size=3, padding=1, groups=32) |
| self.conv_knight = nn.Conv1d(32, 32, kernel_size=3, padding=1, groups=32) |
| self.tension_mlp = nn.Sequential( |
| nn.Linear(32 * 4 + 1, 64), |
| nn.GELU(), |
| nn.Linear(64, dim) |
| ) |
|
|
| def forward(self, board_state, material_weights): |
| x = self.piece_embed(board_state) |
| x_t = x.transpose(1, 2) |
| f1 = self.conv_rank_file(x_t) |
| f2 = self.conv_diag(x_t) |
| f3 = self.conv_antidiag(x_t) |
| f4 = self.conv_knight(x_t) |
| f_all = torch.cat([f1, f2, f3, f4], dim=1).transpose(1, 2) |
| tension_in = torch.cat([f_all, material_weights.unsqueeze(-1)], dim=-1) |
| return self.tension_mlp(tension_in) |
|
|
| class DumbAttention(nn.Module): |
| def __init__(self, dim=192, heads=8, bottleneck=32): |
| super().__init__() |
| self.dim = dim |
| self.heads = heads |
| self.head_dim = dim // heads |
| self.qkv_proj = nn.Linear(dim, dim * 3, bias=False) |
| self.out_proj = nn.Linear(dim, dim, bias=False) |
| self.hadamard_mlp = nn.Sequential( |
| nn.Linear(dim, bottleneck), |
| nn.SiLU(), |
| nn.Linear(bottleneck, 1) |
| ) |
| self.trinity_g = nn.Linear(dim, 16, bias=False) |
| self.temp_mlp = nn.Sequential( |
| nn.Linear(dim, 16), |
| nn.SiLU(), |
| nn.Linear(16, 1) |
| ) |
| self.threat_weight = nn.Parameter(torch.ones(1) * 0.5) |
|
|
| def forward(self, x, mat_diff_matrix): |
| B, N, C = x.shape |
| q, k, v = self.qkv_proj(x).chunk(3, dim=-1) |
| q_h = q.view(B, N, self.heads, self.head_dim).transpose(1, 2) |
| k_h = k.view(B, N, self.heads, self.head_dim).transpose(1, 2) |
| v_h = v.view(B, N, self.heads, self.head_dim).transpose(1, 2) |
| S_base = (q_h @ k_h.transpose(-2, -1)) / math.sqrt(self.head_dim) |
| |
| q_k_hadamard = q.unsqueeze(2) * k.unsqueeze(1) |
| S_tensor = self.hadamard_mlp(q_k_hadamard).squeeze(-1).unsqueeze(1) |
| |
| g_q = torch.sigmoid(self.trinity_g(q)) |
| g_k = torch.sigmoid(self.trinity_g(k)) |
| S_trinity = (g_q @ g_k.transpose(-2, -1)).unsqueeze(1) |
| |
| B_material = F.relu(mat_diff_matrix).unsqueeze(1) * self.threat_weight |
| |
| tau = torch.sigmoid(self.temp_mlp(x.mean(dim=1))) * 0.5 + 0.75 |
| tau = tau.unsqueeze(-1).unsqueeze(-1) |
| |
| S_total = (S_base + S_tensor + S_trinity + B_material) / tau |
| A = F.softmax(S_total, dim=-1) |
| out = (A @ v_h).transpose(1, 2).reshape(B, N, C) |
| return self.out_proj(out) |
|
|
| class DumbFractalFFN(nn.Module): |
| def __init__(self, dim=192, hidden_dim=288): |
| super().__init__() |
| self.w1 = nn.Linear(dim, hidden_dim, bias=False) |
| self.w2 = nn.Linear(dim, hidden_dim, bias=False) |
| self.w3 = nn.Linear(hidden_dim, dim, bias=False) |
|
|
| def forward(self, x): |
| h1 = F.silu(self.w1(x)) |
| h2 = self.w2(x) |
| chunk_size = h2.shape[-1] // 2 |
| h2_a, h2_b = torch.split(h2, chunk_size, dim=-1) |
| fractal_interaction = torch.cat([h2_a * h2_b, h2_b**2], dim=-1) |
| return self.w3(h1 * fractal_interaction) |
|
|
| class DumbBlock(nn.Module): |
| def __init__(self, dim=192): |
| super().__init__() |
| self.norm1 = nn.LayerNorm(dim) |
| self.attn = DumbAttention(dim=dim) |
| self.norm2 = nn.LayerNorm(dim) |
| self.ffn = DumbFractalFFN(dim=dim) |
| self.gate = nn.Parameter(torch.ones(1) * 0.1) |
|
|
| def forward(self, x, mat_diff_matrix): |
| x = x + self.gate * self.attn(self.norm1(x), mat_diff_matrix) |
| x = x + self.gate * self.ffn(self.norm2(x)) |
| return x |
|
|
| class DumbcPreTrainedModel(PreTrainedModel): |
| config_class = DumbcConfig |
| base_model_prefix = "dumbc" |
|
|
| def _init_weights(self, module): |
| if isinstance(module, (nn.Linear, nn.Conv1d)): |
| module.weight.data.normal_(mean=0.0, std=self.config.initializer_range) |
| if module.bias is not None: |
| module.bias.data.zero_() |
|
|
| class DumbcModel(DumbcPreTrainedModel): |
| def __init__(self, config): |
| super().__init__(config) |
| self.config = config |
| self.retina = DumbChessRetina(dim=config.dim) |
| self.layers = nn.ModuleList([DumbBlock(dim=config.dim) for _ in range(config.num_unique_layers)]) |
| self.step_embed = nn.Parameter(torch.randn(config.num_loops, 1, 1, config.dim) * 0.02) |
| |
| self.from_head = nn.Linear(config.dim, 64) |
| self.to_head = nn.Linear(config.dim, 64) |
| self.value_head = nn.Sequential( |
| nn.Linear(config.dim, 64), |
| nn.GELU(), |
| nn.Linear(64, 3) |
| ) |
| self.post_init() |
|
|
| def forward(self, board_state, mat_diff_matrix, material_weights, **kwargs): |
| x = self.retina(board_state, material_weights) |
| for loop_idx in range(self.config.num_loops): |
| x = x + self.step_embed[loop_idx] |
| for layer in self.layers: |
| x = layer(x, mat_diff_matrix) |
| |
| from_logits = self.from_head(x) |
| to_logits = self.to_head(x) |
| policy_matrix = torch.bmm(from_logits, to_logits.transpose(1, 2)) |
| |
| global_pool = x.mean(dim=1) |
| value_logits = self.value_head(global_pool) |
| |
| return {"policy_matrix": policy_matrix, "value_logits": value_logits} |
|
|
