model.py

"""
Cozet: Native SYNAXIM Base Model
=================================
Non-transformer architecture using Symbiotic Gate (M-matrix)
instead of self-attention. Trained from scratch.

Architecture per layer:
  RMSNorm -> Q/K/V proj -> Symbiotic Gate (M update + retrieval) -> O proj -> Residual
  RMSNorm -> SwiGLU MLP -> Residual

(c) 2026 GRRN Research. All rights reserved.
"""

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from dataclasses import dataclass
from typing import Optional, List, Tuple


@dataclass
class CozConfig:
    """Cozet model configuration."""
    hidden_size: int = 1024
    num_layers: int = 12
    num_attention_heads: int = 16
    num_kv_heads: int = 4
    intermediate_size: int = 4096
    vocab_size: int = 32000
    max_seq_len: int = 4096
    rope_theta: float = 10000.0
    rms_norm_eps: float = 1e-6
    memory_decay: float = 0.995
    tie_word_embeddings: bool = True
    
    @property
    def head_dim(self) -> int:
        return self.hidden_size // self.num_attention_heads
    
    def num_params(self) -> int:
        """Estimate total parameter count."""
        D = self.hidden_size
        V = self.vocab_size
        I = self.intermediate_size
        L = self.num_layers
        n_kv = self.num_kv_heads
        hd = self.head_dim
        
        embed = V * D
        per_layer = (
            D * D +                    # q_proj
            D * (n_kv * hd) +          # k_proj
            D * (n_kv * hd) +          # v_proj
            D * D +                    # o_proj
            D * I +                    # gate_proj
            D * I +                    # up_proj
            I * D +                    # down_proj
            2 * D +                    # norms
            2                          # gate params
        )
        lm_head = 0 if self.tie_word_embeddings else V * D
        total = embed + L * per_layer + lm_head + D
        return total


# ---- Presets ----

COZET_SMALL = CozConfig(
    hidden_size=1024, num_layers=12, num_attention_heads=16,
    num_kv_heads=4, intermediate_size=4096, vocab_size=32000,
    max_seq_len=4096, tie_word_embeddings=True,
)

COZET_MEDIUM = CozConfig(
    hidden_size=2048, num_layers=24, num_attention_heads=16,
    num_kv_heads=4, intermediate_size=8192, vocab_size=32000,
    max_seq_len=8192, rope_theta=500000.0, tie_word_embeddings=True,
)

COZET_LARGE = CozConfig(
    hidden_size=4096, num_layers=32, num_attention_heads=32,
    num_kv_heads=8, intermediate_size=14336, vocab_size=32000,
    max_seq_len=8192, rope_theta=1000000.0, tie_word_embeddings=False,
)


class RMSNorm(nn.Module):
    """Root Mean Square Layer Normalization."""
    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:
        norm = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return (x.float() * norm).type_as(x) * self.weight


