modeling_jeeves.py

"""HuggingFace-compatible Jeeves model.

This file gets uploaded to the Hub so users can load with:
    from transformers import AutoModelForCausalLM
    model = AutoModelForCausalLM.from_pretrained("Anurich/Jeeves-Small-75M", trust_remote_code=True)

The architecture is self-contained — no local imports needed.
Features: Looped Transformer + Value Residual Learning + GQA + RoPE + SwiGLU.
"""

import math
from typing import Optional, Tuple

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import GenerationMixin, PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast

from .configuration_jeeves import JeevesConfig


# ---------------------------------------------------------------------------
# Core layers
# ---------------------------------------------------------------------------

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


class SwiGLUFFN(nn.Module):
    """SwiGLU Feed-Forward Network."""
    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.0):
        super().__init__()
        self.gate_proj = nn.Linear(d_model, d_ff, bias=False)
        self.up_proj = nn.Linear(d_model, d_ff, bias=False)
        self.down_proj = nn.Linear(d_ff, d_model, bias=False)
        self.dropout = nn.Dropout(dropout) if dropout > 0 else nn.Identity()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.dropout(self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)))


# ---------------------------------------------------------------------------
# RoPE
# ---------------------------------------------------------------------------

def precompute_rope_freqs(head_dim: int, max_seq_len: int, base: float = 10000.0,
                          device=None) -> torch.Tensor:
    freqs = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    t = torch.arange(max_seq_len, device=device).float()
    freqs = torch.outer(t, freqs)
    return torch.polar(torch.ones_like(freqs), freqs)


def apply_rope(q, k, freqs_cis):
    if q.device.type == 'mps':
        return _apply_rope_real(q, k, freqs_cis)
    q_c = torch.view_as_complex(q.float().reshape(*q.shape[:-1], -1, 2))
    k_c = torch.view_as_complex(k.float().reshape(*k.shape[:-1], -1, 2))
    f = freqs_cis.unsqueeze(0).unsqueeze(2)
    q_r = torch.view_as_real(q_c * f).flatten(-2)
    k_r = torch.view_as_real(k_c * f).flatten(-2)
    return q_r.type_as(q), k_r.type_as(k)


def _apply_rope_real(q, k, freqs_cis):
    cos = freqs_cis.real.unsqueeze(0).unsqueeze(2)
    sin = freqs_cis.imag.unsqueeze(0).unsqueeze(2)

    def _rotate(x):
        pairs = x.float().reshape(*x.shape[:-1], -1, 2)
        r, i = pairs[..., 0], pairs[..., 1]
        out = torch.stack([r * cos - i * sin, r * sin + i * cos], dim=-1).flatten(-2)
        return out.type_as(x)

    return _rotate(q), _rotate(k)


def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
    if n_rep == 1:
        return x
    b, s, kv, d = x.shape
    return x[:, :, :, None, :].expand(b, s, kv, n_rep, d).reshape(b, s, kv * n_rep, d)


# ---------------------------------------------------------------------------
# Attention with Value Residual Learning
# ---------------------------------------------------------------------------

class GQAWithValueResidual(nn.Module):
    """Grouped-Query Attention with optional Value Residual Learning."""

    def __init__(self, config: JeevesConfig):
        super().__init__()
        self.d_model = config.d_model
        self.n_heads = config.n_heads
        self.n_kv_heads = config.n_kv_heads
        self.head_dim = config.head_dim
        self.n_kv_groups = config.n_heads // config.n_kv_heads
        self.use_flash_attention = config.use_flash_attention
        self.use_value_residual = config.use_value_residual

        self.q_proj = nn.Linear(config.d_model, config.n_heads * config.head_dim, bias=False)
        self.k_proj = nn.Linear(config.d_model, config.n_kv_heads * config.head_dim, bias=False)
        self.v_proj = nn.Linear(config.d_model, config.n_kv_heads * config.head_dim, bias=False)
        self.o_proj = nn.Linear(config.n_heads * config.head_dim, config.d_model, bias=False)
        self.attn_dropout = nn.Dropout(config.dropout) if config.dropout > 0 else nn.Identity()

        if config.use_value_residual:
            self.alpha_logit = nn.Parameter(torch.tensor(config.value_residual_alpha_init))

    def forward(self, x, freqs_cis, mask=None, first_layer_v=None):
        batch, seq_len, _ = x.shape
        q = self.q_proj(x).view(batch, seq_len, self.n_heads, self.head_dim)
        k = self.k_proj(x).view(batch, seq_len, self.n_kv_heads, self.head_dim)
        v = self.v_proj(x).view(batch, seq_len, self.n_kv_heads, self.head_dim)
        raw_v = v

        if self.use_value_residual and first_layer_v is not None:
            alpha = torch.sigmoid(self.alpha_logit)
            v = (1.0 - alpha) * v + alpha * first_layer_v

        q, k = apply_rope(q, k, freqs_cis)
        k = repeat_kv(k, self.n_kv_groups)
        v = repeat_kv(v, self.n_kv_groups)

        q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)

        is_accel = q.is_cuda or q.device.type == 'mps'
        if self.use_flash_attention and is_accel:
            attn_out = F.scaled_dot_product_attention(q, k, v, attn_mask=None, is_causal=True)
        else:
            scale = 1.0 / math.sqrt(self.head_dim)
            scores = torch.matmul(q, k.transpose(-2, -1)) * scale
            if mask is not None:
                scores = scores + mask
            w = F.softmax(scores, dim=-1, dtype=torch.float32).type_as(q)
            w = self.attn_dropout(w)
            attn_out = torch.matmul(w, v)

        attn_out = attn_out.transpose(1, 2).contiguous().view(batch, seq_len, -1)
        return self.o_proj(attn_out), raw_v


