train slm: 17.78M params

README.md

@@ -1,3 +1,16 @@
----
-license: mit
----
+# SLM (small human model)
+
+Tiny transformer LM trained from scratch on a toy corpus in Colab.
+
+- params: 17.78M
+- arch: {'n_embd': 384, 'n_layer': 10, 'n_head': 8, 'block_size': 64}
+- vocab: `
+ .abcdefghijklmnopqrstuvwxyz`
+
+## Usage
+
+```python
+from model import load_model, sample
+model, vocab = load_model("mohar07/PunyPunk")
+print(sample(model, vocab, "the quick", new_tokens=200))
+```

config.json

@@ -0,0 +1 @@
+{"arch": {"n_embd": 384, "n_layer": 10, "n_head": 8, "block_size": 64}, "vocab": "\n .abcdefghijklmnopqrstuvwxyz"}

model.py

@@ -0,0 +1,75 @@
+import os
+import json
+from huggingface_hub import snapshot_download
+
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class DenseTransformerBlock(nn.Module):
+    def __init__(self, n_embd, n_head, block_size, dropout):
+        super().__init__()
+        self.ln1 = nn.LayerNorm(n_embd)
+        self.attn = nn.MultiheadAttention(n_embd, n_head, dropout=dropout, batch_first=True)
+        self.ln2 = nn.LayerNorm(n_embd)
+        self.mlp = nn.Sequential(
+            nn.Linear(n_embd, 4 * n_embd),
+            nn.GELU(),
+            nn.Linear(4 * n_embd, n_embd),
+            nn.Dropout(dropout),
+        )
+        self.register_buffer("mask", torch.triu(torch.full((block_size, block_size), float("-inf")), diagonal=1))
+
+    def forward(self, x):
+        a, _ = self.attn(self.ln1(x), self.ln1(x), self.ln1(x), attn_mask=self.mask)
+        x = x + a
+        x = x + self.mlp(self.ln2(x))
+        return x
+
+class SLM(nn.Module):
+    def __init__(self, vocab_size, n_embd=384, n_head=8, n_layer=10, block_size=64, dropout=0.1):
+        super().__init__()
+        self.token_emb = nn.Embedding(vocab_size, n_embd)
+        self.pos_emb = nn.Embedding(block_size, n_embd)
+        self.blocks = nn.Sequential(*[DenseTransformerBlock(n_embd, n_head, block_size, dropout) for _ in range(n_layer)])
+        self.ln_f = nn.LayerNorm(n_embd)
+        self.lm_head = nn.Linear(n_embd, vocab_size, bias=False)
+        self.token_emb.weight = self.lm_head.weight
+        self.block_size = block_size
+
+    def forward(self, idx, targets=None):
+        b, t = idx.shape
+        x = self.token_emb(idx) + self.pos_emb(torch.arange(t, device=idx.device))
+        x = self.blocks(x)
+        logits = self.lm_head(self.ln_f(x))
+        loss = None
+        if targets is not None:
+            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1))
+        return logits, loss
+
+    @torch.no_grad()
+    def generate(self, idx, new_tokens, temperature=1.0):
+        for _ in range(new_tokens):
+            idx_cond = idx[:, -self.block_size :]
+            logits, _ = self.forward(idx_cond)
+            logits = logits[:, -1, :] / temperature
+            probs = F.softmax(logits, dim=-1)
+            idx = torch.cat([idx, torch.multinomial(probs, 1)], dim=1)
+        return idx
+
+def load_model(repo_or_dir):
+    ckpt_dir = repo_or_dir if os.path.isdir(repo_or_dir) else snapshot_download(repo_or_dir)
+    cfg = json.load(open(os.path.join(ckpt_dir, "config.json")))
+    vocab = cfg["vocab"]
+    arch = cfg["arch"]
+    model = SLM(len(vocab), **arch)
+    model.load_state_dict(torch.load(os.path.join(ckpt_dir, "pytorch_model.bin"), map_location="cpu"))
+    return model, vocab
+
+def sample(model, vocab, prompt, new_tokens=200, temperature=1.0):
+    stoi = {c: i for i, c in enumerate(vocab)}
+    itos = {i: c for c, i in stoi.items()}
+    idx = torch.tensor([[stoi[c] for c in prompt]], dtype=torch.long)
+    out = model.generate(idx, new_tokens, temperature=temperature)[0].tolist()
+    return "".join(itos[i] for i in out)

pytorch_model.bin

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