class SymbioticGate(nn.Module):
    """
    The Symbiotic Gate: core SYNAXIM attention replacement.
    
    Instead of softmax(QK^T/sqrt(d)) @ V, computes:
      gate = sigmoid(mean(Q * K) * scale * gate_scale + gate_bias)
      M    = gate * decay * M + (1-gate) * outer(k_norm, v_norm)
      out  = q @ M @ W_o
    
    M is persistent state: O(1) memory, infinite context.
    Fully differentiable for training.
    """
    
    def __init__(self, config: CozConfig, layer_idx: int):
        super().__init__()
        self.D = config.hidden_size
        self.n_heads = config.num_attention_heads
        self.n_kv = config.num_kv_heads
        self.head_dim = config.head_dim
        self.decay = config.memory_decay
        self.layer_idx = layer_idx
        
        # Projections
        self.q_proj = nn.Linear(self.D, self.n_heads * self.head_dim, bias=False)
        self.k_proj = nn.Linear(self.D, self.n_kv * self.head_dim, bias=False)
        self.v_proj = nn.Linear(self.D, self.n_kv * self.head_dim, bias=False)
        self.o_proj = nn.Linear(self.n_heads * self.head_dim, self.D, bias=False)
        
        # Learnable gate parameters
        self.gate_scale = nn.Parameter(torch.ones(1))
        self.gate_bias = nn.Parameter(torch.zeros(1))
        
        # RoPE tables (precomputed, not learned)
        self._rope_built = False
    
    def _build_rope(self, max_seq: int, device: torch.device):
        """Precompute RoPE cos/sin tables."""
        if self._rope_built and max_seq <= self.rope_cos.shape[0]:
            return
        self._rope_built = False  # Rebuild if we need more positions
        half = self.head_dim // 2
        freqs = 1.0 / (10000.0 ** (torch.arange(0, half, device=device).float() / half))
        t = torch.arange(max_seq, device=device).float()
        angles = torch.outer(t, freqs)
        self.register_buffer("rope_cos", angles.cos(), persistent=False)
        self.register_buffer("rope_sin", angles.sin(), persistent=False)
        self._rope_built = True
    
    def _apply_rope(self, x: torch.Tensor, pos: int) -> torch.Tensor:
        """Apply RoPE to (n_heads, head_dim) tensor at position pos."""
        half = self.head_dim // 2
        cos = self.rope_cos[pos, :half]
        sin = self.rope_sin[pos, :half]
        x1, x2 = x[..., :half], x[..., half:]
        return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
    
    def forward(self, h: torch.Tensor, M: torch.Tensor,
                position: int) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Symbiotic Gate forward pass for a single token.
        
        Args:
            h: (D,) hidden state
            M: (D, D) persistent memory matrix
            position: token position for RoPE
        
        Returns:
            output: (D,) attention output
            M_new: (D, D) updated memory
        """
        self._build_rope(position + 1, h.device)
        
        # Project
        q = self.q_proj(h)    # (n_heads * hd,)
        k = self.k_proj(h)    # (n_kv * hd,)
        v = self.v_proj(h)    # (n_kv * hd,)
        
        # Reshape for per-head ops
        q_h = q.view(self.n_heads, self.head_dim)
        k_h = k.view(self.n_kv, self.head_dim)
        v_h = v.view(self.n_kv, self.head_dim)
        
        # RoPE
        q_h = self._apply_rope(q_h, position)
        k_h = self._apply_rope(k_h, position)
        
        # GQA: repeat KV heads
        if self.n_kv < self.n_heads:
            repeat = self.n_heads // self.n_kv
            k_h = k_h.repeat_interleave(repeat, dim=0)
            v_h = v_h.repeat_interleave(repeat, dim=0)
        
        # Gate score
        scale = 1.0 / math.sqrt(self.head_dim)
        gate_score = (q_h * k_h).sum(-1).mean() * scale
        gate_score = gate_score * self.gate_scale + self.gate_bias
        gate = torch.sigmoid(gate_score)
        
        # Flatten to D-space
        k_flat = k_h.reshape(-1)
        v_flat = v_h.reshape(-1)
        
        # Normalize
        k_norm = k_flat / (k_flat.norm() + 1e-8)
        v_norm = v_flat / (v_flat.norm() + 1e-8) * h.norm()
        
        # M-matrix update (differentiable)
        imprint = torch.outer(k_norm, v_norm)
        M_new = gate * self.decay * M + (1.0 - gate) * imprint
        
        # Retrieve
        q_flat = q_h.reshape(-1)
        output = q_flat @ M_new
        
        # Output projection
        result = self.o_proj(output)
        
        return result, M_new


class SynaxBlock(nn.Module):
    """One SYNAXIM layer: SymbioticGate + SwiGLU MLP."""
    
    def __init__(self, config: CozConfig, layer_idx: int):
        super().__init__()
        self.norm_attn = RMSNorm(config.hidden_size, config.rms_norm_eps)
        self.attn = SymbioticGate(config, layer_idx)
        self.norm_mlp = RMSNorm(config.hidden_size, config.rms_norm_eps)
        
        # SwiGLU MLP
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)
    
    def forward(self, h: torch.Tensor, M: torch.Tensor,
                position: int) -> Tuple[torch.Tensor, torch.Tensor]:
        # Attention
        h_normed = self.norm_attn(h)
        attn_out, M_new = self.attn(h_normed, M, position)
        h = h + attn_out
        
        # MLP
        h_normed = self.norm_mlp(h)
        gate_out = self.gate_proj(h_normed)
        up_out = self.up_proj(h_normed)
        mlp_out = self.down_proj(F.silu(gate_out) * up_out)
        h = h + mlp_out
        
        return h, M_new


class CozModel(nn.Module):
    """
    Cozet: Native SYNAXIM base model.
    
    Non-transformer architecture trained from scratch.
    Uses Symbiotic Gate (M-matrix) instead of self-attention.
    O(1) memory at inference. Infinite context. No KV cache.
    """
    
    def __init__(self, config: CozConfig):
        super().__init__()
        self.config = config
        self.D = config.hidden_size
        self.n_layers = config.num_layers
        
        # Embedding
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        
        # SYNAXIM layers
        self.layers = nn.ModuleList([
            SynaxBlock(config, i) for i in range(config.num_layers)
        ])
        
        # Final norm
        self.final_norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
        
        # LM head
        if not config.tie_word_embeddings:
            self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        
        # Initialize weights
        self._init_weights()
    
    def _init_weights(self):
        """SYNAXIM-specific weight initialization."""
        depth_scale = 1.0 / math.sqrt(2 * self.n_layers)
        
        for name, p in self.named_parameters():
            if p.dim() == 1:
                continue  # Norms already initialized to ones
            if "embed" in name:
                nn.init.normal_(p, std=0.02)
            elif "down_proj" in name or "o_proj" in name:
                nn.init.normal_(p, std=0.02 * depth_scale)
            elif "lm_head" in name:
                nn.init.normal_(p, std=0.02 * depth_scale)
            elif p.dim() == 2:
                nn.init.normal_(p, std=0.02)
    
    def init_m_states(self, device: torch.device) -> List[torch.Tensor]:
        """Create fresh M-matrix states for a new sequence."""
        return [torch.zeros(self.D, self.D, device=device)
                for _ in range(self.n_layers)]
    
    def forward_token(self, token_id: int, M_states: List[torch.Tensor],
                       position: int) -> Tuple[torch.Tensor, List[torch.Tensor]]:
        """
        Process one token through the full model.
        
        Returns: (logits, updated_M_states)
        """
        h = self.embed_tokens.weight[token_id]
        
        for i, layer in enumerate(self.layers):
            h, M_states[i] = layer(h, M_states[i], position)
        
        h = self.final_norm(h)
        
        if self.config.tie_word_embeddings:
            logits = h @ self.embed_tokens.weight.T
        else:
            logits = self.lm_head(h)
        
        return logits, M_states
    
    def forward_sequence(self, token_ids: torch.Tensor,
                          chunk_size: int = 256) -> torch.Tensor:
        """
        Process a sequence with truncated BPTT.
        
        Args:
            token_ids: (seq_len,) token IDs
            chunk_size: gradient truncation window
        
        Returns:
            total_loss: scalar loss for the sequence
        """
        device = token_ids.device
        seq_len = token_ids.shape[0]
        M_states = self.init_m_states(device)
        total_loss = torch.tensor(0.0, device=device)
        n_tokens = 0
        
        for chunk_start in range(0, seq_len - 1, chunk_size):
            chunk_end = min(chunk_start + chunk_size, seq_len - 1)
            
            # Detach M at chunk boundaries (truncated BPTT)
            M_states = [m.detach() for m in M_states]
            
            chunk_loss = torch.tensor(0.0, device=device, requires_grad=True)
            
            for t in range(chunk_start, chunk_end):
                logits, M_states = self.forward_token(
                    token_ids[t].item(), M_states, t
                )
                loss = F.cross_entropy(
                    logits.unsqueeze(0),
                    token_ids[t + 1].unsqueeze(0)
                )
                chunk_loss = chunk_loss + loss
                n_tokens += 1
            
            total_loss = total_loss + chunk_loss
        
        return total_loss / max(n_tokens, 1)


def create_model(size: str = "small") -> CozModel:
    """Create a Cozet model by size name."""
    configs = {
        "small": COZET_SMALL,
        "medium": COZET_MEDIUM,
        "large": COZET_LARGE,
    }
    config = configs.get(size, COZET_SMALL)
    model = CozModel(config)
    n_params = sum(p.numel() for p in model.parameters())
    print(f"Cozet-{size}: {n_params:,} parameters")
    print(f"  D={config.hidden_size}, L={config.num_layers}, "
          f"H={config.num_attention_heads}/{config.num_kv_heads} GQA, "
          f"I={config.intermediate_size}, V={config.vocab_size}")
    return model


if __name__ == "__main__":
    # Quick sanity check
    model = create_model("small")
    device = torch.device("cpu")
    
    M_states = model.init_m_states(device)
    
    # Forward one token
    logits, M_states = model.forward_token(42, M_states, 0)
    print(f"Logits shape: {logits.shape}")
    print(f"M[0] norm after 1 token: {M_states[0].norm():.4f}")
    
    # Forward second token
    logits2, M_states = model.forward_token(100, M_states, 1)
    print(f"M[0] norm after 2 tokens: {M_states[0].norm():.4f}")
    
    # Test gradient flow
    M_states_fresh = model.init_m_states(device)
    logits3, _ = model.forward_token(42, M_states_fresh, 0)
    loss = F.cross_entropy(logits3.unsqueeze(0), torch.tensor([100]))
    loss.backward()
    
    grad_norms = {n: p.grad.norm().item() for n, p in model.named_parameters() 
                  if p.grad is not None}
    print(f"Gradient flow: {len(grad_norms)} parameters received gradients")
    print(f"Loss: {loss.item():.4f}")
    print("Sanity check PASSED.")