# ---------------------------------------------------------------------------
# Transformer Block
# ---------------------------------------------------------------------------

class TransformerBlock(nn.Module):
    def __init__(self, config: JeevesConfig):
        super().__init__()
        self.attn_norm = RMSNorm(config.d_model, eps=config.norm_eps)
        self.attention = GQAWithValueResidual(config)
        self.ffn_norm = RMSNorm(config.d_model, eps=config.norm_eps)
        self.ffn = SwiGLUFFN(config.d_model, config.d_ff, config.dropout)

    def forward(self, x, freqs_cis, mask=None, first_layer_v=None):
        h, raw_v = self.attention(self.attn_norm(x), freqs_cis, mask, first_layer_v)
        x = x + h
        x = x + self.ffn(self.ffn_norm(x))
        return x, raw_v


# ---------------------------------------------------------------------------
# Jeeves Model (HuggingFace-compatible)
# ---------------------------------------------------------------------------

class JeevesForCausalLM(PreTrainedModel, GenerationMixin):
    """Jeeves: Looped Transformer + Value Residual Learning.

    Loads native Jeeves weights directly — no conversion needed.
    """
    config_class = JeevesConfig
    supports_gradient_checkpointing = False
    _tied_weights_keys = {"lm_head.weight": "tok_emb.weight"}

    def __init__(self, config: JeevesConfig):
        super().__init__(config)
        self.config = config

        # Embedding
        self.tok_emb = nn.Embedding(config.vocab_size, config.d_model)

        # Layer structure
        if config.loop_block_idx is not None:
            n_early = config.loop_block_idx
            n_late = config.n_layers - config.loop_block_idx - 1
            self.early_layers = nn.ModuleList([TransformerBlock(config) for _ in range(n_early)])
            self.loop_block = TransformerBlock(config)
            self.late_layers = nn.ModuleList([TransformerBlock(config) for _ in range(n_late)])
            self.n_loop_iters = config.n_loop_iters
            self.use_input_injection = config.use_input_injection
            self.looped = True
        else:
            self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)])
            self.looped = False

        self.norm = RMSNorm(config.d_model, eps=config.norm_eps)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)

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

        # Store RoPE params — freqs_cis is computed fresh in forward()
        # to avoid corruption from HF's meta-device initialization
        self._rope_head_dim = config.head_dim
        self._rope_max_seq_len = config.max_seq_len
        self._rope_base = config.rope_base
        self._freqs_cache = None

        self.post_init()

    def get_input_embeddings(self):
        return self.tok_emb

    def set_input_embeddings(self, value):
        self.tok_emb = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def _get_freqs_cis(self, seq_len: int, device: torch.device) -> torch.Tensor:
        """Get RoPE frequencies, computing and caching on first call."""
        if self._freqs_cache is None or self._freqs_cache.device != device:
            self._freqs_cache = precompute_rope_freqs(
                self._rope_head_dim, self._rope_max_seq_len, self._rope_base, device
            )
        return self._freqs_cache[:seq_len]

    def _make_causal_mask(self, seq_len, device):
        mask = torch.full((seq_len, seq_len), float("-inf"), device=device)
        return torch.triu(mask, diagonal=1)

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        **kwargs,
    ) -> CausalLMOutputWithPast:
        if inputs_embeds is None:
            h = self.tok_emb(input_ids)
        else:
            h = inputs_embeds

        batch, seq_len, _ = h.shape
        device = h.device
        freqs_cis = self._get_freqs_cis(seq_len, device)

        mask = None
        is_accel = h.is_cuda or h.device.type == 'mps'
        if not self.config.use_flash_attention or not is_accel:
            mask = self._make_causal_mask(seq_len, device)

        first_layer_v = None

        if self.looped:
            # Early layers
            for i, layer in enumerate(self.early_layers):
                h, raw_v = layer(h, freqs_cis, mask, first_layer_v)
                if i == 0 and self.config.use_value_residual:
                    first_layer_v = raw_v

            # Looped block with input injection
            loop_input = h
            for loop_iter in range(self.n_loop_iters):
                h, _ = self.loop_block(h, freqs_cis, mask, first_layer_v)
                if self.use_input_injection and loop_iter < self.n_loop_iters - 1:
                    h = h + loop_input

            # Late layers
            for layer in self.late_layers:
                h, _ = layer(h, freqs_cis, mask, first_layer_v)
        else:
            for i, layer in enumerate(self.layers):
                h, raw_v = layer(h, freqs_cis, mask, first_layer_v)
                if i == 0 and self.config.use_value_residual:
                    first_layer_v = raw_v

        h = self.norm(h)
        logits = self.lm_head(h)

        loss = None
        if labels is not None:
            loss = F.cross_entropy(
                logits.view(-1, self.config.vocab_size),
                labels.view(-1),
                ignore_index=-100,
            )

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
        )

    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return {"input_ids": input_ids}

    @staticmethod
    def _reorder_cache(past, beam_idx):
        return past