Upload folder using huggingface_hub

README.md

@@ -12,13 +12,7 @@ tags:
 - on-device
 pipeline_tag: text-generation
 widget:
-- text: '<|im_start|>user
-
-    What is the capital of France?<|im_end|>
-
-    <|im_start|>assistant
-
-    '
+- text: "<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"
 ---
 
 # PicoLM-80M-Instruct 🚀
@@ -35,7 +29,7 @@ Trained completely from scratch on Kaggle dual Tesla T4 GPUs with zero budget, P
 - **Parameters:** 80,242,240 (~80.2M)
 - **Context Window:** 2,048 tokens
 - **Vocabulary:** 16,384 (Single-digit regex split, Byte-level BPE)
-- **Format:** Safetensors (FP16)
+- **Format:** Safetensors (FP16) & GGUF
 - **Primary Language:** English + Python Code
 - **License:** Apache 2.0
 
@@ -43,7 +37,7 @@ Trained completely from scratch on Kaggle dual Tesla T4 GPUs with zero budget, P
 
 ## 📊 Empirical Benchmark Results (Verified)
 
-Unlike models with estimated metrics, all scores below were **empirically measured** directly on the model weights using standard log-likelihood evaluations:
+All scores below were **empirically measured** directly on the model weights using standard log-likelihood evaluations:
 
 | Benchmark / Task | Random Baseline | SmolLM2-135M (HF) | Gemma 3 270M (Google) | PicoLM-80M-Instruct (Ours) |
 | :--- | :--- | :--- | :--- | :--- |
@@ -55,54 +49,28 @@ Unlike models with estimated metrics, all scores below were **empirically measur
 
 ---
 
-## 🏗️ Architecture Specifications
-
-- **Layers:** 20 Transformer blocks (Deep & Thin)
-- **Hidden Dimension ($d_{model}$):** 576
-- **Attention Heads (Q):** 9 (Head dim = 64)
-- **KV Heads (GQA):** 3 (3:1 Grouped-Query Attention)
-- **Intermediate Dimension:** 1536 (SwiGLU)
-- **Normalization:** Pre-RMSNorm with **QK-Norm** (prevents attention logit blowup)
-- **Positional Encoding:** RoPE ($	heta = 10000$)
-- **Weights Tying:** Input embeddings and LM Head weights are shared (Zero parameter waste)
-
----
-
-## 📚 Training Curriculum
-
-1. **Pre-Training (Stage 1 - ~307M Tokens):**
-   - 54% Cosmopedia v2 (Textbooks & reasoning)
-   - 23% FineWeb-Edu Dedup (Educational web articles)
-   - 12% Python-Edu (Algorithms and code syntax)
-   - 11% SmolTalk Dialogues
-   - Optimization: WSD (Warmup-Stable-Decay) across 2x Tesla T4 GPUs.
-2. **Chat Alignment (Stage 2 - SFT):**
-   - 15,000 curated dialogues from SmolTalk.
-   - Strict Prompt-Masking (`ignore_index = -100`) applied so only assistant responses receive loss updates.
-
----
+## 💻 Quickstart (Transformers Native)
 
-## 💻 Quickstart Inference
+You can load and chat with PicoLM directly via Hugging Face `transformers`:
 
 ```python
-import torch
-from transformers import PreTrainedTokenizerFast
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
-# Load tokenizer
-tokenizer = PreTrainedTokenizerFast.from_pretrained("aethertp/PicoLM-80M-Instruct")
+model_id = "aethertp/PicoLM-80M-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True).cuda()
 
-# Format prompt using ChatML
-user_prompt = "What is the capital of France?"
-formatted = f"<|im_start|>user\n{user_prompt}<|im_end|>\n<|im_start|>assistant\n"
-input_ids = torch.tensor([tokenizer.encode(formatted, add_special_tokens=False)])
+messages = [{"role": "user", "content": "What is the capital of France?"}]
+prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
 
-# Generate response
-# (Model weights available in model.safetensors)
+outputs = model.generate(**inputs, max_new_tokens=50, temperature=0.6, do_sample=True)
+print(tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:]))
 ```
 
 ---
 
 ## ⚠️ Limitations
 
-- **Factual Depth:** With 80M parameters, the model cannot serve as a reliable encyclopedia for obscure facts. Factual queries should be supported by RAG.
-- **Multi-step Math:** Elementary arithmetic works, but multi-variable algebra requires external verification.
+- **Factual Depth:** With 80M parameters, the model cannot serve as a comprehensive encyclopedia. Factual queries should be supported by RAG.
+- **Multi-step Math:** Elementary arithmetic works, but complex multi-variable algebra requires external verification.

config.json

@@ -3,6 +3,10 @@
     "PicoLMForCausalLM"
   ],
   "model_type": "picolm",
+  "auto_map": {
+    "AutoConfig": "configuration_picolm.PicoLMConfig",
+    "AutoModelForCausalLM": "modeling_picolm.PicoLMForCausalLM"
+  },
   "vocab_size": 16384,
   "hidden_size": 576,
   "num_hidden_layers": 20,
