src/bihmoe/models/structured.py

from __future__ import annotations
import torch
import torch.nn as nn
from .blocks import DenseFFN, MoEFFN, TransformerBlock

class StructuredBiHMoE(nn.Module):
    """
    V2.2: fixes optic chiasm routing by introducing an explicit QUERY token type
    (last non-pad token), and implements a soft chiasm split (alpha-left) so
    query/mid/val flow primarily into the right (global) hemisphere.

    Key idea:
      xl = x * alpha_left + type_L(type)
      xr = x * (1 - alpha_left) + type_R(type)
    """
    def __init__(
        self,
        vocab_size: int,
        d_model: int,
        n_heads: int,
        n_layers_stem: int,
        n_layers_hemi: int,
        d_ff_dense: int,
        d_ff_expert: int,
        n_experts: int,
        top_k: int,
        workspace_tokens: int = 4,
        reconcile_every: int = 2,
        dropout: float = 0.0,
        pad_id: int = 0,
        fuse: str = "mean",
        moe_warmup_steps: int = 0,

        # asymmetry
        left_local_window: int = 0,
        right_noise_std: float = 0.0,

        # callosum
        callosum_competitive: bool = True,
        callosum_tau: float = 1.0,

        # optic chiasm (type + soft split)
        chiasm_enabled: bool = False,
        chiasm_mode: str = "bias",  # "bias" or "split"
        noise_vocab: int = 16,
        key_start: int = -1, key_count: int = 0,
        mid_start: int = -1, mid_count: int = 0,
        val_start: int = -1, val_count: int = 0,
        chiasm_alpha_key: float = 0.9,     # fraction of x sent to LEFT for KEY
        chiasm_alpha_mid: float = 0.1,     # fraction to LEFT for MID
        chiasm_alpha_val: float = 0.1,     # fraction to LEFT for VAL
        chiasm_alpha_query: float = 0.1,   # fraction to LEFT for QUERY (i.e., mostly RIGHT)
        chiasm_alpha_other: float = 0.5,
        output_gate: bool = True,
        output_tau: float = 1.0,
    ):
        super().__init__()
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.workspace_tokens = workspace_tokens
        self.reconcile_every = max(1, reconcile_every)
        self.pad_id = int(pad_id)
        self.fuse = fuse
        self.moe_warmup_steps = int(moe_warmup_steps)

        self.left_local_window = int(left_local_window)
        self.right_noise_std = float(right_noise_std)

        self.callosum_competitive = bool(callosum_competitive)
        self.callosum_tau = float(callosum_tau)

        self.chiasm_enabled = bool(chiasm_enabled)
        self.chiasm_mode = str(chiasm_mode)
        self.noise_vocab = int(noise_vocab)
        self.key_start, self.key_count = int(key_start), int(key_count)
        self.mid_start, self.mid_count = int(mid_start), int(mid_count)
        self.val_start, self.val_count = int(val_start), int(val_count)

        self.chiasm_alpha_key = float(chiasm_alpha_key)
        self.chiasm_alpha_mid = float(chiasm_alpha_mid)
        self.chiasm_alpha_val = float(chiasm_alpha_val)
        self.chiasm_alpha_query = float(chiasm_alpha_query)
        self.chiasm_alpha_other = float(chiasm_alpha_other)
        self.output_gate = bool(output_gate)
        self.output_tau = float(output_tau)

        self.tok = nn.Embedding(vocab_size, d_model)

        # types:
        # 0=PAD, 1=SPECIAL(BOS/SEP), 2=NOISE, 3=KEY, 4=MID, 5=VAL, 6=QUERY, 7=OTHER
        self.n_types = 8
        if self.chiasm_enabled:
            self.type_L = nn.Embedding(self.n_types, d_model)
            self.type_R = nn.Embedding(self.n_types, d_model)

        self.stem = nn.ModuleList([
            TransformerBlock(d_model, n_heads, DenseFFN(d_model, d_ff_dense, dropout=dropout), dropout=dropout)
            for _ in range(n_layers_stem)
        ])

        def hemi_blocks():
            blocks = []
            for i in range(n_layers_hemi):
                if (i % 2) == 1:
                    ffn = MoEFFN(d_model, d_ff_expert, n_experts=n_experts, top_k=top_k, dropout=dropout, council_mlp=True)
                else:
                    ffn = DenseFFN(d_model, d_ff_dense, dropout=dropout)
                blocks.append(TransformerBlock(d_model, n_heads, ffn, dropout=dropout))
            return nn.ModuleList(blocks)

        self.left = hemi_blocks()
        self.right = hemi_blocks()

        self.workspace = nn.Parameter(torch.zeros(workspace_tokens, d_model))

        self.w_to_l = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True)
        self.w_to_r = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True)
        self.write_gate = nn.Linear(d_model, 1)

        self.l_to_w = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True)
        self.r_to_w = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True)

        self.ln = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size)

    def _type_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
        B, T = input_ids.shape
        tid = torch.full_like(input_ids, 7)  # OTHER

        # PAD
        tid = torch.where(input_ids == self.pad_id, torch.zeros_like(tid), tid)
        # SPECIAL (BOS=1, SEP=2)
        tid = torch.where((input_ids == 1) | (input_ids == 2), torch.full_like(tid, 1), tid)

        # NOISE range: [3, 3+noise_vocab)
        n0 = 3
        n1 = n0 + max(0, self.noise_vocab)
        tid = torch.where((input_ids >= n0) & (input_ids < n1), torch.full_like(tid, 2), tid)

        # KEY
        if self.key_count > 0 and self.key_start >= 0:
            ks, ke = self.key_start, self.key_start + self.key_count
            tid = torch.where((input_ids >= ks) & (input_ids < ke), torch.full_like(tid, 3), tid)

        # MID
        if self.mid_count > 0 and self.mid_start >= 0:
            ms, me = self.mid_start, self.mid_start + self.mid_count
            tid = torch.where((input_ids >= ms) & (input_ids < me), torch.full_like(tid, 4), tid)

        # VAL
        if self.val_count > 0 and self.val_start >= 0:
            vs, ve = self.val_start, self.val_start + self.val_count
            tid = torch.where((input_ids >= vs) & (input_ids < ve), torch.full_like(tid, 5), tid)

        # QUERY = last non-pad token (override)
        nonpad = (input_ids != self.pad_id).to(torch.long)
        last = nonpad.sum(dim=1).clamp_min(1) - 1  # (B,)
        b = torch.arange(B, device=input_ids.device)
        tid[b, last] = 6

        return tid

    def _last_nonpad(self, input_ids: torch.Tensor, h: torch.Tensor) -> torch.Tensor:
        nonpad = (input_ids != self.pad_id).to(torch.long)
        lengths = nonpad.sum(dim=1).clamp_min(1) - 1
        idx = lengths.view(-1, 1, 1).expand(-1, 1, h.size(-1))
        return h.gather(dim=1, index=idx).squeeze(1)

    def _band_mask(self, T: int, window: int, device) -> torch.Tensor:
        i = torch.arange(T, device=device).view(T, 1)
        j = torch.arange(T, device=device).view(1, T)
        return (j < (i - window)) | (j > (i + window))  # True = disallow

    def forward(self, input_ids: torch.Tensor, return_aux: bool = False, global_step: int | None = None):
        B, T = input_ids.shape
        key_padding_mask = (input_ids == self.pad_id)
        x = self.tok(input_ids)

        for b in self.stem:
            x = b(x, key_padding_mask=key_padding_mask)

        # Optic chiasm (bias or split)
        if self.chiasm_enabled:
            tid = self._type_ids(input_ids)
            if self.chiasm_mode == "bias":
                # Both hemispheres see the full token embedding; type embeddings bias specialization.
                xl = x + self.type_L(tid)
                xr = x + self.type_R(tid)
            else:
                # Soft split (legacy): attenuates content; can starve the global branch on long/noisy tasks.
                alpha = torch.full((B, T, 1), self.chiasm_alpha_other, device=input_ids.device, dtype=x.dtype)
                alpha = torch.where((tid == 3).unsqueeze(-1), torch.tensor(self.chiasm_alpha_key, device=input_ids.device, dtype=x.dtype), alpha)
                alpha = torch.where((tid == 4).unsqueeze(-1), torch.tensor(self.chiasm_alpha_mid, device=input_ids.device, dtype=x.dtype), alpha)
                alpha = torch.where((tid == 5).unsqueeze(-1), torch.tensor(self.chiasm_alpha_val, device=input_ids.device, dtype=x.dtype), alpha)
                alpha = torch.where((tid == 6).unsqueeze(-1), torch.tensor(self.chiasm_alpha_query, device=input_ids.device, dtype=x.dtype), alpha)
                xl = x * alpha + self.type_L(tid)
                xr = x * (1.0 - alpha) + self.type_R(tid)
        else:
            xl, xr = x, x


        if self.training and self.right_noise_std > 0.0:
            xr = xr + (self.right_noise_std * torch.randn_like(xr))

        w = self.workspace.unsqueeze(0).expand(B, self.workspace_tokens, self.d_model)

        warmup_active = (global_step is not None and global_step < self.moe_warmup_steps)
        warmup_expert = int(global_step) if warmup_active else None

        left_attn_mask = self._band_mask(T, self.left_local_window, device=input_ids.device) if (self.left_local_window and self.left_local_window > 0) else None

        gate_L_mean = None
        gate_R_mean = None

        n = min(len(self.left), len(self.right))
        for i in range(n):
            # Left local
            if isinstance(self.left[i].ffn, MoEFFN) and warmup_active:
                xl = self.left[i](xl, attn_mask=left_attn_mask, key_padding_mask=key_padding_mask, ffn_kwargs={"warmup_expert": warmup_expert})
            else:
                xl = self.left[i](xl, attn_mask=left_attn_mask, key_padding_mask=key_padding_mask)

            # Right global
            if isinstance(self.right[i].ffn, MoEFFN) and warmup_active:
                xr = self.right[i](xr, key_padding_mask=key_padding_mask, ffn_kwargs={"warmup_expert": warmup_expert})
            else:
                xr = self.right[i](xr, key_padding_mask=key_padding_mask)

            if ((i + 1) % self.reconcile_every) == 0:
                dwL, _ = self.w_to_l(w, xl, xl, need_weights=False, key_padding_mask=key_padding_mask)
                dwR, _ = self.w_to_r(w, xr, xr, need_weights=False, key_padding_mask=key_padding_mask)

                if self.callosum_competitive:
                    sL = self.write_gate(dwL.float())
                    sR = self.write_gate(dwR.float())
                    scores = torch.cat([sL, sR], dim=-1)
                    g = torch.softmax(scores / max(1e-6, self.callosum_tau), dim=-1)
                    dw = (g[..., 0:1].to(dwL.dtype) * dwL) + (g[..., 1:2].to(dwR.dtype) * dwR)
                    gate_L_mean = float(g[..., 0].mean().item())
                    gate_R_mean = float(g[..., 1].mean().item())
                else:
                    dw = dwL + dwR

                w = w + dw

                dl, _ = self.l_to_w(xl, w, w, need_weights=False)
                dr, _ = self.r_to_w(xr, w, w, need_weights=False)
                xl = xl + dl
                xr = xr + dr

        y_merge = self.ln(w.mean(dim=1))
        logits_merge = self.head(y_merge)

        y_l_last = self.ln(self._last_nonpad(input_ids, xl))
        y_r_last = self.ln(self._last_nonpad(input_ids, xr))
        logits_l_last = self.head(y_l_last)
        logits_r_last = self.head(y_r_last)

        # Decision-level fusion
        if self.output_gate:
            # confidence = max softmax prob (computed in fp32)
            pL = torch.softmax(logits_l_last.float(), dim=-1)
            pR = torch.softmax(logits_r_last.float(), dim=-1)
            cL = pL.max(dim=-1).values  # (B,)
            cR = pR.max(dim=-1).values  # (B,)
            scores = torch.stack([cL, cR], dim=-1)  # (B,2)
            w = torch.softmax(scores / max(1e-6, self.output_tau), dim=-1)  # (B,2)
            wL = w[:, 0].to(logits_l_last.dtype).unsqueeze(-1)
            wR = w[:, 1].to(logits_r_last.dtype).unsqueeze(-1)
            logits_hemi = (wL * logits_l_last) + (wR * logits_r_last)
            logits = 0.5 * logits_merge + 0.5 * logits_hemi
        else:
            logits = (logits_merge + logits_l_last + logits_r_last) / 3.0 if self.fuse == "mean" else logits_merge


        if not return_aux:
            return logits

        pl = torch.softmax(logits_l_last.float(), dim=-1)
        pr = torch.softmax(logits_r_last.float(), dim=-1)
        kl_lr = torch.sum(pl * (pl.clamp_min(1e-9).log() - pr.clamp_min(1e-9).log()), dim=-1)
        kl_rl = torch.sum(pr * (pr.clamp_min(1e-9).log() - pl.clamp_min(1e-9).log()), dim=-1)
        sym_kl = 0.5 * (kl_lr + kl_rl)

        aux = {
            "sym_kl": sym_kl,
            "warmup_active": bool(warmup_active),
            "gate_L_mean": gate_L_mean,
            "gate_R_mean": gate_R_mean,
            "out_wL_mean": (float(w[:,0].mean().item()) if ('w' in locals()) else None),
            "out_wR_mean": (float(w[:,1].mean().item()) if ('w' in locals()) else None),
        }
        return logits, aux