Upload folder using huggingface_hub

README.md

@@ -12,11 +12,11 @@ tags:
 - small-language-model
 pipeline_tag: text-generation
 model-index:
-- name: Dwarf-15M
+- name: Dwarf-15M-Instruct
   results: []
 ---
 
-# Dwarf-15M
+# Dwarf-15M-Instruct
 
 A **15.54M parameter** shell/bash specialist language model that translates natural language into Linux commands.
 
@@ -25,8 +25,8 @@ A **15.54M parameter** shell/bash specialist language model that translates natu
 ```python
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-tokenizer = AutoTokenizer.from_pretrained("ThingAI/Dwarf-15M", trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained("ThingAI/Dwarf-15M", trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("ThingAI/Dwarf-15M-Instruct", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("ThingAI/Dwarf-15M-Instruct", trust_remote_code=True)
 
 prompt = "<|user|>\nFind all Python files modified in the last 3 days\n<|end|>\n<|assistant|>\n"
 inputs = tokenizer(prompt, return_tensors="pt")
@@ -79,7 +79,7 @@ Dwarf-15M takes natural language descriptions of tasks and produces the correspo
 - English: helpful-instructions, FineWeb — 10.3%
 - CoT: Magpie-Reasoning — 1.1%
 
-**SFT:** 900 curated Linux command pairs, 5 epochs, lr=4e-5. Training time: 19 seconds.
+**SFT:** 557 curated Linux command pairs, 5 epochs, lr=4e-5. Training time: 19 seconds.
 
 **Tokenizer:** [DwarfGoToken](https://huggingface.co/ThingAI/DwarfGoToken) — 8,202 token BPE with syntax-aware pre-tokenization for shell operators (2>&1, &&, >>).
 
@@ -122,7 +122,7 @@ Apache 2.0
 
 ```bibtex
 @misc{dwarf15m2026,
-  title={Dwarf-15M: A Shell Specialist Language Model},
+  title={Dwarf-15M-Instruct: A Shell Specialist Language Model},
   author={ThingsAI},
   year={2026},
   url={https://huggingface.co/ThingAI/Dwarf-15M-Instruct}

config.json

@@ -19,6 +19,6 @@
   "head_dim": 64,
   "torch_dtype": "float32",
   "transformers_version": "4.45.0",
-  "bos_token_id": 1,
-  "eos_token_id": 2
+  "bos_token_id": 0,
+  "eos_token_id": 0
 }

configuration_dwarf.py

@@ -4,6 +4,7 @@ from transformers import PretrainedConfig
 
 class DwarfConfig(PretrainedConfig):
     model_type = "dwarf"
+    has_no_defaults_at_init = True
 
     def __init__(
         self,
@@ -29,8 +30,4 @@ class DwarfConfig(PretrainedConfig):
         self.rope_theta = rope_theta
         self.norm_eps = norm_eps
         self.head_dim = head_dim
-        self.num_hidden_layers = n_layers
-        self.hidden_size = d_model
-        self.num_attention_heads = n_heads
-        self.num_key_value_heads = n_kv_heads
         super().__init__(**kwargs)

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ed70fe03dec2a2fd6469ca337f72cf222d5375024fdeb8ddb51b5e3bad8f76a3
-size 72674368
+oid sha256:748cb8e54a65afc49e9a8e6e76760dd2206c98bb419c88f56cfbc429fe67fdc7
+size 62175720

modeling_dwarf.py

@@ -1,4 +1,5 @@
 """Dwarf-15M: a 15.54M parameter shell/bash specialist language model."""
+import math
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -9,80 +10,67 @@ from .configuration_dwarf import DwarfConfig
 class RMSNorm(nn.Module):
     def __init__(self, dim, eps=1e-5):
         super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim))
         self.eps = eps
-        self.scale = nn.Parameter(torch.ones(dim))
 
     def forward(self, x):
