File size: 3,935 Bytes
a40cbeb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 | """Self-contained PyTorch implementation of tiny-edu-114m (plain PyTorch, no transformers modeling code needed).
import torch
from safetensors.torch import load_file
from modeling_tiny_edu import TinyEdu, generate
model = TinyEdu(); model.load_state_dict(load_file("model.safetensors"), strict=False); model.eval() # strict=False: the output head is tied to the embedding
from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained("gpt2")
print(generate(model, tok, "The capital of France is", max_new=30))
"""
import torch, torch.nn as nn, torch.nn.functional as F
D, NH, NKV, NL, V, FF = 768, 12, 4, 12, 50257, 3072 # width, query heads, key/value heads, layers, vocabulary, MLP width
HD = D // NH # head dimension 64
class Block(nn.Module):
def __init__(self):
super().__init__()
self.ln_1, self.ln_2 = nn.LayerNorm(D, eps=1e-5), nn.LayerNorm(D, eps=1e-5)
self.attn = nn.Module(); self.attn.q_proj = nn.Linear(D, NH * HD, bias=False); self.attn.k_proj = nn.Linear(D, NKV * HD, bias=False)
self.attn.v_proj = nn.Linear(D, NKV * HD, bias=False); self.attn.out_proj = nn.Linear(NH * HD, D, bias=False)
self.mlp = nn.Module(); self.mlp.c_fc = nn.Linear(D, FF, bias=False); self.mlp.c_proj = nn.Linear(FF, D, bias=False)
def forward(self, x, cos, sin):
B, T, _ = x.shape
h = self.ln_1(x)
q = self.attn.q_proj(h).view(B, T, NH, HD); k = self.attn.k_proj(h).view(B, T, NKV, HD); v = self.attn.v_proj(h).view(B, T, NKV, HD)
def rot(t): # rotary embedding, half-split rotation
t1, t2 = t[..., :HD // 2], t[..., HD // 2:]
c, s = cos[None, :T, None, :].to(t.dtype), sin[None, :T, None, :].to(t.dtype)
return torch.cat([t1 * c - t2 * s, t2 * c + t1 * s], -1)
q, k = rot(q), rot(k)
k = k.repeat_interleave(NH // NKV, 2); v = v.repeat_interleave(NH // NKV, 2)
o = F.scaled_dot_product_attention(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=True).transpose(1, 2).reshape(B, T, NH * HD)
x = x + self.attn.out_proj(o)
return x + self.mlp.c_proj(F.gelu(self.mlp.c_fc(self.ln_2(x))))
class TinyEdu(nn.Module):
def __init__(self):
super().__init__()
self.transformer = nn.Module(); self.transformer.wte = nn.Embedding(V, D); self.transformer.h = nn.ModuleList([Block() for _ in range(NL)]); self.transformer.ln_f = nn.LayerNorm(D, eps=1e-5)
inv = 1.0 / (10000 ** (torch.arange(0, HD, 2).float() / HD)); f = torch.outer(torch.arange(4096).float(), inv)
self.register_buffer("cos", f.cos(), persistent=False); self.register_buffer("sin", f.sin(), persistent=False)
def forward(self, ids): # (B, T) token ids -> (B, T, 50257) logits; the output head shares the embedding matrix
x = self.transformer.wte(ids)
for b in self.transformer.h:
x = b(x, self.cos, self.sin)
return self.transformer.ln_f(x) @ self.transformer.wte.weight.T
@torch.no_grad()
def generate(model, tok, prompt, max_new=60, temperature=0.8, top_k=50, device=None):
"""Plain sampling without a cache (fine for a model this small); temperature 0 = greedy."""
device = device or next(model.parameters()).device
ids = torch.tensor([tok.encode(prompt)], device=device)
for _ in range(max_new):
lg = model(ids[:, -1024:])[:, -1].float()
if temperature == 0:
nxt = lg.argmax(-1, keepdim=True)
else:
lg = lg / temperature
if top_k:
lg[lg < lg.topk(top_k).values[:, -1:]] = -float("inf")
nxt = torch.multinomial(lg.softmax(-1), 1)
ids = torch.cat([ids, nxt], 1)
if int(nxt) == 50256:
break
return tok.decode(ids[0].tolist())
|