model.py

"""
MOTHER CORE Reasoning Model - Complete Implementation
Includes: RoPE, RMSNorm, SwiGLU, GQA, MoE, MLA, QK-Norm, Post-Norm, Hidden States
"""

from __future__ import annotations

import math
import warnings
from typing import Optional, Tuple, List, Dict, Any

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

from .config import ModelConfig


# ============================================================
# Helper functions
# ============================================================

def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1 = x[..., ::2]
    x2 = x[..., 1::2]
    return torch.stack((-x2, x1), dim=-1).flatten(-2)


# ============================================================
# RMSNorm
# ============================================================

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

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


# ============================================================
# Rotary Embedding
# ============================================================

class RotaryEmbedding(nn.Module):
    def __init__(self, head_dim: int, max_seq_len: int, theta: float = 10000.0):
        super().__init__()
        self.head_dim = head_dim
        self.max_seq_len = max_seq_len
        assert head_dim % 2 == 0, "head_dim must be even for RoPE"
        inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
        t = torch.arange(max_seq_len).float()
        freqs = torch.outer(t, inv_freq)
        emb = torch.cat([freqs, freqs], dim=-1)
        self.register_buffer("cos_cached", emb.cos(), persistent=False)
        self.register_buffer("sin_cached", emb.sin(), persistent=False)

    def forward(self, q: torch.Tensor, k: torch.Tensor, offset: int = 0) -> Tuple[torch.Tensor, torch.Tensor]:
        seq_len = q.shape[-2]
        if offset + seq_len > self.cos_cached.size(0):
            raise ValueError(
                f"RoPE offset+seq_len ({offset + seq_len}) exceeds max_seq_len ({self.max_seq_len})"
            )
        cos = self.cos_cached[offset:offset + seq_len].to(device=q.device, dtype=q.dtype)
        sin = self.sin_cached[offset:offset + seq_len].to(device=q.device, dtype=q.dtype)
        cos = cos.unsqueeze(0).unsqueeze(0)
        sin = sin.unsqueeze(0).unsqueeze(0)
        q = (q * cos) + (rotate_half(q) * sin)
        k = (k * cos) + (rotate_half(k) * sin)
        return q, k


# ============================================================
# SwiGLU
# ============================================================

class SwiGLU(nn.Module):
    def __init__(self, dim: int, hidden_dim: int, dropout: float = 0.0):
        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)
        self.dropout = nn.Dropout(dropout)

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


# ============================================================
# MoE Components
# ============================================================

class ExpertMLP(nn.Module):
    def __init__(self, dim: int, hidden_dim: int, dropout: float = 0.0):
        super().__init__()
        self.net = SwiGLU(dim, hidden_dim, dropout)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(x)


class MoE(nn.Module):
    def __init__(self, dim: int, hidden_dim: int, n_experts: int, top_k: int, dropout: float = 0.0):
        super().__init__()
        assert 1 <= top_k <= n_experts
        self.n_experts = n_experts
        self.top_k = top_k
        self.gate = nn.Linear(dim, n_experts, bias=False)
        self.experts = nn.ModuleList([ExpertMLP(dim, hidden_dim, dropout) for _ in range(n_experts)])

    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
        b, t, d = x.shape
        n = b * t
        flat = x.reshape(n, d)

        gate_logits = self.gate(flat)
        gate_probs = F.softmax(gate_logits, dim=-1)

        topk_vals, topk_idx = torch.topk(gate_probs, k=self.top_k, dim=-1)
        topk_vals = topk_vals / topk_vals.sum(dim=-1, keepdim=True).clamp_min(1e-9)

        out = torch.zeros_like(flat)

        for expert_id, expert in enumerate(self.experts):
            mask = (topk_idx == expert_id)
            if not mask.any():
                continue

            token_idx, topk_slot = mask.nonzero(as_tuple=True)
            expert_in = flat[token_idx]
            expert_out = expert(expert_in)
            weight = topk_vals[token_idx, topk_slot].unsqueeze(-1)
            out.index_add_(0, token_idx, expert_out * weight)

        importance = gate_probs.mean(dim=0)
        load = torch.zeros(self.n_experts, device=x.device, dtype=x.dtype)
        load.scatter_add_(0, topk_idx.reshape(-1), torch.ones_like(topk_idx.reshape(-1), dtype=x.dtype))
        load = load / load.sum().clamp_min(1e-9)
        aux_loss = self.n_experts * torch.sum(importance * load)

        return out.view(b, t, d), aux_loss


