# tiny-gpt: the model, the training run and every measurement used in # "How Do Language Models Work? I Built One Small Enough to Watch" # https://learnai24.com/ai-basics/how-language-models-work/ # # Requirements: python 3 and torch. Nothing else. Get the corpus with: # curl -O https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt # Save this file as tiny_gpt.py next to input.txt and run: python3 tiny_gpt.py # It trains for about twenty minutes on two CPU threads, then prints every # number that appears in the article. All seeds are fixed, so the first # evaluation should read: train 4.3472 val 4.3459 import math, time, json, os, sys import torch, torch.nn as nn from torch.nn import functional as F torch.manual_seed(1337) torch.set_num_threads(2) DEV='cpu' BLOCK=128; NEMB=128; NHEAD=4; NLAYER=4; BATCH=32; DROP=0.0 text=open('input.txt','r',encoding='utf-8').read() chars=sorted(set(text)); V=len(chars) stoi={c:i for i,c in enumerate(chars)}; itos={i:c for c,i in stoi.items()} enc=lambda s:[stoi[c] for c in s] dec=lambda l:''.join(itos[i] for i in l) data=torch.tensor(enc(text),dtype=torch.long) n=int(0.9*len(data)); train_d=data[:n]; val_d=data[n:] def get_batch(split, g=None): d = train_d if split=='train' else val_d ix = torch.randint(len(d)-BLOCK-1,(BATCH,),generator=g) x = torch.stack([d[i:i+BLOCK] for i in ix]) y = torch.stack([d[i+1:i+BLOCK+1] for i in ix]) return x,y class Head(nn.Module): def __init__(s,hs): super().__init__() s.key=nn.Linear(NEMB,hs,bias=False); s.query=nn.Linear(NEMB,hs,bias=False); s.value=nn.Linear(NEMB,hs,bias=False) s.register_buffer('tril',torch.tril(torch.ones(BLOCK,BLOCK))) def forward(s,x): B,T,C=x.shape k=s.key(x); q=s.query(x) w=q@k.transpose(-2,-1)*k.shape[-1]**-0.5 w=w.masked_fill(s.tril[:T,:T]==0,float('-inf')) w=F.softmax(w,dim=-1) return w@s.value(x) class MHA(nn.Module): def __init__(s,nh,hs): super().__init__(); s.heads=nn.ModuleList([Head(hs) for _ in range(nh)]); s.proj=nn.Linear(NEMB,NEMB) def forward(s,x): return s.proj(torch.cat([h(x) for h in s.heads],dim=-1)) class FF(nn.Module): def __init__(s): super().__init__(); s.net=nn.Sequential(nn.Linear(NEMB,4*NEMB),nn.ReLU(),nn.Linear(4*NEMB,NEMB)) def forward(s,x): return s.net(x) class Block(nn.Module): def __init__(s): super().__init__(); s.sa=MHA(NHEAD,NEMB//NHEAD); s.ff=FF(); s.l1=nn.LayerNorm(NEMB); s.l2=nn.LayerNorm(NEMB) def forward(s,x): x=x+s.sa(s.l1(x)); x=x+s.ff(s.l2(x)); return x class GPT(nn.Module): def __init__(s): super().__init__() s.tok=nn.Embedding(V,NEMB); s.pos=nn.Embedding(BLOCK,NEMB) s.blocks=nn.Sequential(*[Block() for _ in range(NLAYER)]) s.lnf=nn.LayerNorm(NEMB); s.head=nn.Linear(NEMB,V) def forward(s,idx,targets=None): B,T=idx.shape x=s.tok(idx)+s.pos(torch.arange(T)) x=s.lnf(s.blocks(x)); logits=s.head(x) loss=None if targets is not None: loss=F.cross_entropy(logits.view(B*T,V),targets.view(B*T)) return logits,loss @torch.no_grad() def generate(s,idx,maxnew,temp=1.0,g=None): for _ in range(maxnew): logits,_=s(idx[:,-BLOCK:]) logits=logits[:,-1,:]/temp probs=F.softmax(logits,dim=-1) nxt=torch.multinomial(probs,1,generator=g) idx=torch.cat((idx,nxt),dim=1) return idx if __name__=='__main__': MAXSTEPS=5000 m=GPT() print('corpus',len(text),'training split',n,'held-out',len(data)-n,'vocabulary',V,flush=True) print('params',sum(p.numel() for p in m.parameters()),flush=True) opt=torch.optim.AdamW(m.parameters(),lr=1e-3) @torch.no_grad() def evaluate(): m.eval(); out={} g=torch.Generator().manual_seed(99) for split in ['train','val']: ls=torch.zeros(20) for k in range(20): x,y=get_batch(split,g); _,l=m(x,y); ls[k]=l.item() out[split]=ls.mean().item() m.train(); return out @torch.no_grad() def sample(n=260,temp=1.0): m.eval() g=torch.Generator().manual_seed(7) ctx=torch.tensor([[stoi['\n']]],dtype=torch.long) s=dec(m.generate(ctx,n,temp=temp,g=g)[0].tolist()) m.train(); return s log=[] SNAPS=[0,100,500,2000,5000] t0=time.time() for step in range(MAXSTEPS+1): if step in SNAPS: ev=evaluate(); sm=sample() log.append({'step':step,'train':round(ev['train'],4),'val':round(ev['val'],4),'sample':sm}) print('=== step',step,'train %.4f val %.4f (%.0fs)'%(ev['train'],ev['val'],time.time()-t0),flush=True) print(repr(sm),flush=True) json.dump(log,open('log.json','w'),indent=1) torch.save(m.state_dict(),'ck_%d.pt'%step) if step==MAXSTEPS: break x,y=get_batch('train'); _,loss=m(x,y) opt.zero_grad(set_to_none=True); loss.backward(); opt.step() print('done %.0fs'%(time.time()-t0),flush=True) # ---------------- measurements used in the article ---------------- import math, collections, re m.eval() train_text = text[:n] def dist(prompt, k=8): idx = torch.tensor([enc(prompt)], dtype=torch.long)[:, -BLOCK:] with torch.no_grad(): logits, _ = m(idx) p = F.softmax(logits[0, -1, :], dim=-1) v, i = torch.topk(p, k) return [(itos[int(a)], round(float(b) * 100, 2)) for a, b in zip(i, v)], p print('\n--- full distribution after "hear me spea"') top, p = dist('First Citizen:\nBefore we proceed any further, hear me spea') print(top[:5], 'rest %.4f%%' % (100 * (1 - sum(x[1] for x in top[:5]) / 100))) print('--- full distribution after "\\nMENENIUS:\\nThe "') print(dist('\nMENENIUS:\nThe ')[0][:6]) print('\n--- top guess at every position of one speech') line = 'First Citizen:\nBefore we proceed any further, hear me speak.' tops = [] for j in range(3, len(line) - 1): top, p = dist(line[:j], 1) tops.append(top[0][1]) print(repr(line[max(0, j - 6):j]), 'guess', repr(top[0][0]), top[0][1], 'actual', repr(line[j])) print('positions', len(tops), 'min', min(tops), 'max', max(tops)) print('\n--- reference points') # A character-frequency model, fitted on the training split only and scored # on exactly the batches evaluate() uses, so both numbers are comparable to # the train/val losses printed above. cnt = collections.Counter(train_text); NN = len(train_text) lp = {c: math.log(cnt[c] / NN) if cnt.get(c) else -30.0 for c in chars} gb = torch.Generator().manual_seed(99); ref = {} for split in ['train', 'val']: d = train_d if split == 'train' else val_d tot = 0.0; cn = 0 for _ in range(20): for i in torch.randint(len(d) - BLOCK - 1, (BATCH,), generator=gb).tolist(): for j in range(i + 1, i + BLOCK + 1): tot += -lp[itos[d[j].item()]]; cn += 1 ref[split] = tot / cn print('uniform %.4f' % math.log(V)) print('unigram train %.4f val %.4f' % (ref['train'], ref['val'])) def gen(temp, seed, count): gg = torch.Generator().manual_seed(seed) ctx = torch.tensor([[stoi['\n']]], dtype=torch.long) return dec(m.generate(ctx, count, temp=temp, g=gg)[0].tolist()) print('\n--- repetition and vocabulary against temperature (5 seeds each)') for temp in [0.2, 0.5, 1.0, 1.5]: rr = []; dd = [] for sd in [1, 2, 3, 4, 5]: s = gen(temp, sd, 1000) w = s.split(); dd.append(len(set(w)) / max(1, len(w))) seen = set(); dup = 0; tot = 0 for i in range(len(s) - 11): win = s[i:i + 12]; tot += 1 if win in seen: dup += 1 seen.add(win) rr.append(dup / tot) print('temp', temp, 'repeat12 %.2f%%' % (100 * sum(rr) / len(rr)), 'distinct words %.1f%%' % (100 * sum(dd) / len(dd))) print('\n--- share of generated word-like strings that appear in the corpus') words = set(w.lower() for w in re.findall(r"[a-zA-Z']+", text)) print('corpus word types', len(words)) for st in SNAPS[1:]: mm = GPT(); mm.load_state_dict(torch.load('ck_%d.pt' % st)); mm.eval() gg = torch.Generator().manual_seed(7) ctx = torch.tensor([[stoi['\n']]], dtype=torch.long) s = dec(mm.generate(ctx, 2000, temp=1.0, g=gg)[0].tolist()) ws = [w.lower() for w in re.findall(r"[a-zA-Z']+", s)] print('step', st, '%.1f%%' % (100 * sum(1 for w in ws if w in words) / len(ws))) print('\n--- longest passage shared with the training split, 600 characters, 5 seeds') for sd in [7, 11, 23, 42, 101]: s = gen(1.0, sd, 600); best = 0; bs = '' for L in range(4, 150): hit = None for i in range(0, len(s) - L + 1): if s[i:i + L] in train_text: hit = s[i:i + L]; break if hit is None: break best = L; bs = hit print('seed', sd, best, repr(bs))