Upload fine-tuned Pebble-10M-Chat model

config.json

@@ -0,0 +1,32 @@
+{
+  "architectures": [
+    "Pebble10MLM"
+  ],
+  "attention": {
+    "is_causal": true,
+    "rope_theta": 10000.0
+  },
+  "auto_map": {
+    "AutoConfig": "configuration_pebble.PebbleConfig",
+    "AutoModelForCausalLM": "modeling_pebble.PebbleForCausalLM"
+  },
+  "block_pattern": "mmma|mmma",
+  "dtype": "float32",
+  "hidden_size": 384,
+  "hybrid_ratio": "3:1 mamba2:attention",
+  "intermediate_size": 1536,
+  "mamba2": {
+    "d_conv": 4,
+    "d_state": 128,
+    "expand": 2,
+    "headdim": 96,
+    "use_mem_eff_path": true
+  },
+  "max_position_embeddings": 512,
+  "model_type": "pebble_10m",
+  "num_attention_heads": 6,
+  "num_hidden_layers": 8,
+  "rms_norm_eps": 1e-06,
+  "transformers_version": "4.57.1",
+  "vocab_size": 2048
+}

configuration_pebble.py

@@ -0,0 +1,41 @@
+from transformers import PretrainedConfig
+
+class PebbleConfig(PretrainedConfig):
+    model_type = "pebble_10m"
+
+    def __init__(
+        self,
+        vocab_size=2048,
+        hidden_size=384,
+        intermediate_size=1536,
+        num_hidden_layers=8,
+        num_attention_heads=6,
+        block_pattern="mmma|mmma",
+        hybrid_ratio="3:1 mamba2:attention",
+        max_position_embeddings=512,
+        rms_norm_eps=1e-6,
+        tie_word_embeddings=True,
+        mamba2=None,
+        attention=None,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.block_pattern = block_pattern
+        self.hybrid_ratio = hybrid_ratio
+        self.max_position_embeddings = max_position_embeddings
+        self.rms_norm_eps = rms_norm_eps
+        self.tie_word_embeddings = tie_word_embeddings
+        
+        # Default dictionaries if not provided in config.json
+        self.mamba2 = mamba2 or {
+            "d_state": 128, "d_conv": 4, "expand": 2, 
+            "headdim": 96, "use_mem_eff_path": True
+        }
+        self.attention = attention or {
+            "rope_theta": 10000.0, "is_causal": True
+        }
+        super().__init__(**kwargs)

generation_config.json

@@ -0,0 +1,4 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "4.57.1"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:457d3a32ae269f8cb7be5fc155b6425ad8a8fa82c708bb8c3787cf5861f80b9d
+size 44278976

modeling_pebble.py

@@ -0,0 +1,147 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+try:
+    from mamba_ssm import Mamba2
+except ImportError:
+    raise ImportError("mamba-ssm is required. pip install mamba-ssm causal-conv1d")
+
+from .configuration_pebble import PebbleConfig
+
+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):
+        dt = x.dtype
+        xf = x.float()
+        xf = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + self.eps)
+        return self.weight * xf.to(dt)
+
+class AttentionBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        dim = config.hidden_size
+        n_heads = config.num_attention_heads
+        hidden = config.intermediate_size
+        assert dim % n_heads == 0
+        self.nh, self.hd = n_heads, dim // n_heads
+        self.wqkv = nn.Linear(dim, 3 * dim, bias=False)
+        self.wo   = nn.Linear(dim, dim, bias=False)
+        self.fc1  = nn.Linear(dim, hidden, bias=False)
+        self.fc2  = nn.Linear(hidden, dim, bias=False)
+        self.ln1  = RMSNorm(dim, eps=config.rms_norm_eps)
+        self.ln2  = RMSNorm(dim, eps=config.rms_norm_eps)
+        self.rope_theta = config.attention.get("rope_theta", 10000.0)
+
+    def forward(self, x):
+        B, T, C = x.shape
+        h = self.ln1(x)
+
+        qkv = self.wqkv(h).view(B, T, 3, self.nh, self.hd) \
+                          .permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0].float(), qkv[1].float(), qkv[2]
+
+        half = self.hd // 2
+        invf = 1.0 / (self.rope_theta ** (
+            torch.arange(0, half, device=x.device, dtype=torch.float32)
+            * 2.0 / self.hd))
+        ang = torch.outer(
+            torch.arange(T, device=x.device, dtype=torch.float32), invf)
+        cos, sin = ang.cos()[None, None], ang.sin()[None, None]
+
+        q1, q2 = q[..., :half], q[..., half:]
+        k1, k2 = k[..., :half], k[..., half:]
+        q = torch.cat([q1 * cos - q2 * sin,
+                       q1 * sin + q2 * cos], dim=-1).to(v.dtype)
+        k = torch.cat([k1 * cos - k2 * sin,
+                       k1 * sin + k2 * cos], dim=-1).to(v.dtype)
+
+        y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        y = y.transpose(1, 2).reshape(B, T, C)
+
+        x = x + self.wo(y)
+        x = x + self.fc2(F.gelu(self.fc1(self.ln2(x))))
+        return x
+
+class MambaBlock(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.ln = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        mamba_cfg = config.mamba2
+        self.mixer = Mamba2(
+            d_model=config.hidden_size,
+            d_state=mamba_cfg.get("d_state", 128),
+            d_conv=mamba_cfg.get("d_conv", 4),
+            expand=mamba_cfg.get("expand", 2),
+            headdim=mamba_cfg.get("headdim", 96),
+            use_mem_eff_path=mamba_cfg.get("use_mem_eff_path", True),
+        )
+
+    def forward(self, x):
+        return x + self.mixer(self.ln(x))
+
+class PebbleForCausalLM(PreTrainedModel):
+    config_class = PebbleConfig
+    supports_gradient_checkpointing = False
+    _no_split_modules = ["MambaBlock", "AttentionBlock"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        
+        self.wte = nn.Embedding(config.vocab_size, config.hidden_size)
+        
+        # 3:1 Mamba:Attention ratio layout
+        self.blocks = nn.ModuleList([
+            MambaBlock(config) if i % 4 < 3
+            else AttentionBlock(config)
+            for i in range(config.num_hidden_layers)
+        ])
+        
+        self.lnf = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        
+        # Tie weights
+        self.tie_weights()
+
+    def tie_weights(self):
+        if self.config.tie_word_embeddings:
+            self.lm_head.weight = self.wte.weight
+
+    def forward(self, input_ids=None, attention_mask=None, labels=None, past_key_values=None, **kwargs):
+        x = self.wte(input_ids)
+        
+        for blk in self.blocks:
+            x = blk(x)
+            
+        logits = self.lm_head(self.lnf(x))
+
+        loss = None
+        if labels is not None:
+            # Shift so that tokens < n predict n+1
+            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)
+            )
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=past_key_values,
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, **kwargs):
+        # Mamba handles state internally in the mixer, so we don't use past_key_values 
+        # at the model level for now (standard HF generation will still work for greedy/beam).
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+        }

special_tokens_map.json

@@ -0,0 +1,16 @@
+{
+  "bos_token": {
+    "content": "<|eos|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|eos|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,27 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<|eos|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|eos|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|eos|>",
+  "eos_token_id": 0,
+  "extra_special_tokens": {},
+  "model_input_names": [
+    "input_ids"
+  ],
+  "model_max_length": 512,
+  "pad_token": null,
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "unk_token": null,
+  "vocab_size": 2048
+}