modeling_ivme.py

"""HuggingFace Transformers model for Ivme-Conversate-v2.

Reimplements the original IvmeConversateV2 architecture as a PreTrainedModel
so it works with AutoModelForCausalLM, .generate(), and safetensors. Math
(RMSNorm, RoPE, SwiGLU, tied embeddings, full causal attention) is unchanged
from the original; adds an optional KV cache for efficient generation.
"""

from typing import Optional

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

from .configuration_ivme import IvmeConfig


class IvmeRMSNorm(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:
        dtype = x.dtype
        x = x.float()
        rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
        out = x * rms
        return (out.to(dtype)) * self.weight


def _precompute_rope_freqs(head_dim: int, max_seq_len: int, theta: float, device=None):
    assert head_dim % 2 == 0, "RoPE requires an even head_dim"
    freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    positions = torch.arange(max_seq_len, device=device).float()
    angles = torch.outer(positions, freqs)
    return torch.polar(torch.ones_like(angles), angles)


def _apply_rope(x: torch.Tensor, rope_freqs: torch.Tensor) -> torch.Tensor:
    B, H, T, D = x.shape
    x_complex = torch.view_as_complex(x.float().reshape(B, H, T, D // 2, 2))
    freqs = rope_freqs.view(1, 1, T, D // 2)
    x_rotated = x_complex * freqs
    out = torch.view_as_real(x_rotated).reshape(B, H, T, D)
    return out.type_as(x)


class IvmeSelfAttention(nn.Module):
    def __init__(self, config: IvmeConfig, layer_idx: int):
        super().__init__()
        self.layer_idx = layer_idx
        hidden_dim = config.hidden_dim
        self.n_heads = config.n_heads
        self.head_dim = hidden_dim // config.n_heads
        self.dropout = config.dropout

        self.q_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.k_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.v_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)
        self.out_proj = nn.Linear(hidden_dim, hidden_dim, bias=False)

    def forward(self, x, rope_freqs, past_key_value=None):
        B, T, C = x.shape

        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)

        q = _apply_rope(q, rope_freqs)
        k = _apply_rope(k, rope_freqs)

        if past_key_value is not None:
            k, v = past_key_value.update(k, v, self.layer_idx)

        is_causal = past_key_value is None or k.shape[2] == q.shape[2]

        out = F.scaled_dot_product_attention(
            q, k, v, is_causal=is_causal,
            dropout_p=self.dropout if self.training else 0.0,
        )
        out = out.transpose(1, 2).contiguous().view(B, T, C)
        return self.out_proj(out)


class IvmeSwiGLU(nn.Module):
    def __init__(self, config: IvmeConfig):
        super().__init__()
        hidden_dim = config.hidden_dim
        inner_dim = int(hidden_dim * config.ffn_mult * 2 / 3)
        inner_dim = ((inner_dim + 7) // 8) * 8
        self.gate_proj = nn.Linear(hidden_dim, inner_dim, bias=False)
        self.up_proj = nn.Linear(hidden_dim, inner_dim, bias=False)
        self.down_proj = nn.Linear(inner_dim, hidden_dim, bias=False)

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class IvmeBlock(nn.Module):
    def __init__(self, config: IvmeConfig, layer_idx: int):
        super().__init__()
        self.attn_norm = IvmeRMSNorm(config.hidden_dim, eps=config.norm_eps)
        self.attn = IvmeSelfAttention(config, layer_idx)
        self.ffn_norm = IvmeRMSNorm(config.hidden_dim, eps=config.norm_eps)
        self.ffn = IvmeSwiGLU(config)

    def forward(self, x, rope_freqs, past_key_value=None):
        x = x + self.attn(self.attn_norm(x), rope_freqs, past_key_value=past_key_value)
        x = x + self.ffn(self.ffn_norm(x))
        return x


class IvmePreTrainedModel(PreTrainedModel):
    config_class = IvmeConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = False
    _no_split_modules = ["IvmeBlock"]
    _supports_cache_class = True
    _supports_sdpa = True

    def _init_weights(self, 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)


class IvmeModel(IvmePreTrainedModel):
    def __init__(self, config: IvmeConfig):
        super().__init__(config)
        self.tok_embed = nn.Embedding(config.vocab_size, config.hidden_dim)
        self.blocks = nn.ModuleList(
            [IvmeBlock(config, layer_idx=i) for i in range(config.n_layers)]
        )
        self.final_norm = IvmeRMSNorm(config.hidden_dim, eps=config.norm_eps)
        self.post_init()

    def get_input_embeddings(self):
        return self.tok_embed

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

    def forward(self, input_ids, past_key_values=None, use_cache=False, **kwargs):
        B, T = input_ids.shape

        past_len = 0
        if past_key_values is not None and len(past_key_values) > 0:
            past_len = past_key_values.get_seq_length()

        if past_len + T > self.config.context_len:
            raise ValueError(
                f"sequence length {past_len + T} exceeds context_len {self.config.context_len}"
            )

        full_rope_freqs = _precompute_rope_freqs(
            self.config.head_dim, self.config.context_len, self.config.rope_theta,
            device=input_ids.device,
        )
        rope_freqs = full_rope_freqs[past_len: past_len + T]

        x = self.tok_embed(input_ids)
        for block in self.blocks:
            x = block(x, rope_freqs, past_key_value=past_key_values)
        x = self.final_norm(x)
        return x


class IvmeForCausalLM(IvmePreTrainedModel, GenerationMixin):
    _tied_weights_keys = {"lm_head.weight": "model.tok_embed.weight"}

    def __init__(self, config: IvmeConfig):
        super().__init__(config)
        self.model = IvmeModel(config)
        self.lm_head = nn.Linear(config.hidden_dim, config.vocab_size, bias=False)
        self.post_init()
        if config.tie_word_embeddings:
            self.tie_weights()

    def get_input_embeddings(self):
        return self.model.tok_embed

    def set_input_embeddings(self, value):
        self.model.tok_embed = value

    def get_output_embeddings(self):
        return self.lm_head

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

    def forward(
        self, input_ids, attention_mask=None, past_key_values=None,
        labels=None, use_cache=None, return_dict=True, **kwargs,
    ):
        if use_cache and past_key_values is None:
            past_key_values = DynamicCache()

        hidden_states = self.model(
            input_ids,
            past_key_values=past_key_values if use_cache else None,
            use_cache=use_cache,
        )
        logits = self.lm_head(hidden_states)

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100,
            )

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=past_key_values if use_cache else None,
        )


__all__ = ["IvmeConfig", "IvmeModel", "IvmeForCausalLM"]