Add safetensors + Transformers (AutoModelForCausalLM) support

NEW_FILES.md

@@ -0,0 +1,33 @@
+## Transformers / safetensors support
+
+This model can now be loaded with `AutoModelForCausalLM` instead of the
+manual pickle-loading workflow, and weights are available as
+`model.safetensors`.
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model = AutoModelForCausalLM.from_pretrained(
+    "IvmeLabs/Ivme-Conversate-v2-Base", trust_remote_code=True, dtype=torch.float32,
+)
+tokenizer = AutoTokenizer.from_pretrained("IvmeLabs/Ivme-Conversate-v2-Base", trust_remote_code=True)
+model.eval()
+
+inputs = tokenizer("Once upon a time, there was a", return_tensors="pt")
+out = model.generate(
+    **inputs, max_new_tokens=200, do_sample=True,
+    temperature=0.8, top_k=50, pad_token_id=tokenizer.pad_token_id,
+)
+print(tokenizer.decode(out[0], skip_special_tokens=True))
+```
+
+`trust_remote_code=True` is required (custom architecture: RoPE + SwiGLU +
+RMSNorm dense decoder). The original `ckpt_final.pt` pickle checkpoint and
+`model/` architecture source remain in this repo unchanged for backwards
+compatibility.
+
+**Note on batch generation:** use left-padding
+(`tokenizer.padding_side = "left"`) — the model doesn't use an explicit
+attention mask over padded positions, so right-padding within a batch will
+give incorrect results.

config.json

@@ -0,0 +1,26 @@
+{
+  "architectures": [
+    "IvmeForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_ivme.IvmeConfig",
+    "AutoModelForCausalLM": "modeling_ivme.IvmeForCausalLM"
+  },
+  "context_len": 1024,
+  "dropout": 0.0,
+  "dtype": "float32",
+  "ffn_mult": 4.0,
+  "hidden_dim": 384,
+  "hidden_size": 384,
+  "max_position_embeddings": 1024,
+  "model_type": "ivme",
+  "n_heads": 6,
+  "n_layers": 10,
+  "norm_eps": 1e-05,
+  "num_attention_heads": 6,
+  "num_hidden_layers": 10,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.13.1",
+  "vocab_size": 16000
+}

configuration_ivme.py

@@ -0,0 +1,45 @@
+"""HuggingFace Transformers config for Ivme-Conversate-v2."""
+
+from transformers import PretrainedConfig
+
+
+class IvmeConfig(PretrainedConfig):
+    model_type = "ivme"
+
+    def __init__(
+        self,
+        vocab_size: int = 16_000,
+        hidden_dim: int = 384,
+        n_layers: int = 10,
+        n_heads: int = 6,
+        context_len: int = 1024,
+        ffn_mult: float = 4.0,
+        rope_theta: float = 10_000.0,
+        norm_eps: float = 1e-5,
+        tie_embeddings: bool = True,
+        dropout: float = 0.0,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_dim = hidden_dim
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.context_len = context_len
+        self.ffn_mult = ffn_mult
+        self.rope_theta = rope_theta
+        self.norm_eps = norm_eps
+        self.dropout = dropout
+
+        assert hidden_dim % n_heads == 0, "hidden_dim must be divisible by n_heads"
+
+        self.max_position_embeddings = context_len
+        self.num_hidden_layers = n_layers
+        self.num_attention_heads = n_heads
+        self.hidden_size = hidden_dim
+
+        kwargs.setdefault("tie_word_embeddings", tie_embeddings)
+        super().__init__(**kwargs)
+
+    @property
+    def head_dim(self) -> int:
+        return self.hidden_dim // self.n_heads

generation_config.json

@@ -0,0 +1,9 @@
+{
+  "bos_token_id": 0,
+  "eos_token_id": 0,
+  "pad_token_id": 1,
+  "do_sample": true,
+  "temperature": 0.8,
+  "top_k": 50,
+  "max_new_tokens": 200
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:30faf39956673139dc2321a95624c14bc1345ec62f685b71366802dcd4c677f3
+size 119972856

modeling_ivme.py

@@ -0,0 +1,229 @@
+"""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"]

tokenizer_config.json

@@ -0,0 +1,9 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|endoftext|>",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<|unk|>"
+}