"""HuggingFace Transformers model for ExpIvme-DiffusionConversate-v1. A masked/absorbing-state discrete diffusion language model. Architecture (RMSNorm, RoPE, SwiGLU, tied embeddings) inherited from IvmeLabs/Ivme-Conversate-v2-Base, scaled to ~130M params, with bidirectional (non-causal) attention for masked diffusion. """ from dataclasses import dataclass import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel from transformers.modeling_outputs import ModelOutput try: from .configuration_expivme_diffusion import ExpIvmeDiffusionConfig except ImportError: from configuration_expivme_diffusion import ExpIvmeDiffusionConfig def _precompute_rope_freqs(head_dim, max_seq_len, theta, device=None): 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.cos(angles), torch.sin(angles) def _apply_rope(x, rope_cos_sin): cos, sin = rope_cos_sin B, H, T, D = x.shape x1 = x[..., 0::2] x2 = x[..., 1::2] cos = cos.view(1, 1, T, D // 2).to(x.dtype) sin = sin.view(1, 1, T, D // 2).to(x.dtype) out1 = x1 * cos - x2 * sin out2 = x1 * sin + x2 * cos out = torch.stack([out1, out2], dim=-1).reshape(B, H, T, D) return out.type_as(x) class ExpIvmeRMSNorm(nn.Module): def __init__(self, dim, eps=1e-5): super().__init__() self.eps = eps self.weight = nn.Parameter(torch.ones(dim)) def forward(self, x): dtype = x.dtype x = x.float() rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps) return (x * rms).to(dtype) * self.weight class ExpIvmeSelfAttention(nn.Module): def __init__(self, hidden_dim, n_heads, dropout=0.0): super().__init__() self.n_heads = n_heads self.head_dim = hidden_dim // n_heads self.dropout = 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, attn_mask=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) k = _apply_rope(k, rope) out = F.scaled_dot_product_attention( q, k, v, attn_mask=attn_mask, is_causal=False, 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 ExpIvmeSwiGLU(nn.Module): def __init__(self, hidden_dim, ffn_mult): super().__init__() inner_dim = int(hidden_dim * 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 ExpIvmeBlock(nn.Module): def __init__(self, hidden_dim, n_heads, ffn_mult, norm_eps, dropout=0.0): super().__init__() self.attn_norm = ExpIvmeRMSNorm(hidden_dim, eps=norm_eps) self.attn = ExpIvmeSelfAttention(hidden_dim, n_heads, dropout) self.ffn_norm = ExpIvmeRMSNorm(hidden_dim, eps=norm_eps) self.ffn = ExpIvmeSwiGLU(hidden_dim, ffn_mult) def forward(self, x, rope, attn_mask=None): x = x + self.attn(self.attn_norm(x), rope, attn_mask=attn_mask) x = x + self.ffn(self.ffn_norm(x)) return x @dataclass class DiffusionLMOutput(ModelOutput): loss: torch.FloatTensor = None logits: torch.FloatTensor = None class ExpIvmeForDiffusionLMHub(PreTrainedModel): """Single module tree — self.model.* and self.lm_head only.""" config_class = ExpIvmeDiffusionConfig base_model_prefix = "model" _tied_weights_keys = {"lm_head.weight": "model.tok_embed.weight"} def __init__(self, config): super().__init__(config) self.model = nn.Module() self.model.tok_embed = nn.Embedding(config.vocab_size, config.hidden_dim) self.model.blocks = nn.ModuleList([ ExpIvmeBlock(config.hidden_dim, config.n_heads, config.ffn_mult, config.norm_eps, config.dropout) for _ in range(config.n_layers) ]) self.model.final_norm = ExpIvmeRMSNorm(config.hidden_dim, eps=config.norm_eps) self.lm_head = nn.Linear(config.hidden_dim, config.vocab_size, bias=False) self.head_dim = config.hidden_dim // config.n_heads self.rope_theta = config.rope_theta 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 forward(self, input_ids, attention_mask=None, labels=None, mask_positions=None, t=None, return_dict=True, **kw): B, T = input_ids.shape rope = _precompute_rope_freqs(self.head_dim, T, self.rope_theta, device=input_ids.device) sdpa_mask = None if attention_mask is not None: sdpa_mask = torch.zeros(B, 1, 1, T, dtype=torch.float32, device=input_ids.device) sdpa_mask.masked_fill_(attention_mask[:, None, None, :] == 0, float("-inf")) sdpa_mask = sdpa_mask.to(dtype=self.model.tok_embed.weight.dtype) x = self.model.tok_embed(input_ids) for block in self.model.blocks: x = block(x, rope, attn_mask=sdpa_mask) x = self.model.final_norm(x) logits = self.lm_head(x) loss = None if labels is not None and mask_positions is not None: ce = F.cross_entropy( logits.view(-1, self.config.vocab_size), labels.view(-1), ignore_index=-100, reduction="none", ).view(B, T) ce = ce * mask_positions.float() per_example_loss = ce.sum(dim=1) if t is not None: weight = 1.0 / t.clamp(min=1e-3) per_example_loss = per_example_loss * weight n_masked = mask_positions.float().sum(dim=1).clamp(min=1.0) loss = (per_example_loss / n_masked).mean() if not return_dict: return (loss, logits) if loss is not None else (logits,) return DiffusionLMOutput(loss=loss, logits=logits) __all__ = ["ExpIvmeDiffusionConfig", "ExpIvmeForDiffusionLMHub"]