modeling_shivik_m1.py

# modeling_shivik_m1.py
# Minimal HF-compatible wrapper around your local Aries implementation.
# This file exposes ShivikM1Config and ShivikM1ForCausalLM for transformers' auto_map.

import torch
import torch.nn as nn
from transformers import PretrainedConfig, PreTrainedModel

# -------------------------
# Provide minimal config
# -------------------------
class ShivikM1Config(PretrainedConfig):
    model_type = "shivik-m1"
    def __init__(self, **kwargs):
        # defaults (tweak if needed)
        kwargs.setdefault("vocab_size", 49156)
        kwargs.setdefault("d_model", 2048)
        kwargs.setdefault("n_layers", 24)
        kwargs.setdefault("num_heads", 16)
        kwargs.setdefault("num_paths", 3)
        kwargs.setdefault("rotary_dim", 128)
        kwargs.setdefault("context_length", 4096)
        super().__init__(**kwargs)

# -------------------------
# Insert your Aries model implementation here (cleaned)
# I will include a compact, self-contained implementation compatible with the
# code you provided earlier.
# -------------------------
import math
import torch.nn.functional as F

class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.eps = eps
        self.weight = nn.Parameter(torch.ones(dim))
    def forward(self, x):
        norm = x.pow(2).mean(-1, keepdim=True)
        x = x * torch.rsqrt(norm + self.eps)
        return x * self.weight

def apply_rope(x, cos, sin):
    x1 = x[..., ::2]
    x2 = x[..., 1::2]
    x_rot = torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
    return x_rot

class MultiPathAttention(nn.Module):
    def __init__(self, d_model, num_heads, num_paths, rotary_dim):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads
        self.num_paths = num_paths
        self.head_dim = d_model // num_heads
        self.rotary_dim = rotary_dim
        self.qkv_proj = nn.Linear(d_model, 3 * d_model, bias=False)
        self.o_proj = nn.Linear(d_model, d_model, bias=False)
        self.path_weights = nn.Parameter(torch.zeros(num_paths))
    def forward(self, x, cos, sin, mask, past_kv=None):
        B, T, C = x.shape
        qkv = self.qkv_proj(x)
        q, k, v = qkv.chunk(3, dim=-1)
        q = q.view(B, T, self.num_heads, self.head_dim)
        k = k.view(B, T, self.num_heads, self.head_dim)
        v = v.view(B, T, self.num_heads, self.head_dim)
        if self.rotary_dim > 0:
            q[..., :self.rotary_dim] = apply_rope(q[..., :self.rotary_dim], cos, sin)
            k[..., :self.rotary_dim] = apply_rope(k[..., :self.rotary_dim], cos, sin)
        if past_kv is not None:
            past_k, past_v = past_kv
            k = torch.cat([past_k, k], dim=1)
            v = torch.cat([past_v, v], dim=1)
        present = (k, v)
        path_attn = []
        for _ in range(self.num_paths):
            scores = (q @ k.transpose(-2, -1)) / (self.head_dim ** 0.5)
            scores = scores + mask
            att = scores.softmax(-1)
            out = att @ v
            path_attn.append(out)
        weights = F.softmax(self.path_weights, dim=0)
        final = sum(w * p for w, p in zip(weights, path_attn))
        final = final.reshape(B, T, C)
        out = self.o_proj(final)
        return out, present

class SwiGLU(nn.Module):
    def __init__(self, dim, hidden_dim):
        super().__init__()
        self.w1 = nn.Linear(dim, hidden_dim)
        self.w2 = nn.Linear(dim, hidden_dim)
        self.w3 = nn.Linear(hidden_dim, dim)
    def forward(self, x):
        return self.w3(F.silu(self.w1(x)) * self.w2(x))

class AriesBlock(nn.Module):
    def __init__(self, cfg):
        super().__init__()
        self.norm1 = RMSNorm(cfg.d_model)
        self.attn = MultiPathAttention(cfg.d_model, cfg.num_heads, cfg.num_paths, cfg.rotary_dim)
        self.norm2 = RMSNorm(cfg.d_model)
        self.mlp = SwiGLU(cfg.d_model, 4 * cfg.d_model)
    def forward(self, x, cos, sin, mask, past_kv=None):
        h, present = self.attn(self.norm1(x), cos, sin, mask, past_kv)
        x = x + h
        x = x + self.mlp(self.norm2(x))
        return x, present

class ShivikM1Model(nn.Module):
    def __init__(self, cfg: ShivikM1Config):
        super().__init__()
        vocab_size = getattr(cfg, "vocab_size", 49156)
        d_model = getattr(cfg, "d_model", 2048)
        n_layers = getattr(cfg, "n_layers", 24)
        num_heads = getattr(cfg, "num_heads", 16)
        ctxt = getattr(cfg, "context_length", 4096)
        num_paths = getattr(cfg, "num_paths", 3)
        rotary_dim = getattr(cfg, "rotary_dim", 128)
        self.cfg = cfg
        self.token_embed = nn.Embedding(vocab_size, d_model)
        self.pos_embed = nn.Parameter(torch.zeros(1, ctxt, d_model))
        mask = torch.tril(torch.ones(ctxt, ctxt)).unsqueeze(0).unsqueeze(0)
        mask = (mask == 0).float() * -1e4
        self.register_buffer("causal_mask", mask)
        t = torch.arange(ctxt)
        freqs = 1.0 / (10000 ** (torch.arange(0, rotary_dim, 2) / rotary_dim))
        angles = torch.einsum("i,j->ij", t, freqs)
        cos = angles.cos().unsqueeze(1)
        sin = angles.sin().unsqueeze(1)
        self.register_buffer("rope_cos", cos)
        self.register_buffer("rope_sin", sin)
        self.blocks = nn.ModuleList([AriesBlock(cfg) for _ in range(n_layers)])
        self.norm = RMSNorm(d_model)
        self.lm_head = nn.Linear(d_model, vocab_size, bias=False)
        # tie weights
        try:
            self.lm_head.weight = self.token_embed.weight
        except Exception:
            pass
    def forward(self, input_ids, past_kvs=None):
        B, T = input_ids.shape
        x = self.token_embed(input_ids) + self.pos_embed[:, :T]
        mask = self.causal_mask[:, :, :T, :T]
        presents = []
        if past_kvs is None:
            past_kvs = [None] * len(self.blocks)
        for i, block in enumerate(self.blocks):
            x, present = block(x, self.rope_cos[:T], self.rope_sin[:T], mask, past_kvs[i])
            presents.append(present)
        x = self.norm(x)
        logits = self.lm_head(x)
        return {"logits": logits, "present_kvs": presents}

# -------------------------
# Finally wrap as HF PreTrainedModel
# -------------------------
class ShivikM1ForCausalLM(PreTrainedModel):
    config_class = ShivikM1Config
    base_model_prefix = "shivik_m1"
    def __init__(self, config):
        PreTrainedModel.__init__(self, config)
        self.model = ShivikM1Model(config)
    def forward(self, input_ids=None, **kwargs):
        return self.model(input_ids, **kwargs)