Upload folder using huggingface_hub

chat_template.jinja

@@ -0,0 +1 @@
+{% for m in messages %}{{ 'User: ' + m['content'].strip() + '\n' if m['role'] == 'user' else 'Assistant: ' + m['content'].strip() + '\n' }}{% endfor %}{% if add_generation_prompt %}Assistant:{% endif %}

config.json

@@ -0,0 +1,23 @@
+{
+  "architectures": [
+    "MiniLlmForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_mini_llm.MiniLlmConfig",
+    "AutoModelForCausalLM": "modeling_mini_llm.MiniLlmForCausalLM"
+  },
+  "dtype": "float32",
+  "eos_token_id": 0,
+  "hidden_size": 512,
+  "intermediate_size": 1408,
+  "max_position_embeddings": 1024,
+  "model_type": "mini_llm",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 8,
+  "pad_token_id": 0,
+  "rms_norm_eps": 1e-05,
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.19.0",
+  "vocab_size": 32768
+}

configuration_mini_llm.py

@@ -0,0 +1,45 @@
+from transformers import PretrainedConfig
+
+
+class MiniLlmConfig(PretrainedConfig):
+    model_type = "mini_llm"
+
+    def __init__(
+        self,
+        vocab_size=16384,
+        hidden_size=512,
+        num_hidden_layers=8,
+        num_attention_heads=8,
+        intermediate_size=1408,
+        max_position_embeddings=512,
+        rms_norm_eps=1e-5,
+        rope_theta=10000.0,
+        tie_word_embeddings=True,
+        **kwargs,
+    ):
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            eos_token_id=kwargs.pop("eos_token_id", 0),
+            pad_token_id=kwargs.pop("pad_token_id", 0),
+            **kwargs,
+        )
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.max_position_embeddings = max_position_embeddings
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_theta = rope_theta
+
+    @classmethod
+    def from_cfg(cls, cfg, **kwargs):
+        return cls(
+            vocab_size=cfg.vocab,
+            hidden_size=cfg.dim,
+            num_hidden_layers=cfg.layers,
+            num_attention_heads=cfg.heads,
+            intermediate_size=cfg.ffn,
+            max_position_embeddings=cfg.ctx,
+            **kwargs,
+        )

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "do_sample": true,
+  "eos_token_id": 0,
+  "max_new_tokens": 256,
+  "pad_token_id": 0,
+  "repetition_penalty": 1.15,
+  "temperature": 0.5,
+  "top_k": 40,
+  "transformers_version": "5.19.0"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:666c351a69d7ae9a2ca2b349833e111fd9581ed54879638b18e9de8656728df6
+size 170171728

modeling_mini_llm.py

