| """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 |
| HD = D // NH |
|
|
|
|
| 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): |
| 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): |
| 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()) |
|
|
