"""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"]