@@ -12,6 +16,7 @@
   "max_position_embeddings": 2048,
   "rms_norm_eps": 1e-05,
   "rope_theta": 10000.0,
+  "tie_word_embeddings": true,
   "torch_dtype": "float16",
   "transformers_version": "4.45.0"
 }

configuration_picolm.py

@@ -0,0 +1,28 @@
+from transformers import PretrainedConfig
+
+class PicoLMConfig(PretrainedConfig):
+    model_type = "picolm"
+    def __init__(
+        self,
+        vocab_size=16384,
+        hidden_size=576,
+        num_hidden_layers=20,
+        num_attention_heads=9,
+        num_key_value_heads=3,
+        intermediate_size=1536,
+        max_position_embeddings=2048,
+        rms_norm_eps=1e-5,
+        rope_theta=10000.0,
+        tie_word_embeddings=True,
+        **kwargs
+    ):
+        super().__init__(tie_word_embeddings=tie_word_embeddings, **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.num_key_value_heads = num_key_value_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

modeling_picolm.py

@@ -0,0 +1,105 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from .configuration_picolm import PicoLMConfig
+
+class RMSNorm(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):
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps) * self.weight
+
+def precompute_rope_cis(dim: int, max_seq_len: int, theta: float = 10000.0):
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(max_seq_len, dtype=torch.float32)
+    freqs = torch.outer(t, freqs)
+    return torch.cos(freqs), torch.sin(freqs)
+
+def apply_rope(x, cos, sin):
+    B, H, S, D = x.shape
+    cos = cos[:S, :].to(x.device).unsqueeze(0).unsqueeze(0)
+    sin = sin[:S, :].to(x.device).unsqueeze(0).unsqueeze(0)
+    x1, x2 = x[..., : D // 2], x[..., D // 2 :]
+    return torch.cat([x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1)
+
+class Attention(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+        self.head_dim = args.hidden_size // args.num_attention_heads
+        self.n_heads = args.num_attention_heads
+        self.n_kv_heads = args.num_key_value_heads
+        self.num_reps = self.n_heads // self.n_kv_heads
+
+        self.q_proj = nn.Linear(args.hidden_size, self.n_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(args.hidden_size, self.n_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(args.hidden_size, self.n_kv_heads * self.head_dim, bias=False)
+        self.out_proj = nn.Linear(args.n_heads * self.head_dim, args.hidden_size, bias=False)
+        self.q_norm = RMSNorm(self.head_dim, eps=args.rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=args.rms_norm_eps)
+
+    def forward(self, x, cos, sin):
+        B, S, C = x.shape
+        q = apply_rope(self.q_norm(self.q_proj(x).view(B, S, self.n_heads, self.head_dim).transpose(1, 2)), cos, sin)
+        k = apply_rope(self.k_norm(self.k_proj(x).view(B, S, self.n_kv_heads, self.head_dim).transpose(1, 2)), cos, sin)
+        v = self.v_proj(x).view(B, S, self.n_kv_heads, self.head_dim).transpose(1, 2)
+        if self.num_reps > 1:
+            k = k[:, :, None, :, :].expand(B, self.n_kv_heads, self.num_reps, S, self.head_dim).reshape(B, self.n_heads, S, self.head_dim)
+            v = v[:, :, None, :, :].expand(B, self.n_kv_heads, self.num_reps, S, self.head_dim).reshape(B, self.n_heads, S, self.head_dim)
+        attn_out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        return self.out_proj(attn_out.transpose(1, 2).contiguous().view(B, S, C))
+
+class SwiGLUMLP(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.gate_proj = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(args.hidden_size, args.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(args.intermediate_size, args.hidden_size, bias=False)
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+class TransformerBlock(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.attn_norm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.attn = Attention(args)
+        self.mlp_norm = RMSNorm(args.hidden_size, eps=args.rms_norm_eps)
+        self.mlp = SwiGLUMLP(args)
+    def forward(self, x, cos, sin):
+        return x + self.mlp(self.mlp_norm(x + self.attn(self.attn_norm(x), cos, sin)))
+
+class PicoLMForCausalLM(PreTrainedModel):
+    config_class = PicoLMConfig
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.blocks = nn.ModuleList([TransformerBlock(config) for _ in range(config.num_hidden_layers)])
+        self.final_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.lm_head.weight = self.tok_embeddings.weight
+
+        cos, sin = precompute_rope_cis(config.hidden_size // config.num_attention_heads, config.max_position_embeddings, config.rope_theta)
+        self.register_buffer("rope_cos", cos, persistent=False)
+        self.register_buffer("rope_sin", sin, persistent=False)
+
+    def forward(self, input_ids, labels=None, **kwargs):
+        B, S = input_ids.shape
+        x = self.tok_embeddings(input_ids)
+        cos, sin = self.rope_cos[:S], self.rope_sin[:S]
+        for block in self.blocks:
+            x = block(x, cos, sin)
+        logits = self.lm_head(self.final_norm(x))
+
+        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, self.config.vocab_size), shift_labels.view(-1), ignore_index=-100)
+
+        return CausalLMOutputWithPast(loss=loss, logits=logits)