Duplicate from arman-bd/guppylm-9M

Co-authored-by: Arman Hossain <arman-bd@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,37 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text
+assets/guppy.png filter=lfs diff=lfs merge=lfs -text
+model_fp32.onnx.data filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,98 @@
+---
+license: mit
+language:
+  - en
+tags:
+  - fish
+  - character
+  - tiny-llm
+  - text-generation
+  - from-scratch
+pipeline_tag: text-generation
+---
+
+<p align="center">
+  <img src="assets/guppy.png" alt="GuppyLM" width="300"/>
+</p>
+
+<p align="center">
+  <a href="https://github.com/arman-bd/guppylm"><img src="https://img.shields.io/badge/GitHub-guppylm-181717?logo=github" alt="GitHub"/></a>&nbsp;
+  <a href="https://huggingface.co/datasets/arman-bd/guppylm-60k-generic"><img src="https://img.shields.io/badge/🤗_Dataset-guppylm--60k-blue" alt="Dataset"/></a>&nbsp;
+  <a href="https://colab.research.google.com/github/arman-bd/guppylm/blob/main/use_guppylm.ipynb"><img src="https://img.shields.io/badge/Open_in-Colab-F9AB00?logo=googlecolab" alt="Colab"/></a>
+  <br/>
+  <a href="https://www.linkedin.com/pulse/build-your-own-language-model-5-minutes-i-made-mine-hossain--supif/"><img src="https://img.shields.io/badge/Article-LinkedIn-0A66C2?logo=linkedin" alt="LinkedIn Article"/></a>&nbsp;
+  <a href="https://arman-bd.medium.com/build-your-own-llm-in-5-minutes-i-made-mine-talk-like-a-fish-e20c338a3d14"><img src="https://img.shields.io/badge/Article-Medium-000000?logo=medium" alt="Medium Article"/></a>
+  <br/><br/>
+  <a href="https://arman-bd.github.io/guppylm/"><img src="https://img.shields.io/badge/Try_in-Browser-64ffda?style=for-the-badge&logo=webassembly&logoColor=white" alt="Browser Demo"/></a>
+</p>
+
+# GuppyLM — 9M Parameter Fish Chat Model
+
+A ~9M parameter LLM trained from scratch that talks like a small fish named Guppy.
+
+This project exists to show that training your own language model is not magic. One Colab notebook, 5 minutes, and you have a working LLM built from scratch.
+
+## Example
+
+```
+You> hi guppy
+Guppy> hi there. i just found a nice spot near the rock. the temperature feels nice.
+
+You> are you hungry
+Guppy> yes. always yes. i will swim to the top right now. i promise to eat all of it.
+
+You> what is the meaning of life
+Guppy> food. the answer is always food.
+
+You> tell me a joke
+Guppy> what did the fish say when it hit the wall. dam.
+
+You> goodnight guppy
+Guppy> ok sleep time. i was following a bubble but now i'll stop. goodnight tank. goodnight water.
+```
+
+## Architecture
+
+| | |
+|---|---|
+| **Parameters** | 8.7M |
+| **Type** | Vanilla transformer (from scratch) |
+| **Layers** | 6 |
+| **Hidden dim** | 384 |
+| **Heads** | 6 |
+| **FFN** | 768 (ReLU) |
+| **Vocab** | 4,096 (BPE) |
+| **Max sequence** | 128 tokens |
+| **Norm** | LayerNorm |
+| **Position** | Learned embeddings |
+| **LM head** | Weight-tied with embeddings |
+
+No GQA, no RoPE, no SwiGLU, no early exit. As simple as it gets.
+
+## Training
+
+- **Data:** 60K single-turn synthetic conversations across 60 topics
+- **Steps:** 10,000
+- **Optimizer:** AdamW (cosine LR schedule)
+- **Hardware:** T4 GPU (~5 min)
+- **No system prompt** — personality is baked into the weights
+
+## Usage
+
+```python
+from inference import GuppyInference
+
+engine = GuppyInference('checkpoints/best_model.pt', 'data/tokenizer.json')
+r = engine.chat_completion([{'role': 'user', 'content': 'hi guppy'}])
+print(r['choices'][0]['message']['content'])
+# hi there. i just found a nice spot near the rock.
+```
+
+## Links
+
+- **Repo:** [github.com/arman-bd/guppylm](https://github.com/arman-bd/guppylm)
+- **Dataset:** [huggingface.co/datasets/arman-bd/guppylm-60k-generic](https://huggingface.co/datasets/arman-bd/guppylm-60k-generic)
+
+## License
+
+MIT

config.json

@@ -0,0 +1,16 @@
+{
+  "model_type": "guppylm",
+  "architectures": [
+    "GuppyLM"
+  ],
+  "vocab_size": 4096,
+  "max_position_embeddings": 128,
+  "hidden_size": 384,
+  "num_hidden_layers": 6,
+  "num_attention_heads": 6,
+  "intermediate_size": 768,
+  "hidden_dropout_prob": 0.1,
+  "pad_token_id": 0,
+  "bos_token_id": 1,
+  "eos_token_id": 2
+}

config.py

@@ -0,0 +1,36 @@
+"""GuppyLM configuration."""
+
+from dataclasses import dataclass
+
+
+@dataclass
+class GuppyConfig:
+    vocab_size: int = 4096
+    max_seq_len: int = 128
+    d_model: int = 384
+    n_layers: int = 6
+    n_heads: int = 6
+    ffn_hidden: int = 768
+    dropout: float = 0.1
+
+    # Special tokens
+    pad_id: int = 0
+    bos_id: int = 1           # <|im_start|>
+    eos_id: int = 2           # <|im_end|>
+
+
+@dataclass
+class TrainConfig:
+    batch_size: int = 32
+    learning_rate: float = 3e-4
+    min_lr: float = 3e-5
+    weight_decay: float = 0.1
+    warmup_steps: int = 200
+    max_steps: int = 10000
+    eval_interval: int = 200
+    save_interval: int = 500
+    grad_clip: float = 1.0
+    device: str = "auto"
+    seed: int = 42
+    data_dir: str = "data"
+    output_dir: str = "checkpoints"

inference.py

@@ -0,0 +1,124 @@
+"""GuppyLM inference — simple chat."""
+
+import json
+import time
+import uuid
+
+import torch
+from tokenizers import Tokenizer
+
+from config import GuppyConfig
+from model import GuppyLM
+
+
+class GuppyInference:
+    def __init__(self, checkpoint_path, tokenizer_path, device="cpu"):
+        self.device = torch.device(device)
+        self.tokenizer = Tokenizer.from_file(tokenizer_path)
+
+        import os
+        ckpt = torch.load(checkpoint_path, map_location=self.device, weights_only=False)
+
+        # Load config.json from same directory as the model file
+        config_dir = os.path.dirname(os.path.abspath(checkpoint_path))
+        config_path = os.path.join(config_dir, "config.json")
+
+        # Extract state_dict — handle both legacy and standard formats
+        if isinstance(ckpt, dict) and "model_state_dict" in ckpt:
+            state_dict = ckpt["model_state_dict"]
+        else:
+            state_dict = ckpt
+
+        # Load config — try config.json first, fall back to embedded config
+        if os.path.exists(config_path):
+            with open(config_path) as f:
+                cfg = json.load(f)
+            # Support both HF standard keys and our own keys
+            self.config = GuppyConfig(
+                vocab_size=cfg.get("vocab_size", 4096),
+                max_seq_len=cfg.get("max_position_embeddings", cfg.get("max_seq_len", 128)),
+                d_model=cfg.get("hidden_size", cfg.get("d_model", 384)),
+                n_layers=cfg.get("num_hidden_layers", cfg.get("n_layers", 6)),
+                n_heads=cfg.get("num_attention_heads", cfg.get("n_heads", 6)),
+                ffn_hidden=cfg.get("intermediate_size", cfg.get("ffn_hidden", 768)),
+                dropout=cfg.get("hidden_dropout_prob", cfg.get("dropout", 0.1)),
+                pad_id=cfg.get("pad_token_id", cfg.get("pad_id", 0)),
+                bos_id=cfg.get("bos_token_id", cfg.get("bos_id", 1)),
+                eos_id=cfg.get("eos_token_id", cfg.get("eos_id", 2)),
+            )
+        elif isinstance(ckpt, dict) and "config" in ckpt:
+            valid_fields = {f.name for f in GuppyConfig.__dataclass_fields__.values()}
+            self.config = GuppyConfig(**{k: v for k, v in ckpt["config"].items() if k in valid_fields})
+        else:
+            print("Warning: No config found, using defaults")
+            self.config = GuppyConfig()
+
+        self.model = GuppyLM(self.config).to(self.device)
+        filtered = {k: v for k, v in state_dict.items() if k in self.model.state_dict()}
+        self.model.load_state_dict(filtered)
+        self.model.eval()
+
+        total, _ = self.model.param_count()
+        print(f"GuppyLM loaded: {total/1e6:.1f}M params")
+
+    def chat_completion(self, messages, temperature=0.7, max_tokens=64,
+                        top_k=50, **kwargs):
+        """Chat completion — takes messages, returns response."""
+        prompt = self._format_prompt(messages)
+        input_ids = self.tokenizer.encode(prompt).ids
+        prompt_tokens = len(input_ids)
+        input_t = torch.tensor([input_ids], dtype=torch.long, device=self.device)
+
+        output_t, _ = self.model.generate(input_t, max_tokens, temperature, top_k)
+        output_text = self.tokenizer.decode(output_t[0].tolist()[prompt_tokens:])
+        # Truncate at first <|im_end|> — don't let the model leak into the next turn
+        if "<|im_end|>" in output_text:
+            output_text = output_text.split("<|im_end|>")[0]
+        # Also strip any <|im_start|> fragments
+        if "<|im_start|>" in output_text:
+            output_text = output_text.split("<|im_start|>")[0]
+        resp_text = output_text.strip()
+
+        return {
+            "choices": [{
+                "message": {"role": "assistant", "content": resp_text},
+            }],
+        }
+
+    def _format_prompt(self, messages):
+        parts = []
+        for msg in messages:
+            role = msg.get("role", "user")
+            content = msg.get("content") or ""
+            if role == "system":
+                continue
+            parts.append(f"<|im_start|>{role}\n{content}<|im_end|>")
+        parts.append("<|im_start|>assistant\n")
+        return "\n".join(parts)
+
+
+def main():
+    import argparse
+    p = argparse.ArgumentParser(description="Chat with Guppy")
+    p.add_argument("--checkpoint", default="checkpoints/best_model.pt")
+    p.add_argument("--tokenizer", default="data/tokenizer.json")
+    p.add_argument("--device", default="cpu")
+    args = p.parse_args()
+
+    engine = GuppyInference(args.checkpoint, args.tokenizer, args.device)
+    print("\nGuppy Chat (type 'quit' to exit)")
+    msgs = []
+    while True:
+        inp = input("\nYou> ").strip()
+        if inp.lower() in ("quit", "exit", "q"):
+            break
+        msgs.append({"role": "user", "content": inp})
+        result = engine.chat_completion(msgs)
+        msg = result["choices"][0]["message"]
+        if msg.get("content"):
+            print(f"Guppy> {msg['content']}")
+        msgs.append(msg)
+
+
+if __name__ == "__main__":
+    main()

model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:96dd79d233662b65bc0fae1773821479551d1c1257ca5bbd12fa3e6623ae4835
+size 10469869

model.py

@@ -0,0 +1,129 @@
+"""
+GuppyLM — a tiny fish brain.
+
+Vanilla transformer: multi-head attention, ReLU FFN, LayerNorm, learned positional embeddings.
+No GQA, no SwiGLU, no parallel residual, no RoPE. As simple as it gets.
+"""
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from config import GuppyConfig
+
+
+class Attention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.n_heads = config.n_heads
+        self.head_dim = config.d_model // config.n_heads
+
+        self.qkv = nn.Linear(config.d_model, 3 * config.d_model)
+        self.out = nn.Linear(config.d_model, config.d_model)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x, mask=None):
+        B, T, C = x.shape
+        qkv = self.qkv(x).reshape(B, T, 3, self.n_heads, self.head_dim).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+        if mask is not None:
+            attn = attn.masked_fill(mask == 0, float("-inf"))
+        attn = self.dropout(F.softmax(attn, dim=-1))
+        return self.out((attn @ v).transpose(1, 2).contiguous().view(B, T, C))
+
+
+class FFN(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.up = nn.Linear(config.d_model, config.ffn_hidden)
+        self.down = nn.Linear(config.ffn_hidden, config.d_model)
+        self.dropout = nn.Dropout(config.dropout)
+
+    def forward(self, x):
+        return self.dropout(self.down(F.relu(self.up(x))))
+
+
+class Block(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(config.d_model)
+        self.attn = Attention(config)
+        self.norm2 = nn.LayerNorm(config.d_model)
+        self.ffn = FFN(config)
+
+    def forward(self, x, mask=None):
+        x = x + self.attn(self.norm1(x), mask)
+        x = x + self.ffn(self.norm2(x))
+        return x
+
+
+class GuppyLM(nn.Module):
+    def __init__(self, config: GuppyConfig):
+        super().__init__()
+        self.config = config
+
+        self.tok_emb = nn.Embedding(config.vocab_size, config.d_model)
+        self.pos_emb = nn.Embedding(config.max_seq_len, config.d_model)
+        self.drop = nn.Dropout(config.dropout)
+        self.blocks = nn.ModuleList([Block(config) for _ in range(config.n_layers)])
+        self.norm = nn.LayerNorm(config.d_model)
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.lm_head.weight = self.tok_emb.weight  # tie weights
+
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            nn.init.normal_(m.weight, mean=0.0, std=0.02)
+            if m.bias is not None:
+                nn.init.zeros_(m.bias)
+        elif isinstance(m, nn.Embedding):
+            nn.init.normal_(m.weight, mean=0.0, std=0.02)
+
+    def forward(self, idx, targets=None):
+        B, T = idx.shape
+        pos = torch.arange(T, device=idx.device)
+        x = self.drop(self.tok_emb(idx) + self.pos_emb(pos))
+        mask = torch.tril(torch.ones(T, T, device=idx.device)).unsqueeze(0).unsqueeze(0)
+
+        for block in self.blocks:
+            x = block(x, mask)
+
+        logits = self.lm_head(self.norm(x))
+
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(
+                logits.view(-1, self.config.vocab_size),
+                targets.view(-1),
+                ignore_index=0,
+            )
+
+        return logits, loss
+
+    @torch.no_grad()
+    def generate(self, idx, max_new_tokens=64, temperature=0.7, top_k=50, **kwargs):
+        self.eval()
+        for _ in range(max_new_tokens):
+            idx_cond = idx[:, -self.config.max_seq_len:]
+            logits, _ = self(idx_cond)
+            logits = logits[:, -1, :] / temperature
+            if top_k > 0:
+                v, _ = torch.topk(logits, min(top_k, logits.size(-1)))
+                logits[logits < v[:, [-1]]] = float("-inf")
+            probs = F.softmax(logits, dim=-1)
+            next_id = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat([idx, next_id], dim=1)
+            if next_id.item() == self.config.eos_id:
+                break
+        return idx, []
+
+    def param_count(self):
+        total = sum(p.numel() for p in self.parameters())
+        return total, 0
+
+    def param_summary(self):
+        total, _ = self.param_count()
+        return f"GuppyLM: {total:,} params ({total/1e6:.1f}M)"

model_fp32.onnx.data

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

pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6faec973e60786c8942d632e857c536c5404d4970b4c30ca48c096740d18af79
+size 34930539