symbio_model.py

"""SymbioGPT — Multi-organelle GPT with learned per-channel gating.

4 organelles: CausalConv + Monarch + LongConv + CausalSelfAttention
fused via OrganelleGate with learnable temperature.

Architecture: RoPE, RMSNorm, SwiGLU, SkipGate, weight-tied output.
"""
import math
from dataclasses import dataclass
from typing import List, Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F


# ═══════════════════════════════════════════════════════════════════
# Configuration
# ═══════════════════════════════════════════════════════════════════


@dataclass
class SymbioConfig:
    d_model: int = 320
    n_layers: int = 8
    n_heads: int = 5
    head_dim: int = 64
    ffn_mult: int = 4
    dropout: float = 0.0
    context_length: int = 256
    vocab_size: int = 2000
    weight_tying: bool = True
    organelles: Tuple[str, ...] = ("causal_conv", "monarch", "long_conv", "attention")
    conv_kernel_size: int = 4
    n_monarch_heads: int = 1
    gate_temperature_init: float = 1.0
    free_energy_beta: float = 0.001
    per_layer_organelles: Optional[List[Tuple[str, ...]]] = None

    @property
    def p(self) -> int:
        return int(math.isqrt(self.context_length))

    @property
    def n_organelles(self) -> int:
        return len(self.organelles)


# ═══════════════════════════════════════════════════════════════════
# Building blocks
# ═══════════════════════════════════════════════════════════════════


class RMSNorm(nn.Module):
    def __init__(self, dim: int, eps: float = 1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        rms = torch.sqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
        return x / rms * self.weight


class RotaryEmbedding(nn.Module):
    def __init__(self, dim: int, max_seq_len: int = 2048):
        super().__init__()
        freqs = 1.0 / (10000.0 ** (torch.arange(0, dim, 2).float() / dim))
        positions = torch.arange(max_seq_len).float()
        angles = torch.outer(positions, freqs)
        self.register_buffer("cos_cache", angles.cos())
        self.register_buffer("sin_cache", angles.sin())

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        seq_len = x.size(2)
        half = x.size(-1) // 2
        x1, x2 = x[..., :half], x[..., half:]
        cos = self.cos_cache[:seq_len, :half].unsqueeze(0).unsqueeze(0)
        sin = self.sin_cache[:seq_len, :half].unsqueeze(0).unsqueeze(0)
        return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)