-        rms = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
-        return (x.float() * rms).to(x.dtype) * self.scale
+        return x * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps).to(x.dtype) * self.weight
 
 
 class RotaryEmbedding(nn.Module):
-    def __init__(self, head_dim, max_seq_len, theta=10000.0):
+    def __init__(self, dim, max_seq_len=2048, theta=10000.0):
         super().__init__()
-        assert head_dim % 2 == 0
-        self.head_dim = head_dim
-        self.max_seq_len = max_seq_len
-        self.theta = theta
-        self.cos_cache = None
-        self.sin_cache = None
-        self._max = 0
-
-    def _build_cache(self, seq_len, device):
-        inv_freq = 1.0 / (self.theta ** (torch.arange(0, self.head_dim, 2, device=device).float() / self.head_dim))
-        t = torch.arange(seq_len, device=device).float()
-        freqs = torch.outer(t, inv_freq)
-        emb = torch.cat([freqs, freqs], dim=-1)
-        self.cos_cache = emb.cos()[None, None]
-        self.sin_cache = emb.sin()[None, None]
-        self._max = seq_len
-
-    @staticmethod
-    def _rotate_half(x):
-        x1, x2 = x.chunk(2, dim=-1)
-        return torch.cat([-x2, x1], dim=-1)
-
-    def forward(self, q, k):
-        T = q.size(2)
-        if self.cos_cache is None or T > self._max or self.cos_cache.device != q.device:
-            self._build_cache(max(T, self.max_seq_len), q.device)
-        cos = self.cos_cache[:, :, :T, :]
-        sin = self.sin_cache[:, :, :T, :]
-        q = q * cos + self._rotate_half(q) * sin
-        k = k * cos + self._rotate_half(k) * sin
-        return q, k
-
-
-class GroupedQueryAttention(nn.Module):
+        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self._build_cache(max_seq_len)
+
+    def _build_cache(self, seq_len):
+        t = torch.arange(seq_len, dtype=self.inv_freq.dtype)
+        freqs = torch.outer(t, self.inv_freq)
+        self.register_buffer("cos_cache", freqs.cos(), persistent=False)
+        self.register_buffer("sin_cache", freqs.sin(), persistent=False)
+
+    def forward(self, x, offset=0):
+        seq_len = x.shape[1]
+        if offset + seq_len > self.cos_cache.shape[0]:
+            self._build_cache(offset + seq_len)
+        cos = self.cos_cache[offset:offset + seq_len]
+        sin = self.sin_cache[offset:offset + seq_len]
+        return cos, sin
+
+
+def apply_rope(x, cos, sin):
+    d = x.shape[-1]
+    x1, x2 = x[..., :d//2], x[..., d//2:]
+    cos = cos[:x.shape[-2], :d//2].unsqueeze(0).unsqueeze(0)
+    sin = sin[:x.shape[-2], :d//2].unsqueeze(0).unsqueeze(0)
+    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)
+
+
+class GQAAttention(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.n_heads = config.n_heads
         self.n_kv_heads = config.n_kv_heads
-        self.n_groups = config.n_heads // config.n_kv_heads
         self.head_dim = config.head_dim
-
-        self.q_proj = nn.Linear(config.d_model, config.n_heads * config.head_dim, bias=True)
-        self.k_proj = nn.Linear(config.d_model, config.n_kv_heads * config.head_dim, bias=True)
-        self.v_proj = nn.Linear(config.d_model, config.n_kv_heads * config.head_dim, bias=True)
+        self.q_proj = nn.Linear(config.d_model, config.n_heads * config.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.d_model, config.n_kv_heads * config.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.d_model, config.n_kv_heads * config.head_dim, bias=False)
         self.o_proj = nn.Linear(config.n_heads * config.head_dim, config.d_model, bias=False)
+        self.group_size = config.n_heads // config.n_kv_heads
 
-        self.rope = RotaryEmbedding(config.head_dim, config.max_seq_len, config.rope_theta)
-
-    def forward(self, x):
+    def forward(self, x, cos, sin):
         B, T, _ = 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_kv_heads, self.head_dim).transpose(1, 2)
         v = self.v_proj(x).view(B, T, self.n_kv_heads, self.head_dim).transpose(1, 2)
-
-        q, k = self.rope(q, k)
-
-        if self.n_groups > 1:
-            k = k.repeat_interleave(self.n_groups, dim=1)
-            v = v.repeat_interleave(self.n_groups, dim=1)
-
-        out = F.scaled_dot_product_attention(q, k, v, attn_mask=None, is_causal=True)
-        out = out.transpose(1, 2).contiguous().view(B, T, self.n_heads * self.head_dim)
-        return self.o_proj(out)
+        q = apply_rope(q, cos, sin)
+        k = apply_rope(k, cos, sin)
+        if self.group_size > 1:
+            k = k.unsqueeze(2).expand(-1, -1, self.group_size, -1, -1).reshape(B, self.n_heads, T, self.head_dim)
+            v = v.unsqueeze(2).expand(-1, -1, self.group_size, -1, -1).reshape(B, self.n_heads, T, self.head_dim)
+        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        return self.o_proj(out.transpose(1, 2).reshape(B, T, -1))
 
 
 class SwiGLUFFN(nn.Module):
@@ -99,20 +87,19 @@ class SwiGLUFFN(nn.Module):
 class DwarfBlock(nn.Module):
     def __init__(self, config):
         super().__init__()
-        self.norm_attn = RMSNorm(config.d_model, config.norm_eps)
-        self.attn = GroupedQueryAttention(config)
-        self.norm_ffn = RMSNorm(config.d_model, config.norm_eps)
+        self.attn_norm = RMSNorm(config.d_model, config.norm_eps)
+        self.attn = GQAAttention(config)
+        self.ff_norm = RMSNorm(config.d_model, config.norm_eps)
         self.ffn = SwiGLUFFN(config)
 
-    def forward(self, x):
-        x = x + self.attn(self.norm_attn(x))
-        x = x + self.ffn(self.norm_ffn(x))
+    def forward(self, x, cos, sin):
+        x = x + self.attn(self.attn_norm(x), cos, sin)
+        x = x + self.ffn(self.ff_norm(x))
         return x
 
 
 class DwarfForCausalLM(PreTrainedModel, GenerationMixin):
     config_class = DwarfConfig
-    _tied_weights_keys = ["lm_head.weight"]
 
     def __init__(self, config):
         super().__init__(config)
@@ -120,27 +107,15 @@ class DwarfForCausalLM(PreTrainedModel, GenerationMixin):
         self.layers = nn.ModuleList([DwarfBlock(config) for _ in range(config.n_layers)])
         self.norm = RMSNorm(config.d_model, config.norm_eps)
         self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.lm_head.weight = self.embed_tokens.weight  # weight tying
+        self.rope = RotaryEmbedding(config.head_dim, config.max_seq_len, config.rope_theta)
         self.post_init()
 
-    def tie_weights(self, **kwargs):
-        self.lm_head.weight = self.embed_tokens.weight
-
-    def get_input_embeddings(self):
-        return self.embed_tokens
-
-    def set_input_embeddings(self, value):
-        self.embed_tokens = 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, labels=None, **kwargs):
         x = self.embed_tokens(input_ids)
+        cos, sin = self.rope(x)
         for layer in self.layers:
-            x = layer(x)
+            x = layer(x, cos, sin)
         logits = self.lm_head(self.norm(x))
 
         loss = None
@@ -155,4 +130,7 @@ class DwarfForCausalLM(PreTrainedModel, GenerationMixin):
         return CausalLMOutput(loss=loss, logits=logits)
 
     def prepare_inputs_for_generation(self, input_ids, **kwargs):
-        return {"input_ids": input_ids}
+        return {"input_ids": input_ids}
+
+    def count_parameters(self):
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)