@@ -0,0 +1,151 @@
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+try:
+    from .configuration_mini_llm import MiniLlmConfig
+except ImportError:
+    from configuration_mini_llm import MiniLlmConfig
+
+
+class MiniLlmRMSNorm(nn.Module):
+    def __init__(self, d, eps=1e-5):
+        super().__init__()
+        self.w = nn.Parameter(torch.ones(d))
+        self.eps = eps
+
+    def forward(self, x):
+        n = torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
+        return (x.float() * n).type_as(x) * self.w
+
+
+def mini_rope_cache(T, hd, device, theta=10000.0):
+    inv = 1.0 / (theta ** (torch.arange(0, hd, 2, device=device).float() / hd))
+    f = torch.outer(torch.arange(T, device=device).float(), inv)
+    return f.cos(), f.sin()
+
+
+def mini_apply_rope(x, cos, sin):
+    x1, x2 = x[..., ::2], x[..., 1::2]
+    cos, sin = cos[: x.size(2)], sin[: x.size(2)]
+    out = torch.stack([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+    return out.flatten(-2).type_as(x)
+
+
+class MiniLlmAttention(nn.Module):
+    def __init__(self, config: MiniLlmConfig):
+        super().__init__()
+        d = config.hidden_size
+        self.h = config.num_attention_heads
+        self.qkv = nn.Linear(d, 3 * d, bias=False)
+        self.o = nn.Linear(d, d, bias=False)
+
+    def forward(self, x, cos, sin):
+        B, T, C = x.shape
+        q, k, v = self.qkv(x).view(B, T, 3, self.h, C // self.h).permute(2, 0, 3, 1, 4)
+        q, k, v = q.contiguous(), k.contiguous(), v.contiguous()
+        q, k = mini_apply_rope(q, cos, sin), mini_apply_rope(k, cos, sin)
+        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        return self.o(y.transpose(1, 2).reshape(B, T, C))
+
+
+class MiniLlmSwiGLU(nn.Module):
+    def __init__(self, config: MiniLlmConfig):
+        super().__init__()
+        d, ffn = config.hidden_size, config.intermediate_size
+        self.w1 = nn.Linear(d, ffn, bias=False)
+        self.w3 = nn.Linear(d, ffn, bias=False)
+        self.w2 = nn.Linear(ffn, d, bias=False)
+
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class MiniLlmBlock(nn.Module):
+    def __init__(self, config: MiniLlmConfig):
+        super().__init__()
+        self.n1, self.n2 = MiniLlmRMSNorm(config.hidden_size, config.rms_norm_eps), MiniLlmRMSNorm(
+            config.hidden_size, config.rms_norm_eps
+        )
+        self.attn, self.mlp = MiniLlmAttention(config), MiniLlmSwiGLU(config)
+
+    def forward(self, x, cos, sin):
+        x = x + self.attn(self.n1(x), cos, sin)
+        return x + self.mlp(self.n2(x))
+
+
+class MiniLlmForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = MiniLlmConfig
+    base_model_prefix = "mini_llm"
+    supports_gradient_checkpointing = False
+    _tied_weights_keys = {"head.weight": "emb.weight"}
+
+    def __init__(self, config: MiniLlmConfig):
+        super().__init__(config)
+        self.config = config
+        self.emb = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.blocks = nn.ModuleList(MiniLlmBlock(config) for _ in range(config.num_hidden_layers))
+        self.norm = MiniLlmRMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.head.weight = self.emb.weight
+        cos, sin = mini_rope_cache(
+            config.max_position_embeddings,
+            config.hidden_size // config.num_attention_heads,
+            "cpu",
+            config.rope_theta,
+        )
+        self.register_buffer("cos", cos, persistent=True)
+        self.register_buffer("sin", sin, persistent=True)
+        self.post_init()
+        for n, p in self.named_parameters():
+            if n.endswith("attn.o.weight") or n.endswith("mlp.w2.weight"):
+                nn.init.normal_(p, 0.0, 0.02 / math.sqrt(2 * config.num_hidden_layers))
+
+    def _init_weights(self, module):
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            nn.init.normal_(module.weight, 0.0, 0.02)
+
+    def get_input_embeddings(self):
+        return self.emb
+
+    def set_input_embeddings(self, value):
+        self.emb = value
+
+    def get_output_embeddings(self):
+        return self.head
+
+    def set_output_embeddings(self, value):
+        self.head = value
+
+    def forward(self, input_ids=None, labels=None, attention_mask=None, **kwargs):
+        kwargs.pop("past_key_values", None)
+        kwargs.pop("use_cache", None)
+        kwargs.pop("position_ids", None)
+        kwargs.pop("inputs_embeds", None)
+        if input_ids is None:
+            raise ValueError("input_ids is required")
+        x = self.emb(input_ids)
+        for b in self.blocks:
+            x = b(x, self.cos, self.sin)
+        logits = self.head(self.norm(x))
+        loss = None
+        if labels is not None:
+            loss = F.cross_entropy(logits.float().view(-1, logits.size(-1)), labels.view(-1))
+        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=None)
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        return {"input_ids": input_ids}
+
+
+try:
+    from transformers import AutoConfig, AutoModelForCausalLM
+
+    AutoConfig.register("mini_llm", MiniLlmConfig)
+    AutoModelForCausalLM.register(MiniLlmConfig, MiniLlmForCausalLM)
+except Exception:
+    pass

tokenizer_config.json

@@ -0,0 +1,9 @@
+{
+  "backend": "tokenizers",
+  "bos_token": null,
+  "eos_token": "<eos>",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "<eos>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": null
+}