# ============================================================
# Attention Components
# ============================================================

class CausalSelfAttention(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        assert config.dim % config.n_heads == 0
        assert config.n_heads % config.n_kv_heads == 0

        self.dim = config.dim
        self.n_heads = config.n_heads
        self.n_kv_heads = config.n_kv_heads
        self.head_dim = config.dim // config.n_heads
        self.scale = self.head_dim ** -0.5
        self.dropout = config.dropout
        self.use_sdpa = config.use_flash_if_available and hasattr(F, "scaled_dot_product_attention")
        self.use_qk_norm = config.use_qk_norm

        self.wq = nn.Linear(config.dim, config.n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(config.dim, config.n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(config.dim, config.n_kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(config.dim, config.dim, bias=False)

        self.rope = RotaryEmbedding(self.head_dim, config.max_seq_len, config.rope_theta)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)

        if self.use_qk_norm:
            self.qk_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)

    def _causal_mask(self, q_len: int, k_len: int, device: torch.device) -> torch.Tensor:
        i = torch.arange(q_len, device=device)[:, None]
        j = torch.arange(k_len, device=device)[None, :]
        return (j <= i + (k_len - q_len)).view(1, 1, q_len, k_len)

    def _sliding_window_mask(self, q_len: int, k_len: int, offset: int, window_size: int, device: torch.device) -> torch.Tensor:
        q_abs = torch.arange(offset, offset + q_len, device=device)[:, None]
        k_abs = torch.arange(k_len, device=device)[None, :]
        mask = (k_abs <= q_abs) & (k_abs >= q_abs - window_size + 1)
        return mask.view(1, 1, q_len, k_len)

    def forward(
        self,
        x: torch.Tensor,
        past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        use_cache: bool = False,
        attention_mask: Optional[torch.Tensor] = None,
        window_size: Optional[int] = None,
    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
        b, t, d = x.shape

        q = self.wq(x).view(b, t, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.wk(x).view(b, t, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.wv(x).view(b, t, self.n_kv_heads, self.head_dim).transpose(1, 2)

        offset = 0 if past_kv is None else past_kv[0].shape[-2]
        q, k = self.rope(q, k, offset=offset)

        if self.use_qk_norm:
            q = self.qk_norm(q)
            k = self.qk_norm(k)

        if past_kv is not None:
            pk, pv = past_kv
            k = torch.cat([pk, k], dim=-2)
            v = torch.cat([pv, v], dim=-2)

        new_past = (k, v) if use_cache else None

        if self.n_heads != self.n_kv_heads:
            repeat = self.n_heads // self.n_kv_heads
            k = k.repeat_interleave(repeat, dim=1)
            v = v.repeat_interleave(repeat, dim=1)

        # Build mask
        if attention_mask is not None:
            if attention_mask.dim() == 2:
                pad_len = attention_mask.shape[-1]
                if pad_len < k.shape[-2]:
                    extra = torch.ones((b, k.shape[-2] - pad_len), device=attention_mask.device, dtype=attention_mask.dtype)
                    attention_mask = torch.cat([attention_mask, extra], dim=-1)
                elif pad_len > k.shape[-2]:
                    attention_mask = attention_mask[:, :k.shape[-2]]
                mask = attention_mask[:, None, None, :]
                causal = self._causal_mask(t, k.shape[-2], x.device)
                mask = mask & causal
            elif attention_mask.dim() == 4:
                mask = attention_mask.bool()
            else:
                raise ValueError(f"Unsupported attention_mask shape: {attention_mask.shape}")
        elif window_size is not None and not use_cache:
            mask = self._sliding_window_mask(t, k.shape[-2], offset, window_size, x.device)
        else:
            mask = None

        if mask is not None:
            additive_mask = torch.zeros_like(mask, dtype=x.dtype)
            additive_mask = additive_mask.masked_fill(~mask, float("-inf"))
        else:
            additive_mask = None

        if self.use_sdpa:
            y = F.scaled_dot_product_attention(
                q, k, v,
                attn_mask=additive_mask,
                dropout_p=self.dropout if self.training else 0.0,
                is_causal=(mask is None and past_kv is None),
            )
        else:
            scores = (q @ k.transpose(-2, -1)) * self.scale
            if mask is None:
                mask = self._causal_mask(t, k.shape[-2], x.device)
            scores = scores.masked_fill(~mask, float("-inf"))
            attn = F.softmax(scores, dim=-1)
            attn = self.attn_dropout(attn)
            y = attn @ v

        y = y.transpose(1, 2).contiguous().view(b, t, d)
        y = self.resid_dropout(self.wo(y))
        return y, new_past


class MLAAttention(nn.Module):
    """EXPERIMENTAL: Multi-head Latent Attention (inspired by DeepSeek-V3)"""
    
    def __init__(self, config: ModelConfig):
        super().__init__()
        warnings.warn("MLAAttention is experimental and may not be as stable as standard attention.", UserWarning)
        assert config.dim % config.n_heads == 0
        assert config.n_heads % config.n_kv_heads == 0

        self.dim = config.dim
        self.n_heads = config.n_heads
        self.n_kv_heads = config.n_kv_heads
        self.head_dim = config.dim // config.n_heads
        self.latent_dim = config.mla_latent_dim
        self.scale = self.head_dim ** -0.5
        self.dropout = config.dropout
        self.use_sdpa = config.use_flash_if_available and hasattr(F, "scaled_dot_product_attention")
        self.use_qk_norm = config.use_qk_norm

        self.wq = nn.Linear(config.dim, config.n_heads * self.head_dim, bias=False)
        self.wkv = nn.Linear(config.dim, self.latent_dim, bias=False)
        self.wk_up = nn.Linear(self.latent_dim, config.n_kv_heads * self.head_dim, bias=False)
        self.wv_up = nn.Linear(self.latent_dim, config.n_kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(config.dim, config.dim, bias=False)

        self.rope = RotaryEmbedding(self.head_dim, config.max_seq_len, config.rope_theta)
        self.attn_dropout = nn.Dropout(config.dropout)
        self.resid_dropout = nn.Dropout(config.dropout)

        if self.use_qk_norm:
            self.qk_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)

    def _causal_mask(self, q_len: int, k_len: int, device: torch.device) -> torch.Tensor:
        i = torch.arange(q_len, device=device)[:, None]
        j = torch.arange(k_len, device=device)[None, :]
        return (j <= i + (k_len - q_len)).view(1, 1, q_len, k_len)

    def _sliding_window_mask(self, q_len: int, k_len: int, offset: int, window_size: int, device: torch.device) -> torch.Tensor:
        q_abs = torch.arange(offset, offset + q_len, device=device)[:, None]
        k_abs = torch.arange(k_len, device=device)[None, :]
        mask = (k_abs <= q_abs) & (k_abs >= q_abs - window_size + 1)
        return mask.view(1, 1, q_len, k_len)

    def forward(
        self,
        x: torch.Tensor,
        past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        use_cache: bool = False,
        attention_mask: Optional[torch.Tensor] = None,
        window_size: Optional[int] = None,
    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]:
        b, t, d = x.shape

        q = self.wq(x).view(b, t, self.n_heads, self.head_dim).transpose(1, 2)
        latent = self.wkv(x)
        k = self.wk_up(latent).view(b, t, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.wv_up(latent).view(b, t, self.n_kv_heads, self.head_dim).transpose(1, 2)

        offset = 0 if past_kv is None else past_kv[0].shape[-2]
        q, k = self.rope(q, k, offset=offset)

        if self.use_qk_norm:
            q = self.qk_norm(q)
            k = self.qk_norm(k)

        if past_kv is not None:
            pk, pv = past_kv
            k = torch.cat([pk, k], dim=-2)
            v = torch.cat([pv, v], dim=-2)

        new_past = (k, v) if use_cache else None

        if self.n_heads != self.n_kv_heads:
            repeat = self.n_heads // self.n_kv_heads
            k = k.repeat_interleave(repeat, dim=1)
            v = v.repeat_interleave(repeat, dim=1)

        if attention_mask is not None:
            if attention_mask.dim() == 2:
                pad_len = attention_mask.shape[-1]
                if pad_len < k.shape[-2]:
                    extra = torch.ones((b, k.shape[-2] - pad_len), device=attention_mask.device, dtype=attention_mask.dtype)
                    attention_mask = torch.cat([attention_mask, extra], dim=-1)
                elif pad_len > k.shape[-2]:
                    attention_mask = attention_mask[:, :k.shape[-2]]
                mask = attention_mask[:, None, None, :]
                causal = self._causal_mask(t, k.shape[-2], x.device)
                mask = mask & causal
            elif attention_mask.dim() == 4:
                mask = attention_mask.bool()
            else:
                raise ValueError(f"Unsupported attention_mask shape: {attention_mask.shape}")
        elif window_size is not None and not use_cache:
            mask = self._sliding_window_mask(t, k.shape[-2], offset, window_size, x.device)
        else:
            mask = None

        if mask is not None:
            additive_mask = torch.zeros_like(mask, dtype=x.dtype)
            additive_mask = additive_mask.masked_fill(~mask, float("-inf"))
        else:
            additive_mask = None

        if self.use_sdpa:
            y = F.scaled_dot_product_attention(
                q, k, v,
                attn_mask=additive_mask,
                dropout_p=self.dropout if self.training else 0.0,
                is_causal=(mask is None and past_kv is None),
            )
        else:
            scores = (q @ k.transpose(-2, -1)) * self.scale
            if mask is None:
                mask = self._causal_mask(t, k.shape[-2], x.device)
            scores = scores.masked_fill(~mask, float("-inf"))
            attn = F.softmax(scores, dim=-1)
            attn = self.attn_dropout(attn)
            y = attn @ v

        y = y.transpose(1, 2).contiguous().view(b, t, d)
        y = self.resid_dropout(self.wo(y))
        return y, new_past


# ============================================================
# Transformer Block
# ============================================================

class MotherCoreBlock(nn.Module):
    def __init__(self, config: ModelConfig, layer_idx: int):
        super().__init__()
        hidden_dim = int(config.dim * config.ff_mult)

        if config.use_mla:
            self.attn = MLAAttention(config)
        else:
            self.attn = CausalSelfAttention(config)

        self.use_moe_here = config.use_moe and (layer_idx % config.moe_every == 0)
        if self.use_moe_here:
            self.ff = MoE(dim=config.dim, hidden_dim=hidden_dim, n_experts=config.n_experts,
                          top_k=config.moe_top_k, dropout=config.dropout)
        else:
            self.ff = SwiGLU(config.dim, hidden_dim, dropout=config.dropout)

        self.residual_scale = config.residual_scale
        self.window_size = config.window_size if config.sliding_window_attention else None
        self.post_norm = config.use_post_norm

        self.norm1 = RMSNorm(config.dim, eps=config.rms_norm_eps)
        self.norm2 = RMSNorm(config.dim, eps=config.rms_norm_eps)

        if self.post_norm:
            self.norm_attn = RMSNorm(config.dim, eps=config.rms_norm_eps)
            self.norm_ff = RMSNorm(config.dim, eps=config.rms_norm_eps)

    def forward(
        self,
        x: torch.Tensor,
        past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
        use_cache: bool = False,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]], torch.Tensor]:
        aux_loss = x.new_zeros(())

        if self.post_norm:
            attn_out, new_past = self.attn(x, past_kv=past_kv, use_cache=use_cache,
                                            attention_mask=attention_mask, window_size=self.window_size)
            x = self.norm_attn(x + self.residual_scale * attn_out)

            if self.use_moe_here:
                ff_out, moe_aux = self.ff(x)
                aux_loss = aux_loss + moe_aux
            else:
                ff_out = self.ff(x)
            x = self.norm_ff(x + self.residual_scale * ff_out)
        else:
            attn_out, new_past = self.attn(self.norm1(x), past_kv=past_kv, use_cache=use_cache,
                                            attention_mask=attention_mask, window_size=self.window_size)
            x = x + self.residual_scale * attn_out

            if self.use_moe_here:
                ff_out, moe_aux = self.ff(self.norm2(x))
                aux_loss = aux_loss + moe_aux
            else:
                ff_out = self.ff(self.norm2(x))
            x = x + self.residual_scale * ff_out

        return x, new_past, aux_loss


# ============================================================
# Main MOTHER CORE Model
# ============================================================

class MotherCoreModel(nn.Module):
    def __init__(self, config: ModelConfig):
        super().__init__()
        self.config = config
        self.gradient_checkpointing = False

        self.tok_emb = nn.Embedding(config.vocab_size, config.dim)
        self.drop = nn.Dropout(config.dropout)
        self.blocks = nn.ModuleList([MotherCoreBlock(config, i) for i in range(config.n_layers)])
        self.norm_f = RMSNorm(config.dim, eps=config.rms_norm_eps)
        self.lm_head = nn.Linear(config.dim, config.vocab_size, bias=False)
        
        # Memory gate for RAG control
        self.memory_gate = nn.Linear(config.dim, 1)
        self.memory_gate_loss_weight = 0.1

        if config.tie_embeddings:
            self.lm_head.weight = self.tok_emb.weight

        self.apply(self._init_weights)

        for name, p in self.named_parameters():
            if name.endswith("wo.weight") or name.endswith("w3.weight"):
                nn.init.normal_(p, mean=0.0, std=0.02 / math.sqrt(2 * config.n_layers))

    def _init_weights(self, module: nn.Module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=0.02)

    def enable_gradient_checkpointing(self):
        self.gradient_checkpointing = True

    def forward(
        self,
        input_ids: torch.Tensor,
        labels: Optional[torch.Tensor] = None,
        past_key_values: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None,
        use_cache: bool = False,
        attention_mask: Optional[torch.Tensor] = None,
        output_hidden_states: bool = False,
    ) -> Dict[str, Any]:
        b, t = input_ids.shape
        if t > self.config.max_seq_len:
            raise ValueError(f"Sequence length {t} exceeds max_seq_len={self.config.max_seq_len}")

        if self.training:
            use_cache = False

        x = self.drop(self.tok_emb(input_ids))
        total_aux_loss = x.new_zeros(())
        new_past = [] if use_cache else None
        hidden_states = [] if output_hidden_states else None

        if past_key_values is None:
            past_key_values = [None] * len(self.blocks)

        for block, past_kv in zip(self.blocks, past_key_values):
            if self.gradient_checkpointing and self.training:
                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        return module(*inputs)
                    return custom_forward

                x, block_past, aux = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(block),
                    x, past_kv, use_cache, attention_mask,
                    use_reentrant=False,
                )
            else:
                x, block_past, aux = block(x, past_kv=past_kv, use_cache=use_cache,
                                            attention_mask=attention_mask)

            total_aux_loss = total_aux_loss + aux
            if use_cache:
                new_past.append(block_past)
            if output_hidden_states:
                hidden_states.append(x)

        x = self.norm_f(x)
        
        # === MEMORY GATE ===
        last_hidden = x[:, -1, :]  # last token
        gate_logits = self.memory_gate(last_hidden)
        gate = torch.sigmoid(gate_logits)  # [B, 1]
        logits = self.lm_head(x)

        loss = None
        if labels is not None:
            ce_loss = F.cross_entropy(
                logits[:, :-1, :].reshape(-1, logits.size(-1)),
                labels[:, 1:].reshape(-1),
                ignore_index=-100,
            )
            # === GATE SUPERVISION ===
            loss = ce_loss
            if self.config.use_moe:
                loss = loss + self.config.aux_loss_alpha * total_aux_loss

        return {
            "logits": logits,
            "loss": loss,
            "aux_loss": total_aux_loss,
            "past_key_values": new_past,
            "hidden_states": hidden_states,
            "last_hidden_state": x,
            "gate": gate.detach(),
        }

    @torch.no_grad()
    def generate(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int = 128,
        temperature: float = 0.8,
        top_k: Optional[int] = 50,
        eos_token_id: Optional[int] = None,
    ) -> torch.Tensor:
        self.eval()
        cur = input_ids
        past = None

        for _ in range(max_new_tokens):
            if past is None:
                out = self(cur, use_cache=True)
            else:
                out = self(cur[:, -1:], past_key_values=past, use_cache=True)

            logits = out["logits"][:, -1, :] / max(temperature, 1e-5)
            past = out["past_key_values"]

            if top_k is not None:
                vals, _ = torch.topk(logits, min(top_k, logits.size(-1)))
                logits[logits < vals[:, [-1]]] = -float("inf")

            probs = F.softmax(logits, dim=-1)
            next_token = torch.multinomial(probs, num_samples=1)
            cur = torch.cat([cur, next_token], dim=1)

            if eos_token_id is not None and (next_token == eos_token_id).all():
                break

        return cur