Edge: train a language model small enough to run on a Raspberry Pi

import os
import torch
import torch.nn as nn
import torch.nn.functional as F
 
torch.manual_seed(1337)
device = "cuda" if torch.cuda.is_available() else "cpu"
print("device:", device)

text = open("data/tinyshakespeare.txt").read()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: "".join(itos[i] for i in ids)
 
data = torch.tensor(encode(text), dtype=torch.long)
split = int(0.9 * len(data))
train, val = data[:split], data[split:]
print(f"{len(text):,} characters, vocab of {len(chars)}")

class Block(nn.Module):
    def __init__(self, d, heads):
        super().__init__()
        self.ln1 = nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, heads, batch_first=True)
        self.ln2 = nn.LayerNorm(d)
        self.mlp = nn.Sequential(nn.Linear(d, 4 * d), nn.GELU(), nn.Linear(4 * d, d))
 
    def forward(self, x):
        T = x.size(1)
        mask = torch.triu(torch.ones(T, T, dtype=torch.bool, device=x.device), 1)
        h = self.ln1(x)
        x = x + self.attn(h, h, h, attn_mask=mask, need_weights=False)[0]
        return x + self.mlp(self.ln2(x))
 
class TinyGPT(nn.Module):
    def __init__(self, vocab, d=128, heads=4, layers=4, ctx=128):
        super().__init__()
        self.ctx = ctx
        self.tok = nn.Embedding(vocab, d)
        self.pos = nn.Embedding(ctx, d)
        self.blocks = nn.Sequential(*[Block(d, heads) for _ in range(layers)])
        self.ln = nn.LayerNorm(d)
        self.head = nn.Linear(d, vocab)
 
    def forward(self, idx):
        pos = torch.arange(idx.size(1), device=idx.device)
        x = self.blocks(self.tok(idx) + self.pos(pos))
        return self.head(self.ln(x))
 
model = TinyGPT(len(chars)).to(device)
print(f"{sum(p.numel() for p in model.parameters()) / 1e6:.2f}M parameters")

def batch(split, bs=32):
    d = train if split == "train" else val
    i = torch.randint(len(d) - model.ctx - 1, (bs,))
    x = torch.stack([d[j:j + model.ctx] for j in i])
    y = torch.stack([d[j + 1:j + model.ctx + 1] for j in i])
    return x.to(device), y.to(device)
 
opt = torch.optim.AdamW(model.parameters(), lr=3e-4)
for step in range(5001):
    x, y = batch("train")
    loss = F.cross_entropy(model(x).flatten(0, 1), y.flatten())
    opt.zero_grad()
    loss.backward()
    opt.step()
    if step % 500 == 0:
        print(f"step {step:>4}  loss {loss.item():.3f}")

@torch.no_grad()
def generate(prompt, n=200, temp=0.8):
    idx = torch.tensor([encode(prompt)], device=device)
    for _ in range(n):
        logits = model(idx[:, -model.ctx:])[:, -1] / temp
        nxt = torch.multinomial(F.softmax(logits, -1), 1)
        idx = torch.cat([idx, nxt], 1)
    return decode(idx[0].tolist())
 
print(generate("ROMEO:"))

model.eval().cpu()
dummy = torch.zeros(1, model.ctx, dtype=torch.long)
torch.onnx.export(model, dummy, "edge.onnx",
                  input_names=["tokens"], output_names=["logits"])
print(f"saved edge.onnx ({os.path.getsize('edge.onnx') / 1e6:.1f} MB)")