modeling_ivme_s_v2_instruct.py

"""
Modeling file for Ivme-Conversate-S-v2-Instruct.

Standard decoder-only Transformer architecture, deliberately matching
Ivme-Conversate-v2-Base's proven recipe (pulled directly from its real
config.json): tied embeddings, standard multi-head attention (no GQA, no
DIFF), RoPE, SwiGLU, RMSNorm, pre-norm. No architectural novelty by design --
this model tests a DATA strategy (instruct-heavy, single-epoch pretraining)
in isolation, on infrastructure already proven stable.

Trained on ~900M tokens, single epoch, instruct-heavy mix (UltraChat-200k as
the dominant 45% share, plus SODA, UltraInteract, orca-math, dolly-15k,
sql-create-context) -- all permissively licensed (MIT/CC-BY/CC-BY-SA), no
CC-BY-NC sources, matching v2-Base's Apache-2.0 license.

Uses standard HF tied-embedding conventions (get_output_embeddings /
set_output_embeddings + config.tie_word_embeddings), so PreTrainedModel's
own tie_weights() machinery handles the tie correctly through from_pretrained
-- more robust than manual weight assignment, since it's re-applied
automatically by HF's own loading path rather than needing to survive it.
"""

import math

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutput

try:
    from .configuration_ivme_s_v2_instruct import IvmeConversateSV2InstructConfig
except ImportError:
    from configuration_ivme_s_v2_instruct import IvmeConversateSV2InstructConfig


def build_rope_cache(dim, max_seq_len, base=10000.0):
    assert dim % 2 == 0
    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
    t = torch.arange(max_seq_len).float()
    freqs = torch.outer(t, inv_freq)
    emb = torch.cat([freqs, freqs], dim=-1)
    return emb.cos(), emb.sin()


def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)


def apply_rope(x, cos, sin):
    T = x.shape[-2]
    cos = cos[:T].unsqueeze(0).unsqueeze(0).to(x.dtype)
    sin = sin[:T].unsqueeze(0).unsqueeze(0).to(x.dtype)
    return x * cos + rotate_half(x) * sin


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

    def forward(self, x):
        norm = x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return x * norm * self.weight


class StandardAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        assert d_model % n_heads == 0
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.wqkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.wo = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x, rope_cos, rope_sin):
        B, T, D = x.shape
        qkv = self.wqkv(x)
        q, k, v = qkv.split(D, dim=-1)
        q = q.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = k.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        v = v.view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        q = apply_rope(q, rope_cos, rope_sin)
        k = apply_rope(k, rope_cos, rope_sin)
        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        out = out.transpose(1, 2).contiguous().view(B, T, D)
        return self.wo(out)


class SwiGLU(nn.Module):
    def __init__(self, d_model, d_ff):
        super().__init__()
        self.w_gate = nn.Linear(d_model, d_ff, bias=False)
        self.w_up = nn.Linear(d_model, d_ff, bias=False)
        self.w_down = nn.Linear(d_ff, d_model, bias=False)

    def forward(self, x):
        return self.w_down(F.silu(self.w_gate(x)) * self.w_up(x))


class Block(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, eps=1e-5):
        super().__init__()
        self.norm1 = RMSNorm(d_model, eps)
        self.attn = StandardAttention(d_model, n_heads)
        self.norm2 = RMSNorm(d_model, eps)
        self.ffn = SwiGLU(d_model, d_ff)

    def forward(self, x, rope_cos, rope_sin):
        x = x + self.attn(self.norm1(x), rope_cos, rope_sin)
        x = x + self.ffn(self.norm2(x))
        return x


class IvmeConversateSV2InstructModel(PreTrainedModel):
    """HF-compatible wrapper. Load with:
        AutoModelForCausalLM.from_pretrained(repo_id, trust_remote_code=True)
    """

    config_class = IvmeConversateSV2InstructConfig
    # Explicit declarative tied-weights mapping -- confirmed via direct
    # inspection of transformers' PreTrainedModel.get_expanded_tied_weights_keys
    # that get_input_embeddings()/get_output_embeddings() ALONE do not trigger
    # automatic tying in this version; the class needs _tied_weights_keys set
    # explicitly (same convention used by e.g. GPT2LMHeadModel:
    # {'lm_head.weight': 'transformer.wte.weight'}). Verified this actually
    # ties the weights via post_init() -> init_weights() -> tie_weights():
    # an earlier version of this file relied on get_output_embeddings() alone
    # and the weights were NOT tied (model.tok_embed.weight is model.lm_head.
    # weight was False) despite tie_word_embeddings=True in config.
    _tied_weights_keys = {"lm_head.weight": "tok_embed.weight"}

    def __init__(self, config: IvmeConversateSV2InstructConfig):
        super().__init__(config)
        self.tok_embed = nn.Embedding(config.vocab_size, config.d_model)
        nn.init.normal_(self.tok_embed.weight, mean=0.0, std=0.02)

        self.blocks = nn.ModuleList([
            Block(config.d_model, config.n_heads, config.d_ff, config.norm_eps)
            for _ in range(config.n_layers)
        ])
        self.norm_f = RMSNorm(config.d_model, config.norm_eps)
        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)

        head_dim = config.d_model // config.n_heads
        cos, sin = build_rope_cache(head_dim, config.max_seq_len, config.rope_theta)
        self.register_buffer("rope_cos", cos, persistent=True)
        self.register_buffer("rope_sin", sin, persistent=True)

        self.post_init()

    def get_input_embeddings(self):
        return self.tok_embed

    def set_input_embeddings(self, value):
        self.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 can_generate(self):
        return True

    def forward(self, input_ids, labels=None, **kwargs):
        x = self.tok_embed(input_ids)
        for block in self.blocks:
            x = block(x, self.rope_cos, self.rope_sin)
        x = self.norm_f(x)
        logits = self.lm_head(x)

        loss = None
        if labels is not None:
            loss = F.cross_entropy(
                logits[:, :-1, :].reshape(-1, self.config.vocab_size),
                labels[:, 1:].reshape(-1),
            )

        return CausalLMOutput(loss=loss, logits=logits)