class SwiGLU(nn.Module):
    def __init__(self, d_model: int, ffn_mult: int = 4):
        super().__init__()
        raw_hidden = 2 * d_model * ffn_mult // 3
        hidden_dim = max(64, (raw_hidden // 64) * 64)
        self.w1 = nn.Linear(d_model, hidden_dim, bias=False)
        self.v = nn.Linear(d_model, hidden_dim, bias=False)
        self.w2 = nn.Linear(hidden_dim, d_model, bias=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.w2(F.silu(self.w1(x)) * self.v(x))


class CausalSelfAttention(nn.Module):
    def __init__(self, d_model: int, n_heads: int, head_dim: int, dropout: float = 0.0):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = head_dim
        total_dim = n_heads * head_dim
        self.wq = nn.Linear(d_model, total_dim, bias=False)
        self.wk = nn.Linear(d_model, total_dim, bias=False)
        self.wv = nn.Linear(d_model, total_dim, bias=False)
        self.wo = nn.Linear(total_dim, d_model, bias=False)
        self.attn_dropout = nn.Dropout(dropout) if dropout > 0.0 else nn.Identity()

    def forward(self, x, rope, mask=None):
        B, T, D = x.shape
        H, HD = self.n_heads, self.head_dim
        q = self.wq(x).view(B, T, H, HD).transpose(1, 2)
        k = self.wk(x).view(B, T, H, HD).transpose(1, 2)
        v = self.wv(x).view(B, T, H, HD).transpose(1, 2)
        q = rope(q)
        k = rope(k)
        scale = 1.0 / math.sqrt(HD)
        attn = torch.matmul(q, k.transpose(-2, -1)) * scale
        if mask is not None:
            attn = attn + mask
        attn = F.softmax(attn, dim=-1)
        attn = self.attn_dropout(attn)
        out = torch.matmul(attn, v)
        out = out.transpose(1, 2).contiguous().view(B, T, H * HD)
        return self.wo(out)


# ═══════════════════════════════════════════════════════════════════
# Organelles
# ═══════════════════════════════════════════════════════════════════


class CausalDepthwiseConv1d(nn.Module):
    def __init__(self, channels: int, kernel_size: int = 4):
        super().__init__()
        self.channels = channels
        self.kernel_size = kernel_size
        self.weight = nn.Parameter(torch.empty(channels, 1, kernel_size))
        nn.init.normal_(self.weight, std=math.sqrt(1.0 / kernel_size))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, D = x.shape
        x_t = x.transpose(1, 2)
        x_padded = F.pad(x_t, (self.kernel_size - 1, 0))
        out = F.conv1d(x_padded, self.weight, groups=D)
        return out.transpose(1, 2)


class MonarchMatrix(nn.Module):
    def __init__(self, seq_len: int):
        super().__init__()
        p = int(math.isqrt(seq_len))
        self.seq_len = seq_len
        self.p = p
        scale = math.sqrt(2.0 / (p + p))
        self.L1 = nn.Parameter(torch.randn(p, p, p) * scale)
        self.L2 = nn.Parameter(torch.randn(p, p, p) * scale)

    @staticmethod
    def _julia_batched_mul(A, B):
        return torch.bmm(A.permute(2, 0, 1), B.permute(2, 0, 1)).permute(1, 2, 0)

    def realize(self):
        p = self.p
        T = self.seq_len
        I_T = torch.eye(T, device=self.L1.device, dtype=self.L1.dtype)
        x = I_T.reshape(p, p, T)
        x = x.permute(0, 2, 1)
        x = self._julia_batched_mul(self.L2, x)
        x = x.permute(0, 2, 1)
        x = x.permute(1, 0, 2)
        x = x.permute(0, 2, 1)
        x = self._julia_batched_mul(self.L1, x)
        x = x.permute(0, 2, 1)
        x = x.permute(1, 0, 2)
        return x.reshape(T, T)

    def forward(self, x, causal_mask=None):
        B, T, D_head = x.shape
        M = self.realize()[:T, :T]
        if causal_mask is not None:
            M = M * causal_mask[:T, :T]
        x_flat = x.permute(1, 0, 2).reshape(T, B * D_head)
        y_flat = M @ x_flat
        return y_flat.reshape(T, B, D_head).permute(1, 0, 2)


class LongConv(nn.Module):
    def __init__(self, channels: int, seq_len: int):
        super().__init__()
        self.channels = channels
        self.seq_len = seq_len
        self.kernel = nn.Parameter(torch.empty(channels, 1, seq_len))
        self._init_weights()

    def _init_weights(self):
        scale = math.sqrt(1.0 / self.seq_len)
        nn.init.normal_(self.kernel, mean=0.0, std=scale)
        with torch.no_grad():
            decay = torch.exp(-0.1 * torch.arange(self.seq_len, dtype=torch.float32))
            self.kernel.mul_(decay.unsqueeze(0).unsqueeze(0))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, T, D = x.shape
        K = self.seq_len
        x_t = x.transpose(1, 2)
        x_padded = F.pad(x_t, (K - 1, 0))
        out = F.conv1d(x_padded, self.kernel, groups=D)
        return out.transpose(1, 2)


# ═══════════════════════════════════════════════════════════════════
# Gate and mixer
# ═══════════════════════════════════════════════════════════════════


class OrganelleGate(nn.Module):
    def __init__(self, dim: int, n_organelles: int, temperature_init: float = 1.0):
        super().__init__()
        self.n_organelles = n_organelles
        self.logits = nn.Parameter(torch.zeros(n_organelles, dim))
        self.temperature = nn.Parameter(torch.tensor([temperature_init]))

    def forward(self, organelle_outputs, organelle_mask=None):
        logits = self.logits
        if organelle_mask is not None:
            mask_additive = torch.zeros_like(logits)
            for i in range(self.n_organelles):
                if not organelle_mask[i]:
                    mask_additive[i, :] = -1e10
            logits = logits + mask_additive
        tau = self.temperature.clamp(min=0.01)
        weights = F.softmax(logits / tau, dim=0)
        out = torch.zeros_like(organelle_outputs[0])
        for i in range(self.n_organelles):
            w = weights[i].unsqueeze(0).unsqueeze(0)
            out = out + w * organelle_outputs[i]
        return out


class SkipGate(nn.Module):
    def __init__(self):
        super().__init__()
        self.scale = nn.Parameter(torch.ones(1))

    def forward(self, x):
        return self.scale * x


class SymbioSequenceMixer(nn.Module):
    def __init__(self, config: SymbioConfig):
        super().__init__()
        self.config = config
        d = config.d_model
        T = config.context_length
        self.organelle_names = list(config.organelles)
        self.organelle_modules = nn.ModuleDict()
        for name in self.organelle_names:
            if name == "causal_conv":
                self.organelle_modules[name] = CausalDepthwiseConv1d(d, config.conv_kernel_size)
            elif name == "monarch":
                self.organelle_modules[name] = nn.ModuleList(
                    [MonarchMatrix(T) for _ in range(config.n_monarch_heads)]
                )
            elif name == "long_conv":
                self.organelle_modules[name] = LongConv(d, T)
            elif name == "attention":
                self.organelle_modules[name] = CausalSelfAttention(
                    d, config.n_heads, config.head_dim, config.dropout
                )
        self.gate = OrganelleGate(d, len(self.organelle_names), config.gate_temperature_init)
        if "monarch" in self.organelle_names:
            self.register_buffer("monarch_causal_mask", torch.tril(torch.ones(T, T)))

    def forward(self, x, rope, attn_mask=None, organelle_mask=None):
        B, T, D = x.shape
        outputs = []
        for name in self.organelle_names:
            if name == "causal_conv":
                out = self.organelle_modules[name](x)
            elif name == "monarch":
                heads = self.organelle_modules[name]
                n_mh = len(heads)
                hd = D // n_mh
                slices = []
                for i, monarch in enumerate(heads):
                    x_slice = x[:, :, i * hd : (i + 1) * hd]
                    y_slice = monarch(x_slice, self.monarch_causal_mask)
                    slices.append(y_slice)
                out = torch.cat(slices, dim=-1)
            elif name == "long_conv":
                out = self.organelle_modules[name](x)
            elif name == "attention":
                out = self.organelle_modules[name](x, rope, attn_mask)
            outputs.append(out)
        return self.gate(tuple(outputs), organelle_mask)


# ═══════════════════════════════════════════════════════════════════
# SymbioBlock and SymbioGPT
# ═══════════════════════════════════════════════════════════════════


class SymbioBlock(nn.Module):
    def __init__(self, config: SymbioConfig, layer_organelles=None):
        super().__init__()
        d = config.d_model
        if layer_organelles is not None:
            from dataclasses import replace
            layer_config = replace(config, organelles=layer_organelles)
        else:
            layer_config = config
        self.ln1 = RMSNorm(d)
        self.seq_mixer = SymbioSequenceMixer(layer_config)
        self.skip1 = SkipGate()
        self.ln2 = RMSNorm(d)
        self.ffn = SwiGLU(d, config.ffn_mult)
        self.skip2 = SkipGate()

    def forward(self, x, rope, attn_mask=None, organelle_mask=None):
        x = x + self.skip1(self.seq_mixer(self.ln1(x), rope, attn_mask, organelle_mask))
        x = x + self.skip2(self.ffn(self.ln2(x)))
        return x


class SymbioGPT(nn.Module):
    def __init__(self, config: SymbioConfig):
        super().__init__()
        self.config = config
        self.tok_emb = nn.Embedding(config.vocab_size, config.d_model)
        self.rope = RotaryEmbedding(config.head_dim, config.context_length)
        blocks = []
        for i in range(config.n_layers):
            layer_org = None
            if config.per_layer_organelles is not None:
                layer_org = config.per_layer_organelles[i]
            blocks.append(SymbioBlock(config, layer_org))
        self.blocks = nn.ModuleList(blocks)
        self.ln_f = RMSNorm(config.d_model)
        if config.weight_tying:
            self.head = None
        else:
            self.head = nn.Linear(config.d_model, config.vocab_size, bias=False)

    def forward(self, input_ids, organelle_mask=None):
        B, T = input_ids.shape
        x = self.tok_emb(input_ids)
        attn_mask = torch.triu(
            torch.full((T, T), float("-inf"), device=x.device, dtype=x.dtype), diagonal=1
        )
        for block in self.blocks:
            x = block(x, self.rope, attn_mask, organelle_mask)
        x = self.ln_f(x)
        if self.head is not None:
            return self.head(x)
        return F.linear(x, self.tok_emb.weight)