Add model code

model.py

@@ -0,0 +1,137 @@
+"""TinyChat-5M: 5.1M param transformer. d=192, 6L, 3H, head_dim=64, SwiGLU 4x, RoPE, RMSNorm, tied embeddings."""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, d, eps=1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(d))
+        self.eps = eps
+
+    def forward(self, x):
+        rms = x.pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
+        return x * rms * self.weight
+
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([-x2, x1], dim=-1)
+
+
+class Attention(nn.Module):
+    def __init__(self, d, n_heads, head_dim):
+        super().__init__()
+        self.n_heads = n_heads
+        self.head_dim = head_dim
+        self.wq = nn.Linear(d, n_heads * head_dim, bias=False)
+        self.wk = nn.Linear(d, n_heads * head_dim, bias=False)
+        self.wv = nn.Linear(d, n_heads * head_dim, bias=False)
+        self.wo = nn.Linear(n_heads * head_dim, d, bias=False)
+
+    def forward(self, x, cos, sin):
+        B, S, _ = x.shape
+        q = self.wq(x).view(B, S, self.n_heads, self.head_dim).transpose(1, 2)
+        k = self.wk(x).view(B, S, self.n_heads, self.head_dim).transpose(1, 2)
+        v = self.wv(x).view(B, S, self.n_heads, self.head_dim).transpose(1, 2)
+        cos_f = cos[:S].unsqueeze(0).unsqueeze(0)
+        sin_f = sin[:S].unsqueeze(0).unsqueeze(0)
+        q = (q * cos_f) + (rotate_half(q) * sin_f)
+        k = (k * cos_f) + (rotate_half(k) * sin_f)
+        scale = self.head_dim ** -0.5
+        scores = (q @ k.transpose(-2, -1)) * scale
+        mask = torch.triu(torch.ones(S, S, device=x.device, dtype=torch.bool), diagonal=1)
+        scores = scores.masked_fill(mask, float('-inf'))
+        attn = F.softmax(scores, dim=-1)
+        out = (attn @ v).transpose(1, 2).reshape(B, S, -1)
+        return self.wo(out)
+
+
+class FFN(nn.Module):
+    def __init__(self, d, mult):
+        super().__init__()
+        hidden = d * mult
+        self.w1 = nn.Linear(d, hidden, bias=False)
+        self.w2 = nn.Linear(hidden, d, bias=False)
+        self.w3 = nn.Linear(d, hidden, bias=False)
+
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class Block(nn.Module):
+    def __init__(self, d, n_heads, head_dim, ffn_mult):
+        super().__init__()
+        self.norm1 = RMSNorm(d)
+        self.attn = Attention(d, n_heads, head_dim)
+        self.norm2 = RMSNorm(d)
+        self.ffn = FFN(d, ffn_mult)
+
+    def forward(self, x, cos, sin):
+        x = x + self.attn(self.norm1(x), cos, sin)
+        x = x + self.ffn(self.norm2(x))
+        return x
+
+
+class TinyLM(nn.Module):
+    def __init__(self, vocab=4096, d=192, n_layers=6, n_heads=3, head_dim=64, ffn_mult=4):
+        super().__init__()
+        self.tok_emb = nn.Embedding(vocab, d)
+        self.layers = nn.ModuleList([Block(d, n_heads, head_dim, ffn_mult) for _ in range(n_layers)])
+        self.norm_f = RMSNorm(d)
+        self.head = nn.Linear(d, vocab, bias=False)
+        # Tied embeddings
+        self.head.weight = self.tok_emb.weight
+        # RoPE
+        freqs = 1.0 / (10000 ** (torch.arange(0, head_dim, 2).float() / head_dim))
+        t = torch.arange(2048).float()
+        outer = torch.outer(t, freqs)
+        angles = torch.cat([outer, outer], dim=-1)
+        self.register_buffer("cos", angles.cos())
+        self.register_buffer("sin", angles.sin())
+
+    def forward(self, x, y=None):
+        h = self.tok_emb(x)
+        for layer in self.layers:
+            h = layer(h, self.cos, self.sin)
+        h = self.norm_f(h)
+        logits = self.head(h)
+        if y is not None:
+            return F.cross_entropy(logits.reshape(-1, logits.size(-1)), y.reshape(-1))
+        return logits
+
+
+def load_model(path="model.pt", device="cpu"):
+    model = TinyLM()
+    ckpt = torch.load(path, map_location=device, weights_only=True)
+    model.load_state_dict(ckpt["model"])
+    model.eval()
+    return model
+
+
+def generate(model, tokenizer, prompt, max_new=100, temperature=0.8, rep_penalty=1.5):
+    ids = tokenizer.encode(prompt, add_special_tokens=False).ids
+    ids = torch.tensor([ids])
+    with torch.no_grad():
+        for _ in range(max_new):
+            inp = ids[:, -512:]
+            logits = model(inp)[0, -1]
+            for tid in ids[0].unique():
+                if logits[tid] > 0:
+                    logits[tid] /= rep_penalty
+                else:
+                    logits[tid] *= rep_penalty
+            logits = logits / temperature
+            probs = F.softmax(logits, dim=-1)
+            next_id = torch.multinomial(probs, 1).item()
+            ids = torch.cat([ids, torch.tensor([[next_id]])], dim=1)
+    return tokenizer.decode(ids[0].tolist(), skip_special_tokens=True)
+
+
+if __name__ == "__main__":
+    from tokenizers import Tokenizer
+    tok = Tokenizer.from_file("tokenizer.json")
+    model = load_model()
+    print(f"Params: {sum(p.numel() for p in model.parameters()):,}")
+    print(generate(model, tok, "[INST] Hello! [/INST]", max_new=100))