Language Modelling

import torch

def build_corpus(seed=0, sentences=500):
"""A small language with real grammar, generated from a seed.

Written into the page rather than downloaded, so the week runs
offline and everybody trains on identical characters. It is
regular enough that a model can genuinely learn it, and long
enough that counting pairs of characters cannot."""
g = torch.Generator().manual_seed(seed)
def pick(xs):
return xs[int(torch.randint(len(xs), (1,), generator=g))]
subject = ['the model', 'the network', 'the optimiser',
'the gradient', 'the loss', 'attention',
'the encoder', 'the decoder', 'each layer',
'the batch', 'the learning rate', 'dropout']
verb = ['reads', 'predicts', 'updates', 'normalises',
'averages', 'scores', 'reduces', 'controls',
'stabilises', 'copies']
obj = ['the sequence', 'the next token', 'every position',
'the weights', 'the hidden state', 'the input',
'the residual path', 'the attention weights',
'the training data', 'the validation set']
tail = ['during training', 'at every step', 'in a single pass',
'before the softmax', 'after the residual connection',
'without a mask', 'across the batch', 'once per epoch']
joiner = [', and ', ', so ', ', because ', '. ', '. ', '. ']
out = []
for _ in range(sentences):
out.append('%s %s %s %s%s'
% (pick(subject), pick(verb), pick(obj),
pick(tail), pick(joiner)))
return ''.join(out)

CORPUS = build_corpus()
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

text = CORPUS.strip()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
VOCAB = len(chars)

def encode(s):
return torch.tensor([stoi[c] for c in s], dtype=torch.long)

def decode(t):
return ''.join(itos[int(i)] for i in t)

data = encode(text)
BLOCK_SIZE = 32

def windows(seq, block):
"""Every position is a training example: x is block characters,
y is the same block shifted one to the right."""
starts = range(0, len(seq) - block - 1, 4) # stride 4
x = torch.stack([seq[i:i + block] for i in starts])
y = torch.stack([seq[i + 1:i + block + 1] for i in starts])
return x, y

split = int(len(data) * 0.9)
X_tr, Y_tr = windows(data[:split], BLOCK_SIZE)
X_va, Y_va = windows(data[split:], BLOCK_SIZE)
train_loader = DataLoader(TensorDataset(X_tr, Y_tr), batch_size=64,
shuffle=True)
print('characters %d' % len(text))
print('vocabulary %d: %s' % (VOCAB, ''.join(chars)))
print()
print(text[:180])
print()
print('encoded:', encode('the model')[:12].tolist())
print('decoded:', decode(encode('the model')))

import torch

def build_corpus(seed=0, sentences=500):
"""A small language with real grammar, generated from a seed.

Written into the page rather than downloaded, so the week runs
offline and everybody trains on identical characters. It is
regular enough that a model can genuinely learn it, and long
enough that counting pairs of characters cannot."""
g = torch.Generator().manual_seed(seed)
def pick(xs):
return xs[int(torch.randint(len(xs), (1,), generator=g))]
subject = ['the model', 'the network', 'the optimiser',
'the gradient', 'the loss', 'attention',
'the encoder', 'the decoder', 'each layer',
'the batch', 'the learning rate', 'dropout']
verb = ['reads', 'predicts', 'updates', 'normalises',
'averages', 'scores', 'reduces', 'controls',
'stabilises', 'copies']
obj = ['the sequence', 'the next token', 'every position',
'the weights', 'the hidden state', 'the input',
'the residual path', 'the attention weights',
'the training data', 'the validation set']
tail = ['during training', 'at every step', 'in a single pass',
'before the softmax', 'after the residual connection',
'without a mask', 'across the batch', 'once per epoch']
joiner = [', and ', ', so ', ', because ', '. ', '. ', '. ']
out = []
for _ in range(sentences):
out.append('%s %s %s %s%s'
% (pick(subject), pick(verb), pick(obj),
pick(tail), pick(joiner)))
return ''.join(out)

CORPUS = build_corpus()
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

text = CORPUS.strip()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
VOCAB = len(chars)

def encode(s):
return torch.tensor([stoi[c] for c in s], dtype=torch.long)

def decode(t):
return ''.join(itos[int(i)] for i in t)

data = encode(text)
BLOCK_SIZE = 32

def windows(seq, block):
"""Every position is a training example: x is block characters,
y is the same block shifted one to the right."""
starts = range(0, len(seq) - block - 1, 4) # stride 4
x = torch.stack([seq[i:i + block] for i in starts])
y = torch.stack([seq[i + 1:i + block + 1] for i in starts])
return x, y

split = int(len(data) * 0.9)
X_tr, Y_tr = windows(data[:split], BLOCK_SIZE)
X_va, Y_va = windows(data[split:], BLOCK_SIZE)
train_loader = DataLoader(TensorDataset(X_tr, Y_tr), batch_size=64,
shuffle=True)
x, y = X_tr[0], Y_tr[0]
print('input :', repr(decode(x[:32])))
print('target:', repr(decode(y[:32])))
print('\nthe target is the input shifted one place left.')
print()
for i in range(5):
print('given %-22r predict %r'
% (decode(x[:i + 1]), decode(y[i:i + 1])))
print('\n%d windows of %d characters, so %d predictions per epoch'
% (len(X_tr), BLOCK_SIZE, len(X_tr) * BLOCK_SIZE))

import torch

def build_corpus(seed=0, sentences=500):
"""A small language with real grammar, generated from a seed.

Written into the page rather than downloaded, so the week runs
offline and everybody trains on identical characters. It is
regular enough that a model can genuinely learn it, and long
enough that counting pairs of characters cannot."""
g = torch.Generator().manual_seed(seed)
def pick(xs):
return xs[int(torch.randint(len(xs), (1,), generator=g))]
subject = ['the model', 'the network', 'the optimiser',
'the gradient', 'the loss', 'attention',
'the encoder', 'the decoder', 'each layer',
'the batch', 'the learning rate', 'dropout']
verb = ['reads', 'predicts', 'updates', 'normalises',
'averages', 'scores', 'reduces', 'controls',
'stabilises', 'copies']
obj = ['the sequence', 'the next token', 'every position',
'the weights', 'the hidden state', 'the input',
'the residual path', 'the attention weights',
'the training data', 'the validation set']
tail = ['during training', 'at every step', 'in a single pass',
'before the softmax', 'after the residual connection',
'without a mask', 'across the batch', 'once per epoch']
joiner = [', and ', ', so ', ', because ', '. ', '. ', '. ']
out = []
for _ in range(sentences):
out.append('%s %s %s %s%s'
% (pick(subject), pick(verb), pick(obj),
pick(tail), pick(joiner)))
return ''.join(out)

CORPUS = build_corpus()
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

text = CORPUS.strip()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
VOCAB = len(chars)

def encode(s):
return torch.tensor([stoi[c] for c in s], dtype=torch.long)

def decode(t):
return ''.join(itos[int(i)] for i in t)

data = encode(text)
BLOCK_SIZE = 32

def windows(seq, block):
"""Every position is a training example: x is block characters,
y is the same block shifted one to the right."""
starts = range(0, len(seq) - block - 1, 4) # stride 4
x = torch.stack([seq[i:i + block] for i in starts])
y = torch.stack([seq[i + 1:i + block + 1] for i in starts])
return x, y

split = int(len(data) * 0.9)
X_tr, Y_tr = windows(data[:split], BLOCK_SIZE)
X_va, Y_va = windows(data[split:], BLOCK_SIZE)
train_loader = DataLoader(TensorDataset(X_tr, Y_tr), batch_size=64,
shuffle=True)
import math
from collections import Counter

counts = Counter(text)
total = sum(counts.values())
uniform = math.log(VOCAB)
unigram = -sum((n / total) * math.log(n / total) for n in counts.values())

print('loss if the model guesses uniformly %.4f' % uniform)
print('loss from character frequencies alone %.4f' % unigram)
print('\nperplexity is exp(loss), the effective number of choices:')
print(' uniform %.1f' % math.exp(uniform))
print(' unigram %.1f' % math.exp(unigram))
print('\nanything above the unigram number means the model has not')
print('learned even which letters are common.')

import torch

def build_corpus(seed=0, sentences=500):
"""A small language with real grammar, generated from a seed.

Written into the page rather than downloaded, so the week runs
offline and everybody trains on identical characters. It is
regular enough that a model can genuinely learn it, and long
enough that counting pairs of characters cannot."""
g = torch.Generator().manual_seed(seed)
def pick(xs):
return xs[int(torch.randint(len(xs), (1,), generator=g))]
subject = ['the model', 'the network', 'the optimiser',
'the gradient', 'the loss', 'attention',
'the encoder', 'the decoder', 'each layer',
'the batch', 'the learning rate', 'dropout']
verb = ['reads', 'predicts', 'updates', 'normalises',
'averages', 'scores', 'reduces', 'controls',
'stabilises', 'copies']
obj = ['the sequence', 'the next token', 'every position',
'the weights', 'the hidden state', 'the input',
'the residual path', 'the attention weights',
'the training data', 'the validation set']
tail = ['during training', 'at every step', 'in a single pass',
'before the softmax', 'after the residual connection',
'without a mask', 'across the batch', 'once per epoch']
joiner = [', and ', ', so ', ', because ', '. ', '. ', '. ']
out = []
for _ in range(sentences):
out.append('%s %s %s %s%s'
% (pick(subject), pick(verb), pick(obj),
pick(tail), pick(joiner)))
return ''.join(out)

CORPUS = build_corpus()
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

text = CORPUS.strip()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
VOCAB = len(chars)

def encode(s):
return torch.tensor([stoi[c] for c in s], dtype=torch.long)

def decode(t):
return ''.join(itos[int(i)] for i in t)

data = encode(text)
BLOCK_SIZE = 32

def windows(seq, block):
"""Every position is a training example: x is block characters,
y is the same block shifted one to the right."""
starts = range(0, len(seq) - block - 1, 4) # stride 4
x = torch.stack([seq[i:i + block] for i in starts])
y = torch.stack([seq[i + 1:i + block + 1] for i in starts])
return x, y

split = int(len(data) * 0.9)
X_tr, Y_tr = windows(data[:split], BLOCK_SIZE)
X_va, Y_va = windows(data[split:], BLOCK_SIZE)
train_loader = DataLoader(TensorDataset(X_tr, Y_tr), batch_size=64,
shuffle=True)
import math

# The simplest possible language model: for each character, count what
# followed it, and predict from those counts.
counts = torch.zeros(VOCAB, VOCAB)
for a, b in zip(data[:split][:-1], data[:split][1:]):
counts[a, b] += 1
probs = (counts + 1) / (counts + 1).sum(1, keepdim=True) # smoothed

val = data[split:]
logp = probs[val[:-1], val[1:]].log().mean()
print('bigram validation loss %.4f, perplexity %.1f'
% (-logp.item(), math.exp(-logp.item())))

torch.manual_seed(0)
out, current = [], stoi['t']
for _ in range(120):
current = int(torch.multinomial(probs[current], 1))
out.append(current)
print('\ngenerated from bigram counts:')
print(repr(decode(out)))

import torch

def build_corpus(seed=0, sentences=500):
"""A small language with real grammar, generated from a seed.

Written into the page rather than downloaded, so the week runs
offline and everybody trains on identical characters. It is
regular enough that a model can genuinely learn it, and long
enough that counting pairs of characters cannot."""
g = torch.Generator().manual_seed(seed)
def pick(xs):
return xs[int(torch.randint(len(xs), (1,), generator=g))]
subject = ['the model', 'the network', 'the optimiser',
'the gradient', 'the loss', 'attention',
'the encoder', 'the decoder', 'each layer',
'the batch', 'the learning rate', 'dropout']
verb = ['reads', 'predicts', 'updates', 'normalises',
'averages', 'scores', 'reduces', 'controls',
'stabilises', 'copies']
obj = ['the sequence', 'the next token', 'every position',
'the weights', 'the hidden state', 'the input',
'the residual path', 'the attention weights',
'the training data', 'the validation set']
tail = ['during training', 'at every step', 'in a single pass',
'before the softmax', 'after the residual connection',
'without a mask', 'across the batch', 'once per epoch']
joiner = [', and ', ', so ', ', because ', '. ', '. ', '. ']
out = []
for _ in range(sentences):
out.append('%s %s %s %s%s'
% (pick(subject), pick(verb), pick(obj),
pick(tail), pick(joiner)))
return ''.join(out)

CORPUS = build_corpus()
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

text = CORPUS.strip()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
VOCAB = len(chars)

def encode(s):
return torch.tensor([stoi[c] for c in s], dtype=torch.long)

def decode(t):
return ''.join(itos[int(i)] for i in t)

data = encode(text)
BLOCK_SIZE = 32

def windows(seq, block):
"""Every position is a training example: x is block characters,
y is the same block shifted one to the right."""
starts = range(0, len(seq) - block - 1, 4) # stride 4
x = torch.stack([seq[i:i + block] for i in starts])
y = torch.stack([seq[i + 1:i + block + 1] for i in starts])
return x, y

split = int(len(data) * 0.9)
X_tr, Y_tr = windows(data[:split], BLOCK_SIZE)
X_va, Y_va = windows(data[split:], BLOCK_SIZE)
train_loader = DataLoader(TensorDataset(X_tr, Y_tr), batch_size=64,
shuffle=True)
class GPT(nn.Module):
def __init__(self, vocab=None, width=96, heads=4, blocks=2,
block_size=BLOCK_SIZE, dropout=0.1):
super().__init__()
vocab = vocab or VOCAB
self.block_size = block_size
self.tok = nn.Embedding(vocab, width)
self.pos = nn.Embedding(block_size, width)
self.drop = nn.Dropout(dropout)
layer = nn.TransformerEncoderLayer(
d_model=width, nhead=heads, dim_feedforward=width * 4,
dropout=dropout, batch_first=True, norm_first=True,
activation='gelu')
self.blocks = nn.TransformerEncoder(layer, num_layers=blocks)
self.norm = nn.LayerNorm(width)
self.head = nn.Linear(width, vocab, bias=False)
self.head.weight = self.tok.weight

def forward(self, idx):
T = idx.shape[1]
mask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
h = self.drop(self.tok(idx) + self.pos(torch.arange(T)))
return self.head(self.norm(self.blocks(h, mask=mask)))

@torch.no_grad()
def generate(self, idx, steps, temperature=1.0, top_k=None):
self.eval()
for _ in range(steps):
window = idx[:, -self.block_size:]
logits = self(window)[:, -1, :] / temperature
if top_k:
kth = logits.topk(min(top_k, logits.shape[-1])).values[:, -1:]
logits = logits.masked_fill(logits < kth, float('-inf'))
nxt = torch.multinomial(logits.softmax(-1), 1)
idx = torch.cat([idx, nxt], dim=1)
return idx

def train(model, epochs=5, lr=3e-3, warmup=100):
torch.manual_seed(0)
opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
loss_fn = nn.CrossEntropyLoss()
step = 0
for _ in range(epochs):
model.train()
for xb, yb in train_loader:
step += 1
for group in opt.param_groups:
group['lr'] = lr * min(1.0, step / warmup)
opt.zero_grad()
loss = loss_fn(model(xb).reshape(-1, VOCAB), yb.reshape(-1))
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
model.eval()
with torch.no_grad():
val = loss_fn(model(X_va).reshape(-1, VOCAB),
Y_va.reshape(-1)).item()
return loss.item(), val
import time
torch.manual_seed(0)
model = GPT()
print('parameters %d' % sum(p.numel() for p in model.parameters()))
start = time.time()
train_loss, val_loss = train(model, epochs=5)
import math
print('train loss %.4f validation loss %.4f perplexity %.1f'
% (train_loss, val_loss, math.exp(val_loss)))
print('%.0f seconds' % (time.time() - start))

import torch

def build_corpus(seed=0, sentences=500):
"""A small language with real grammar, generated from a seed.

Written into the page rather than downloaded, so the week runs
offline and everybody trains on identical characters. It is
regular enough that a model can genuinely learn it, and long
enough that counting pairs of characters cannot."""
g = torch.Generator().manual_seed(seed)
def pick(xs):
return xs[int(torch.randint(len(xs), (1,), generator=g))]
subject = ['the model', 'the network', 'the optimiser',
'the gradient', 'the loss', 'attention',
'the encoder', 'the decoder', 'each layer',
'the batch', 'the learning rate', 'dropout']
verb = ['reads', 'predicts', 'updates', 'normalises',
'averages', 'scores', 'reduces', 'controls',
'stabilises', 'copies']
obj = ['the sequence', 'the next token', 'every position',
'the weights', 'the hidden state', 'the input',
'the residual path', 'the attention weights',
'the training data', 'the validation set']
tail = ['during training', 'at every step', 'in a single pass',
'before the softmax', 'after the residual connection',
'without a mask', 'across the batch', 'once per epoch']
joiner = [', and ', ', so ', ', because ', '. ', '. ', '. ']
out = []
for _ in range(sentences):
out.append('%s %s %s %s%s'
% (pick(subject), pick(verb), pick(obj),
pick(tail), pick(joiner)))
return ''.join(out)

CORPUS = build_corpus()
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

text = CORPUS.strip()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
VOCAB = len(chars)

def encode(s):
return torch.tensor([stoi[c] for c in s], dtype=torch.long)

def decode(t):
return ''.join(itos[int(i)] for i in t)

data = encode(text)
BLOCK_SIZE = 32

def windows(seq, block):
"""Every position is a training example: x is block characters,
y is the same block shifted one to the right."""
starts = range(0, len(seq) - block - 1, 4) # stride 4
x = torch.stack([seq[i:i + block] for i in starts])
y = torch.stack([seq[i + 1:i + block + 1] for i in starts])
return x, y

split = int(len(data) * 0.9)
X_tr, Y_tr = windows(data[:split], BLOCK_SIZE)
X_va, Y_va = windows(data[split:], BLOCK_SIZE)
train_loader = DataLoader(TensorDataset(X_tr, Y_tr), batch_size=64,
shuffle=True)
class GPT(nn.Module):
def __init__(self, vocab=None, width=96, heads=4, blocks=2,
block_size=BLOCK_SIZE, dropout=0.1):
super().__init__()
vocab = vocab or VOCAB
self.block_size = block_size
self.tok = nn.Embedding(vocab, width)
self.pos = nn.Embedding(block_size, width)
self.drop = nn.Dropout(dropout)
layer = nn.TransformerEncoderLayer(
d_model=width, nhead=heads, dim_feedforward=width * 4,
dropout=dropout, batch_first=True, norm_first=True,
activation='gelu')
self.blocks = nn.TransformerEncoder(layer, num_layers=blocks)
self.norm = nn.LayerNorm(width)
self.head = nn.Linear(width, vocab, bias=False)
self.head.weight = self.tok.weight

def forward(self, idx):
T = idx.shape[1]
mask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
h = self.drop(self.tok(idx) + self.pos(torch.arange(T)))
return self.head(self.norm(self.blocks(h, mask=mask)))

@torch.no_grad()
def generate(self, idx, steps, temperature=1.0, top_k=None):
self.eval()
for _ in range(steps):
window = idx[:, -self.block_size:]
logits = self(window)[:, -1, :] / temperature
if top_k:
kth = logits.topk(min(top_k, logits.shape[-1])).values[:, -1:]
logits = logits.masked_fill(logits < kth, float('-inf'))
nxt = torch.multinomial(logits.softmax(-1), 1)
idx = torch.cat([idx, nxt], dim=1)
return idx

def train(model, epochs=5, lr=3e-3, warmup=100):
torch.manual_seed(0)
opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
loss_fn = nn.CrossEntropyLoss()
step = 0
for _ in range(epochs):
model.train()
for xb, yb in train_loader:
step += 1
for group in opt.param_groups:
group['lr'] = lr * min(1.0, step / warmup)
opt.zero_grad()
loss = loss_fn(model(xb).reshape(-1, VOCAB), yb.reshape(-1))
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
model.eval()
with torch.no_grad():
val = loss_fn(model(X_va).reshape(-1, VOCAB),
Y_va.reshape(-1)).item()
return loss.item(), val
torch.manual_seed(0)
model = GPT()
train(model, epochs=5)

prompt = encode('the model ').unsqueeze(0)
for temperature in [0.4, 0.8, 1.2]:
torch.manual_seed(0)
out = model.generate(prompt, 140, temperature=temperature)
print('temperature %.1f:' % temperature)
print(' ' + repr(decode(out[0])))
print()

import torch

def build_corpus(seed=0, sentences=500):
"""A small language with real grammar, generated from a seed.

Written into the page rather than downloaded, so the week runs
offline and everybody trains on identical characters. It is
regular enough that a model can genuinely learn it, and long
enough that counting pairs of characters cannot."""
g = torch.Generator().manual_seed(seed)
def pick(xs):
return xs[int(torch.randint(len(xs), (1,), generator=g))]
subject = ['the model', 'the network', 'the optimiser',
'the gradient', 'the loss', 'attention',
'the encoder', 'the decoder', 'each layer',
'the batch', 'the learning rate', 'dropout']
verb = ['reads', 'predicts', 'updates', 'normalises',
'averages', 'scores', 'reduces', 'controls',
'stabilises', 'copies']
obj = ['the sequence', 'the next token', 'every position',
'the weights', 'the hidden state', 'the input',
'the residual path', 'the attention weights',
'the training data', 'the validation set']
tail = ['during training', 'at every step', 'in a single pass',
'before the softmax', 'after the residual connection',
'without a mask', 'across the batch', 'once per epoch']
joiner = [', and ', ', so ', ', because ', '. ', '. ', '. ']
out = []
for _ in range(sentences):
out.append('%s %s %s %s%s'
% (pick(subject), pick(verb), pick(obj),
pick(tail), pick(joiner)))
return ''.join(out)

CORPUS = build_corpus()
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

text = CORPUS.strip()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
VOCAB = len(chars)

def encode(s):
return torch.tensor([stoi[c] for c in s], dtype=torch.long)

def decode(t):
return ''.join(itos[int(i)] for i in t)

data = encode(text)
BLOCK_SIZE = 32

def windows(seq, block):
"""Every position is a training example: x is block characters,
y is the same block shifted one to the right."""
starts = range(0, len(seq) - block - 1, 4) # stride 4
x = torch.stack([seq[i:i + block] for i in starts])
y = torch.stack([seq[i + 1:i + block + 1] for i in starts])
return x, y

split = int(len(data) * 0.9)
X_tr, Y_tr = windows(data[:split], BLOCK_SIZE)
X_va, Y_va = windows(data[split:], BLOCK_SIZE)
train_loader = DataLoader(TensorDataset(X_tr, Y_tr), batch_size=64,
shuffle=True)
class GPT(nn.Module):
def __init__(self, vocab=None, width=96, heads=4, blocks=2,
block_size=BLOCK_SIZE, dropout=0.1):
super().__init__()
vocab = vocab or VOCAB
self.block_size = block_size
self.tok = nn.Embedding(vocab, width)
self.pos = nn.Embedding(block_size, width)
self.drop = nn.Dropout(dropout)
layer = nn.TransformerEncoderLayer(
d_model=width, nhead=heads, dim_feedforward=width * 4,
dropout=dropout, batch_first=True, norm_first=True,
activation='gelu')
self.blocks = nn.TransformerEncoder(layer, num_layers=blocks)
self.norm = nn.LayerNorm(width)
self.head = nn.Linear(width, vocab, bias=False)
self.head.weight = self.tok.weight

def forward(self, idx):
T = idx.shape[1]
mask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
h = self.drop(self.tok(idx) + self.pos(torch.arange(T)))
return self.head(self.norm(self.blocks(h, mask=mask)))

@torch.no_grad()
def generate(self, idx, steps, temperature=1.0, top_k=None):
self.eval()
for _ in range(steps):
window = idx[:, -self.block_size:]
logits = self(window)[:, -1, :] / temperature
if top_k:
kth = logits.topk(min(top_k, logits.shape[-1])).values[:, -1:]
logits = logits.masked_fill(logits < kth, float('-inf'))
nxt = torch.multinomial(logits.softmax(-1), 1)
idx = torch.cat([idx, nxt], dim=1)
return idx

def train(model, epochs=5, lr=3e-3, warmup=100):
torch.manual_seed(0)
opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
loss_fn = nn.CrossEntropyLoss()
step = 0
for _ in range(epochs):
model.train()
for xb, yb in train_loader:
step += 1
for group in opt.param_groups:
group['lr'] = lr * min(1.0, step / warmup)
opt.zero_grad()
loss = loss_fn(model(xb).reshape(-1, VOCAB), yb.reshape(-1))
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
model.eval()
with torch.no_grad():
val = loss_fn(model(X_va).reshape(-1, VOCAB),
Y_va.reshape(-1)).item()
return loss.item(), val
import math
results = []
for block in [8, 16, 32]:
torch.manual_seed(0)
Xa, Ya = windows(data[:split], block)
Xb, Yb = windows(data[split:], block)
loader = DataLoader(TensorDataset(Xa, Ya), batch_size=64, shuffle=True)
model = GPT(block_size=block)
opt = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=0.01)
loss_fn = nn.CrossEntropyLoss()
for _ in range(5):
model.train()
for xb, yb in loader:
opt.zero_grad()
loss_fn(model(xb).reshape(-1, VOCAB), yb.reshape(-1)).backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
model.eval()
with torch.no_grad():
val = loss_fn(model(Xb).reshape(-1, VOCAB), Yb.reshape(-1)).item()
results.append((block, val))

print('%10s %14s %12s' % ('context', 'val loss', 'perplexity'))
for block, val in results:
print('%10d %14.4f %12.1f' % (block, val, math.exp(val)))

import torch

def build_corpus(seed=0, sentences=500):
"""A small language with real grammar, generated from a seed.

Written into the page rather than downloaded, so the week runs
offline and everybody trains on identical characters. It is
regular enough that a model can genuinely learn it, and long
enough that counting pairs of characters cannot."""
g = torch.Generator().manual_seed(seed)
def pick(xs):
return xs[int(torch.randint(len(xs), (1,), generator=g))]
subject = ['the model', 'the network', 'the optimiser',
'the gradient', 'the loss', 'attention',
'the encoder', 'the decoder', 'each layer',
'the batch', 'the learning rate', 'dropout']
verb = ['reads', 'predicts', 'updates', 'normalises',
'averages', 'scores', 'reduces', 'controls',
'stabilises', 'copies']
obj = ['the sequence', 'the next token', 'every position',
'the weights', 'the hidden state', 'the input',
'the residual path', 'the attention weights',
'the training data', 'the validation set']
tail = ['during training', 'at every step', 'in a single pass',
'before the softmax', 'after the residual connection',
'without a mask', 'across the batch', 'once per epoch']
joiner = [', and ', ', so ', ', because ', '. ', '. ', '. ']
out = []
for _ in range(sentences):
out.append('%s %s %s %s%s'
% (pick(subject), pick(verb), pick(obj),
pick(tail), pick(joiner)))
return ''.join(out)

CORPUS = build_corpus()
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

text = CORPUS.strip()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
VOCAB = len(chars)

def encode(s):
return torch.tensor([stoi[c] for c in s], dtype=torch.long)

def decode(t):
return ''.join(itos[int(i)] for i in t)

data = encode(text)
BLOCK_SIZE = 32

def windows(seq, block):
"""Every position is a training example: x is block characters,
y is the same block shifted one to the right."""
starts = range(0, len(seq) - block - 1, 4) # stride 4
x = torch.stack([seq[i:i + block] for i in starts])
y = torch.stack([seq[i + 1:i + block + 1] for i in starts])
return x, y

split = int(len(data) * 0.9)
X_tr, Y_tr = windows(data[:split], BLOCK_SIZE)
X_va, Y_va = windows(data[split:], BLOCK_SIZE)
train_loader = DataLoader(TensorDataset(X_tr, Y_tr), batch_size=64,
shuffle=True)
class GPT(nn.Module):
def __init__(self, vocab=None, width=96, heads=4, blocks=2,
block_size=BLOCK_SIZE, dropout=0.1):
super().__init__()
vocab = vocab or VOCAB
self.block_size = block_size
self.tok = nn.Embedding(vocab, width)
self.pos = nn.Embedding(block_size, width)
self.drop = nn.Dropout(dropout)
layer = nn.TransformerEncoderLayer(
d_model=width, nhead=heads, dim_feedforward=width * 4,
dropout=dropout, batch_first=True, norm_first=True,
activation='gelu')
self.blocks = nn.TransformerEncoder(layer, num_layers=blocks)
self.norm = nn.LayerNorm(width)
self.head = nn.Linear(width, vocab, bias=False)
self.head.weight = self.tok.weight

def forward(self, idx):
T = idx.shape[1]
mask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
h = self.drop(self.tok(idx) + self.pos(torch.arange(T)))
return self.head(self.norm(self.blocks(h, mask=mask)))

@torch.no_grad()
def generate(self, idx, steps, temperature=1.0, top_k=None):
self.eval()
for _ in range(steps):
window = idx[:, -self.block_size:]
logits = self(window)[:, -1, :] / temperature
if top_k:
kth = logits.topk(min(top_k, logits.shape[-1])).values[:, -1:]
logits = logits.masked_fill(logits < kth, float('-inf'))
nxt = torch.multinomial(logits.softmax(-1), 1)
idx = torch.cat([idx, nxt], dim=1)
return idx

def train(model, epochs=5, lr=3e-3, warmup=100):
torch.manual_seed(0)
opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
loss_fn = nn.CrossEntropyLoss()
step = 0
for _ in range(epochs):
model.train()
for xb, yb in train_loader:
step += 1
for group in opt.param_groups:
group['lr'] = lr * min(1.0, step / warmup)
opt.zero_grad()
loss = loss_fn(model(xb).reshape(-1, VOCAB), yb.reshape(-1))
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
model.eval()
with torch.no_grad():
val = loss_fn(model(X_va).reshape(-1, VOCAB),
Y_va.reshape(-1)).item()
return loss.item(), val
import math
print('%8s %8s %10s %12s %12s'
% ('width', 'blocks', 'params', 'train loss', 'val loss'))
for width, blocks in [(48, 2), (96, 2), (96, 4)]:
torch.manual_seed(0)
model = GPT(width=width, blocks=blocks)
tr, va = train(model, epochs=5)
print('%8d %8d %10d %12.4f %12.4f'
% (width, blocks, sum(p.numel() for p in model.parameters()),
tr, va))

import torch

def build_corpus(seed=0, sentences=500):
"""A small language with real grammar, generated from a seed.

Written into the page rather than downloaded, so the week runs
offline and everybody trains on identical characters. It is
regular enough that a model can genuinely learn it, and long
enough that counting pairs of characters cannot."""
g = torch.Generator().manual_seed(seed)
def pick(xs):
return xs[int(torch.randint(len(xs), (1,), generator=g))]
subject = ['the model', 'the network', 'the optimiser',
'the gradient', 'the loss', 'attention',
'the encoder', 'the decoder', 'each layer',
'the batch', 'the learning rate', 'dropout']
verb = ['reads', 'predicts', 'updates', 'normalises',
'averages', 'scores', 'reduces', 'controls',
'stabilises', 'copies']
obj = ['the sequence', 'the next token', 'every position',
'the weights', 'the hidden state', 'the input',
'the residual path', 'the attention weights',
'the training data', 'the validation set']
tail = ['during training', 'at every step', 'in a single pass',
'before the softmax', 'after the residual connection',
'without a mask', 'across the batch', 'once per epoch']
joiner = [', and ', ', so ', ', because ', '. ', '. ', '. ']
out = []
for _ in range(sentences):
out.append('%s %s %s %s%s'
% (pick(subject), pick(verb), pick(obj),
pick(tail), pick(joiner)))
return ''.join(out)

CORPUS = build_corpus()
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

text = CORPUS.strip()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
VOCAB = len(chars)

def encode(s):
return torch.tensor([stoi[c] for c in s], dtype=torch.long)

def decode(t):
return ''.join(itos[int(i)] for i in t)

data = encode(text)
BLOCK_SIZE = 32

def windows(seq, block):
"""Every position is a training example: x is block characters,
y is the same block shifted one to the right."""
starts = range(0, len(seq) - block - 1, 4) # stride 4
x = torch.stack([seq[i:i + block] for i in starts])
y = torch.stack([seq[i + 1:i + block + 1] for i in starts])
return x, y

split = int(len(data) * 0.9)
X_tr, Y_tr = windows(data[:split], BLOCK_SIZE)
X_va, Y_va = windows(data[split:], BLOCK_SIZE)
train_loader = DataLoader(TensorDataset(X_tr, Y_tr), batch_size=64,
shuffle=True)
class GPT(nn.Module):
def __init__(self, vocab=None, width=96, heads=4, blocks=2,
block_size=BLOCK_SIZE, dropout=0.1):
super().__init__()
vocab = vocab or VOCAB
self.block_size = block_size
self.tok = nn.Embedding(vocab, width)
self.pos = nn.Embedding(block_size, width)
self.drop = nn.Dropout(dropout)
layer = nn.TransformerEncoderLayer(
d_model=width, nhead=heads, dim_feedforward=width * 4,
dropout=dropout, batch_first=True, norm_first=True,
activation='gelu')
self.blocks = nn.TransformerEncoder(layer, num_layers=blocks)
self.norm = nn.LayerNorm(width)
self.head = nn.Linear(width, vocab, bias=False)
self.head.weight = self.tok.weight

def forward(self, idx):
T = idx.shape[1]
mask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
h = self.drop(self.tok(idx) + self.pos(torch.arange(T)))
return self.head(self.norm(self.blocks(h, mask=mask)))

@torch.no_grad()
def generate(self, idx, steps, temperature=1.0, top_k=None):
self.eval()
for _ in range(steps):
window = idx[:, -self.block_size:]
logits = self(window)[:, -1, :] / temperature
if top_k:
kth = logits.topk(min(top_k, logits.shape[-1])).values[:, -1:]
logits = logits.masked_fill(logits < kth, float('-inf'))
nxt = torch.multinomial(logits.softmax(-1), 1)
idx = torch.cat([idx, nxt], dim=1)
return idx

def train(model, epochs=5, lr=3e-3, warmup=100):
torch.manual_seed(0)
opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
loss_fn = nn.CrossEntropyLoss()
step = 0
for _ in range(epochs):
model.train()
for xb, yb in train_loader:
step += 1
for group in opt.param_groups:
group['lr'] = lr * min(1.0, step / warmup)
opt.zero_grad()
loss = loss_fn(model(xb).reshape(-1, VOCAB), yb.reshape(-1))
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
model.eval()
with torch.no_grad():
val = loss_fn(model(X_va).reshape(-1, VOCAB),
Y_va.reshape(-1)).item()
return loss.item(), val
torch.manual_seed(0)
model = GPT()
train(model, epochs=5)
model.eval()

for prompt in ['the gradient ', 'attention ', 'if the loss ']:
idx = encode(prompt).unsqueeze(0)
with torch.no_grad():
probs = model(idx)[0, -1].softmax(-1)
top = probs.topk(4)
guesses = ', '.join('%r %.2f' % (itos[int(i)], v)
for v, i in zip(top.values, top.indices))
print('%-16r -> %s' % (prompt, guesses))

import torch

def build_corpus(seed=0, sentences=500):
"""A small language with real grammar, generated from a seed.

Written into the page rather than downloaded, so the week runs
offline and everybody trains on identical characters. It is
regular enough that a model can genuinely learn it, and long
enough that counting pairs of characters cannot."""
g = torch.Generator().manual_seed(seed)
def pick(xs):
return xs[int(torch.randint(len(xs), (1,), generator=g))]
subject = ['the model', 'the network', 'the optimiser',
'the gradient', 'the loss', 'attention',
'the encoder', 'the decoder', 'each layer',
'the batch', 'the learning rate', 'dropout']
verb = ['reads', 'predicts', 'updates', 'normalises',
'averages', 'scores', 'reduces', 'controls',
'stabilises', 'copies']
obj = ['the sequence', 'the next token', 'every position',
'the weights', 'the hidden state', 'the input',
'the residual path', 'the attention weights',
'the training data', 'the validation set']
tail = ['during training', 'at every step', 'in a single pass',
'before the softmax', 'after the residual connection',
'without a mask', 'across the batch', 'once per epoch']
joiner = [', and ', ', so ', ', because ', '. ', '. ', '. ']
out = []
for _ in range(sentences):
out.append('%s %s %s %s%s'
% (pick(subject), pick(verb), pick(obj),
pick(tail), pick(joiner)))
return ''.join(out)

CORPUS = build_corpus()
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

text = CORPUS.strip()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
VOCAB = len(chars)

def encode(s):
return torch.tensor([stoi[c] for c in s], dtype=torch.long)

def decode(t):
return ''.join(itos[int(i)] for i in t)

data = encode(text)
BLOCK_SIZE = 32

def windows(seq, block):
"""Every position is a training example: x is block characters,
y is the same block shifted one to the right."""
starts = range(0, len(seq) - block - 1, 4) # stride 4
x = torch.stack([seq[i:i + block] for i in starts])
y = torch.stack([seq[i + 1:i + block + 1] for i in starts])
return x, y

split = int(len(data) * 0.9)
X_tr, Y_tr = windows(data[:split], BLOCK_SIZE)
X_va, Y_va = windows(data[split:], BLOCK_SIZE)
train_loader = DataLoader(TensorDataset(X_tr, Y_tr), batch_size=64,
shuffle=True)
class GPT(nn.Module):
def __init__(self, vocab=None, width=96, heads=4, blocks=2,
block_size=BLOCK_SIZE, dropout=0.1):
super().__init__()
vocab = vocab or VOCAB
self.block_size = block_size
self.tok = nn.Embedding(vocab, width)
self.pos = nn.Embedding(block_size, width)
self.drop = nn.Dropout(dropout)
layer = nn.TransformerEncoderLayer(
d_model=width, nhead=heads, dim_feedforward=width * 4,
dropout=dropout, batch_first=True, norm_first=True,
activation='gelu')
self.blocks = nn.TransformerEncoder(layer, num_layers=blocks)
self.norm = nn.LayerNorm(width)
self.head = nn.Linear(width, vocab, bias=False)
self.head.weight = self.tok.weight

def forward(self, idx):
T = idx.shape[1]
mask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
h = self.drop(self.tok(idx) + self.pos(torch.arange(T)))
return self.head(self.norm(self.blocks(h, mask=mask)))

@torch.no_grad()
def generate(self, idx, steps, temperature=1.0, top_k=None):
self.eval()
for _ in range(steps):
window = idx[:, -self.block_size:]
logits = self(window)[:, -1, :] / temperature
if top_k:
kth = logits.topk(min(top_k, logits.shape[-1])).values[:, -1:]
logits = logits.masked_fill(logits < kth, float('-inf'))
nxt = torch.multinomial(logits.softmax(-1), 1)
idx = torch.cat([idx, nxt], dim=1)
return idx

def train(model, epochs=5, lr=3e-3, warmup=100):
torch.manual_seed(0)
opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
loss_fn = nn.CrossEntropyLoss()
step = 0
for _ in range(epochs):
model.train()
for xb, yb in train_loader:
step += 1
for group in opt.param_groups:
group['lr'] = lr * min(1.0, step / warmup)
opt.zero_grad()
loss = loss_fn(model(xb).reshape(-1, VOCAB), yb.reshape(-1))
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
model.eval()
with torch.no_grad():
val = loss_fn(model(X_va).reshape(-1, VOCAB),
Y_va.reshape(-1)).item()
return loss.item(), val
import math
torch.manual_seed(0)
model = GPT()
tr, va = train(model, epochs=5)

print('cross entropy loss %.4f' % va)
print('perplexity %.2f' % math.exp(va))
print('bits per character %.3f' % (va / math.log(2)))
print()
print('for comparison:')
print(' uniform over %d characters is %.3f bits'
% (VOCAB, math.log(VOCAB) / math.log(2)))
print(' good character models of English reach roughly 1.1 to 1.3 bits')
print(' on large corpora, which is the number to have in mind.')

import torch

def build_corpus(seed=0, sentences=500):
"""A small language with real grammar, generated from a seed.

Written into the page rather than downloaded, so the week runs
offline and everybody trains on identical characters. It is
regular enough that a model can genuinely learn it, and long
enough that counting pairs of characters cannot."""
g = torch.Generator().manual_seed(seed)
def pick(xs):
return xs[int(torch.randint(len(xs), (1,), generator=g))]
subject = ['the model', 'the network', 'the optimiser',
'the gradient', 'the loss', 'attention',
'the encoder', 'the decoder', 'each layer',
'the batch', 'the learning rate', 'dropout']
verb = ['reads', 'predicts', 'updates', 'normalises',
'averages', 'scores', 'reduces', 'controls',
'stabilises', 'copies']
obj = ['the sequence', 'the next token', 'every position',
'the weights', 'the hidden state', 'the input',
'the residual path', 'the attention weights',
'the training data', 'the validation set']
tail = ['during training', 'at every step', 'in a single pass',
'before the softmax', 'after the residual connection',
'without a mask', 'across the batch', 'once per epoch']
joiner = [', and ', ', so ', ', because ', '. ', '. ', '. ']
out = []
for _ in range(sentences):
out.append('%s %s %s %s%s'
% (pick(subject), pick(verb), pick(obj),
pick(tail), pick(joiner)))
return ''.join(out)

CORPUS = build_corpus()
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

text = CORPUS.strip()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
VOCAB = len(chars)

def encode(s):
return torch.tensor([stoi[c] for c in s], dtype=torch.long)

def decode(t):
return ''.join(itos[int(i)] for i in t)

data = encode(text)
BLOCK_SIZE = 32

def windows(seq, block):
"""Every position is a training example: x is block characters,
y is the same block shifted one to the right."""
starts = range(0, len(seq) - block - 1, 4) # stride 4
x = torch.stack([seq[i:i + block] for i in starts])
y = torch.stack([seq[i + 1:i + block + 1] for i in starts])
return x, y

split = int(len(data) * 0.9)
X_tr, Y_tr = windows(data[:split], BLOCK_SIZE)
X_va, Y_va = windows(data[split:], BLOCK_SIZE)
train_loader = DataLoader(TensorDataset(X_tr, Y_tr), batch_size=64,
shuffle=True)
class GPT(nn.Module):
def __init__(self, vocab=None, width=96, heads=4, blocks=2,
block_size=BLOCK_SIZE, dropout=0.1):
super().__init__()
vocab = vocab or VOCAB
self.block_size = block_size
self.tok = nn.Embedding(vocab, width)
self.pos = nn.Embedding(block_size, width)
self.drop = nn.Dropout(dropout)
layer = nn.TransformerEncoderLayer(
d_model=width, nhead=heads, dim_feedforward=width * 4,
dropout=dropout, batch_first=True, norm_first=True,
activation='gelu')
self.blocks = nn.TransformerEncoder(layer, num_layers=blocks)
self.norm = nn.LayerNorm(width)
self.head = nn.Linear(width, vocab, bias=False)
self.head.weight = self.tok.weight

def forward(self, idx):
T = idx.shape[1]
mask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
h = self.drop(self.tok(idx) + self.pos(torch.arange(T)))
return self.head(self.norm(self.blocks(h, mask=mask)))

@torch.no_grad()
def generate(self, idx, steps, temperature=1.0, top_k=None):
self.eval()
for _ in range(steps):
window = idx[:, -self.block_size:]
logits = self(window)[:, -1, :] / temperature
if top_k:
kth = logits.topk(min(top_k, logits.shape[-1])).values[:, -1:]
logits = logits.masked_fill(logits < kth, float('-inf'))
nxt = torch.multinomial(logits.softmax(-1), 1)
idx = torch.cat([idx, nxt], dim=1)
return idx

def train(model, epochs=5, lr=3e-3, warmup=100):
torch.manual_seed(0)
opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
loss_fn = nn.CrossEntropyLoss()
step = 0
for _ in range(epochs):
model.train()
for xb, yb in train_loader:
step += 1
for group in opt.param_groups:
group['lr'] = lr * min(1.0, step / warmup)
opt.zero_grad()
loss = loss_fn(model(xb).reshape(-1, VOCAB), yb.reshape(-1))
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
model.eval()
with torch.no_grad():
val = loss_fn(model(X_va).reshape(-1, VOCAB),
Y_va.reshape(-1)).item()
return loss.item(), val
torch.manual_seed(0)
model = GPT()
train(model, epochs=5)

# A model can have a fine loss and still produce degenerate text if the
# sampling is wrong. Same weights, three settings.
prompt = encode('the loss ').unsqueeze(0)
for name, kwargs in [('greedy (temperature 0.01)', dict(temperature=0.01)),
('temperature 0.8', dict(temperature=0.8)),
('temperature 0.8, top-k 5',
dict(temperature=0.8, top_k=5)),
('temperature 2.0', dict(temperature=2.0))]:
torch.manual_seed(0)
out = model.generate(prompt, 90, **kwargs)
print('%-26s %r' % (name, decode(out[0])))
print()

import torch

def build_corpus(seed=0, sentences=500):
"""A small language with real grammar, generated from a seed.

Written into the page rather than downloaded, so the week runs
offline and everybody trains on identical characters. It is
regular enough that a model can genuinely learn it, and long
enough that counting pairs of characters cannot."""
g = torch.Generator().manual_seed(seed)
def pick(xs):
return xs[int(torch.randint(len(xs), (1,), generator=g))]
subject = ['the model', 'the network', 'the optimiser',
'the gradient', 'the loss', 'attention',
'the encoder', 'the decoder', 'each layer',
'the batch', 'the learning rate', 'dropout']
verb = ['reads', 'predicts', 'updates', 'normalises',
'averages', 'scores', 'reduces', 'controls',
'stabilises', 'copies']
obj = ['the sequence', 'the next token', 'every position',
'the weights', 'the hidden state', 'the input',
'the residual path', 'the attention weights',
'the training data', 'the validation set']
tail = ['during training', 'at every step', 'in a single pass',
'before the softmax', 'after the residual connection',
'without a mask', 'across the batch', 'once per epoch']
joiner = [', and ', ', so ', ', because ', '. ', '. ', '. ']
out = []
for _ in range(sentences):
out.append('%s %s %s %s%s'
% (pick(subject), pick(verb), pick(obj),
pick(tail), pick(joiner)))
return ''.join(out)

CORPUS = build_corpus()
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

text = CORPUS.strip()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
VOCAB = len(chars)

def encode(s):
return torch.tensor([stoi[c] for c in s], dtype=torch.long)

def decode(t):
return ''.join(itos[int(i)] for i in t)

data = encode(text)
BLOCK_SIZE = 32

def windows(seq, block):
"""Every position is a training example: x is block characters,
y is the same block shifted one to the right."""
starts = range(0, len(seq) - block - 1, 4) # stride 4
x = torch.stack([seq[i:i + block] for i in starts])
y = torch.stack([seq[i + 1:i + block + 1] for i in starts])
return x, y

split = int(len(data) * 0.9)
X_tr, Y_tr = windows(data[:split], BLOCK_SIZE)
X_va, Y_va = windows(data[split:], BLOCK_SIZE)
train_loader = DataLoader(TensorDataset(X_tr, Y_tr), batch_size=64,
shuffle=True)
class GPT(nn.Module):
def __init__(self, vocab=None, width=96, heads=4, blocks=2,
block_size=BLOCK_SIZE, dropout=0.1):
super().__init__()
vocab = vocab or VOCAB
self.block_size = block_size
self.tok = nn.Embedding(vocab, width)
self.pos = nn.Embedding(block_size, width)
self.drop = nn.Dropout(dropout)
layer = nn.TransformerEncoderLayer(
d_model=width, nhead=heads, dim_feedforward=width * 4,
dropout=dropout, batch_first=True, norm_first=True,
activation='gelu')
self.blocks = nn.TransformerEncoder(layer, num_layers=blocks)
self.norm = nn.LayerNorm(width)
self.head = nn.Linear(width, vocab, bias=False)
self.head.weight = self.tok.weight

def forward(self, idx):
T = idx.shape[1]
mask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
h = self.drop(self.tok(idx) + self.pos(torch.arange(T)))
return self.head(self.norm(self.blocks(h, mask=mask)))

@torch.no_grad()
def generate(self, idx, steps, temperature=1.0, top_k=None):
self.eval()
for _ in range(steps):
window = idx[:, -self.block_size:]
logits = self(window)[:, -1, :] / temperature
if top_k:
kth = logits.topk(min(top_k, logits.shape[-1])).values[:, -1:]
logits = logits.masked_fill(logits < kth, float('-inf'))
nxt = torch.multinomial(logits.softmax(-1), 1)
idx = torch.cat([idx, nxt], dim=1)
return idx

def train(model, epochs=5, lr=3e-3, warmup=100):
torch.manual_seed(0)
opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
loss_fn = nn.CrossEntropyLoss()
step = 0
for _ in range(epochs):
model.train()
for xb, yb in train_loader:
step += 1
for group in opt.param_groups:
group['lr'] = lr * min(1.0, step / warmup)
opt.zero_grad()
loss = loss_fn(model(xb).reshape(-1, VOCAB), yb.reshape(-1))
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
model.eval()
with torch.no_grad():
val = loss_fn(model(X_va).reshape(-1, VOCAB),
Y_va.reshape(-1)).item()
return loss.item(), val
torch.manual_seed(0)
model = GPT()
train(model, epochs=5)

def longest_copied(generated, source, minimum=8):
best = ''
for i in range(len(generated)):
for j in range(i + minimum, len(generated) + 1):
piece = generated[i:j]
if piece in source and len(piece) > len(best):
best = piece
elif piece not in source:
break
return best

torch.manual_seed(0)
out = decode(model.generate(encode('the ').unsqueeze(0), 200,
temperature=0.8)[0])
copied = longest_copied(out, text)
print('generated %d characters' % len(out))
print('longest stretch also in the training text: %d characters'
% len(copied))
print(repr(copied))

import torch

def build_corpus(seed=0, sentences=500):
"""A small language with real grammar, generated from a seed.

Written into the page rather than downloaded, so the week runs
offline and everybody trains on identical characters. It is
regular enough that a model can genuinely learn it, and long
enough that counting pairs of characters cannot."""
g = torch.Generator().manual_seed(seed)
def pick(xs):
return xs[int(torch.randint(len(xs), (1,), generator=g))]
subject = ['the model', 'the network', 'the optimiser',
'the gradient', 'the loss', 'attention',
'the encoder', 'the decoder', 'each layer',
'the batch', 'the learning rate', 'dropout']
verb = ['reads', 'predicts', 'updates', 'normalises',
'averages', 'scores', 'reduces', 'controls',
'stabilises', 'copies']
obj = ['the sequence', 'the next token', 'every position',
'the weights', 'the hidden state', 'the input',
'the residual path', 'the attention weights',
'the training data', 'the validation set']
tail = ['during training', 'at every step', 'in a single pass',
'before the softmax', 'after the residual connection',
'without a mask', 'across the batch', 'once per epoch']
joiner = [', and ', ', so ', ', because ', '. ', '. ', '. ']
out = []
for _ in range(sentences):
out.append('%s %s %s %s%s'
% (pick(subject), pick(verb), pick(obj),
pick(tail), pick(joiner)))
return ''.join(out)

CORPUS = build_corpus()
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

text = CORPUS.strip()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
VOCAB = len(chars)

def encode(s):
return torch.tensor([stoi[c] for c in s], dtype=torch.long)

def decode(t):
return ''.join(itos[int(i)] for i in t)

data = encode(text)
BLOCK_SIZE = 32

def windows(seq, block):
"""Every position is a training example: x is block characters,
y is the same block shifted one to the right."""
starts = range(0, len(seq) - block - 1, 4) # stride 4
x = torch.stack([seq[i:i + block] for i in starts])
y = torch.stack([seq[i + 1:i + block + 1] for i in starts])
return x, y

split = int(len(data) * 0.9)
X_tr, Y_tr = windows(data[:split], BLOCK_SIZE)
X_va, Y_va = windows(data[split:], BLOCK_SIZE)
train_loader = DataLoader(TensorDataset(X_tr, Y_tr), batch_size=64,
shuffle=True)
words = text.split()
vocab_words = sorted(set(words))
print('characters: vocabulary %d, sequence length %d'
% (VOCAB, len(text)))
print('words: vocabulary %d, sequence length %d'
% (len(vocab_words), len(words)))
print()
print('a word vocabulary is short sequences and a huge embedding table,')
print('and it has no answer at all for a word it has never seen.')
print('characters are the opposite: tiny table, very long sequences,')
print('and attention costs the square of the length.')

import torch

def build_corpus(seed=0, sentences=500):
"""A small language with real grammar, generated from a seed.

Written into the page rather than downloaded, so the week runs
offline and everybody trains on identical characters. It is
regular enough that a model can genuinely learn it, and long
enough that counting pairs of characters cannot."""
g = torch.Generator().manual_seed(seed)
def pick(xs):
return xs[int(torch.randint(len(xs), (1,), generator=g))]
subject = ['the model', 'the network', 'the optimiser',
'the gradient', 'the loss', 'attention',
'the encoder', 'the decoder', 'each layer',
'the batch', 'the learning rate', 'dropout']
verb = ['reads', 'predicts', 'updates', 'normalises',
'averages', 'scores', 'reduces', 'controls',
'stabilises', 'copies']
obj = ['the sequence', 'the next token', 'every position',
'the weights', 'the hidden state', 'the input',
'the residual path', 'the attention weights',
'the training data', 'the validation set']
tail = ['during training', 'at every step', 'in a single pass',
'before the softmax', 'after the residual connection',
'without a mask', 'across the batch', 'once per epoch']
joiner = [', and ', ', so ', ', because ', '. ', '. ', '. ']
out = []
for _ in range(sentences):
out.append('%s %s %s %s%s'
% (pick(subject), pick(verb), pick(obj),
pick(tail), pick(joiner)))
return ''.join(out)

CORPUS = build_corpus()
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

text = CORPUS.strip()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
VOCAB = len(chars)

def encode(s):
return torch.tensor([stoi[c] for c in s], dtype=torch.long)

def decode(t):
return ''.join(itos[int(i)] for i in t)

data = encode(text)
BLOCK_SIZE = 32

def windows(seq, block):
"""Every position is a training example: x is block characters,
y is the same block shifted one to the right."""
starts = range(0, len(seq) - block - 1, 4) # stride 4
x = torch.stack([seq[i:i + block] for i in starts])
y = torch.stack([seq[i + 1:i + block + 1] for i in starts])
return x, y

split = int(len(data) * 0.9)
X_tr, Y_tr = windows(data[:split], BLOCK_SIZE)
X_va, Y_va = windows(data[split:], BLOCK_SIZE)
train_loader = DataLoader(TensorDataset(X_tr, Y_tr), batch_size=64,
shuffle=True)
from collections import Counter

# Start from characters and repeatedly merge the commonest adjacent pair.
tokens = list(text[:900])
merges = []
for step in range(12):
pairs = Counter(zip(tokens, tokens[1:]))
if not pairs:
break
best, count = pairs.most_common(1)[0]
merged = best[0] + best[1]
merges.append((best, count))
out, i = [], 0
while i < len(tokens):
if i < len(tokens) - 1 and (tokens[i], tokens[i + 1]) == best:
out.append(merged)
i += 2
else:
out.append(tokens[i])
i += 1
tokens = out

print('%-6s %-14s %8s' % ('merge', 'pair joined', 'count'))
for i, (pair, count) in enumerate(merges, 1):
print('%-6d %-14r %8d' % (i, pair[0] + pair[1], count))
print('\n900 characters became %d tokens' % len(tokens))
print('first 20:', tokens[:20])

from tokenizers import Tokenizer, models, trainers, pre_tokenizers

text = ('the model reads the sequence and predicts the next token. '
'training is a loop: forward, loss, backward, step. the gradient '
'tells each weight which way to move and the optimiser decides '
'how far. attention scores every position against every other '
'position, and the scores become weights.') * 8

tok = Tokenizer(models.BPE(unk_token='[UNK]'))
tok.pre_tokenizer = pre_tokenizers.Whitespace()
trainer = trainers.BpeTrainer(vocab_size=120,
special_tokens=['[UNK]', '[PAD]'])
tok.train_from_iterator([text], trainer)

print('vocabulary size %d' % tok.get_vocab_size())
for s in ['the gradient', 'backpropagation', 'zqx']:
enc = tok.encode(s)
print('%-18r -> %s' % (s, enc.tokens))

import torch

def build_corpus(seed=0, sentences=500):
"""A small language with real grammar, generated from a seed.

Written into the page rather than downloaded, so the week runs
offline and everybody trains on identical characters. It is
regular enough that a model can genuinely learn it, and long
enough that counting pairs of characters cannot."""
g = torch.Generator().manual_seed(seed)
def pick(xs):
return xs[int(torch.randint(len(xs), (1,), generator=g))]
subject = ['the model', 'the network', 'the optimiser',
'the gradient', 'the loss', 'attention',
'the encoder', 'the decoder', 'each layer',
'the batch', 'the learning rate', 'dropout']
verb = ['reads', 'predicts', 'updates', 'normalises',
'averages', 'scores', 'reduces', 'controls',
'stabilises', 'copies']
obj = ['the sequence', 'the next token', 'every position',
'the weights', 'the hidden state', 'the input',
'the residual path', 'the attention weights',
'the training data', 'the validation set']
tail = ['during training', 'at every step', 'in a single pass',
'before the softmax', 'after the residual connection',
'without a mask', 'across the batch', 'once per epoch']
joiner = [', and ', ', so ', ', because ', '. ', '. ', '. ']
out = []
for _ in range(sentences):
out.append('%s %s %s %s%s'
% (pick(subject), pick(verb), pick(obj),
pick(tail), pick(joiner)))
return ''.join(out)

CORPUS = build_corpus()
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

text = CORPUS.strip()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
VOCAB = len(chars)

def encode(s):
return torch.tensor([stoi[c] for c in s], dtype=torch.long)

def decode(t):
return ''.join(itos[int(i)] for i in t)

data = encode(text)
BLOCK_SIZE = 32

def windows(seq, block):
"""Every position is a training example: x is block characters,
y is the same block shifted one to the right."""
starts = range(0, len(seq) - block - 1, 4) # stride 4
x = torch.stack([seq[i:i + block] for i in starts])
y = torch.stack([seq[i + 1:i + block + 1] for i in starts])
return x, y

split = int(len(data) * 0.9)
X_tr, Y_tr = windows(data[:split], BLOCK_SIZE)
X_va, Y_va = windows(data[split:], BLOCK_SIZE)
train_loader = DataLoader(TensorDataset(X_tr, Y_tr), batch_size=64,
shuffle=True)
# Training on padded documents wastes compute on padding. Real training
# concatenates everything and cuts fixed-length windows out of the stream.
documents = [text[:120], text[200:260], text[400:700]]
lengths = [len(d) for d in documents]
print('document lengths', lengths)

padded = max(lengths) * len(documents)
packed = sum(lengths)
print('padded to a rectangle: %d positions, %d of them real (%.0f%%)'
% (padded, packed, 100 * packed / padded))

stream = encode(''.join(documents))
block = 48
n_windows = len(stream) // block
print('packed into a stream: %d windows of %d, no padding at all'
% (n_windows, block))

def flops_per_token(width, blocks, context, vocab):
attn_proj = 4 * width * width
attn_scores = 2 * context * width
ff = 8 * width * width
per_block = attn_proj + attn_scores + ff
return 2 * (blocks * per_block + width * vocab)

print('%10s %10s %14s %14s' % ('context', 'width', 'per token', 'share attn'))
for context in [128, 1024, 8192]:
for width in [768]:
total = flops_per_token(width, 12, context, 50257)
attn = 2 * 12 * 2 * context * width
print('%10d %10d %14s %13.1f%%'
% (context, width, '{:,}'.format(total),
100 * attn / total))

def budget(params, tokens):
flops = 6 * params * tokens
memory = params * 16 / 1e9
return flops, memory

print('%-16s %14s %14s %12s'
% ('model', 'parameters', 'tokens', 'memory (GB)'))
for name, params in [('small', 124e6), ('medium', 355e6), ('large', 7e9)]:
tokens = params * 20
flops, memory = budget(params, tokens)
print('%-16s %14s %14s %12.1f'
% (name, '{:,.0f}'.format(params), '{:,.0f}'.format(tokens),
memory))
print('%-16s total compute about %.2e operations' % ('', flops))

import torch

def build_corpus(seed=0, sentences=500):
"""A small language with real grammar, generated from a seed.

Written into the page rather than downloaded, so the week runs
offline and everybody trains on identical characters. It is
regular enough that a model can genuinely learn it, and long
enough that counting pairs of characters cannot."""
g = torch.Generator().manual_seed(seed)
def pick(xs):
return xs[int(torch.randint(len(xs), (1,), generator=g))]
subject = ['the model', 'the network', 'the optimiser',
'the gradient', 'the loss', 'attention',
'the encoder', 'the decoder', 'each layer',
'the batch', 'the learning rate', 'dropout']
verb = ['reads', 'predicts', 'updates', 'normalises',
'averages', 'scores', 'reduces', 'controls',
'stabilises', 'copies']
obj = ['the sequence', 'the next token', 'every position',
'the weights', 'the hidden state', 'the input',
'the residual path', 'the attention weights',
'the training data', 'the validation set']
tail = ['during training', 'at every step', 'in a single pass',
'before the softmax', 'after the residual connection',
'without a mask', 'across the batch', 'once per epoch']
joiner = [', and ', ', so ', ', because ', '. ', '. ', '. ']
out = []
for _ in range(sentences):
out.append('%s %s %s %s%s'
% (pick(subject), pick(verb), pick(obj),
pick(tail), pick(joiner)))
return ''.join(out)

CORPUS = build_corpus()
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

text = CORPUS.strip()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
VOCAB = len(chars)

def encode(s):
return torch.tensor([stoi[c] for c in s], dtype=torch.long)

def decode(t):
return ''.join(itos[int(i)] for i in t)

data = encode(text)
BLOCK_SIZE = 32

def windows(seq, block):
"""Every position is a training example: x is block characters,
y is the same block shifted one to the right."""
starts = range(0, len(seq) - block - 1, 4) # stride 4
x = torch.stack([seq[i:i + block] for i in starts])
y = torch.stack([seq[i + 1:i + block + 1] for i in starts])
return x, y

split = int(len(data) * 0.9)
X_tr, Y_tr = windows(data[:split], BLOCK_SIZE)
X_va, Y_va = windows(data[split:], BLOCK_SIZE)
train_loader = DataLoader(TensorDataset(X_tr, Y_tr), batch_size=64,
shuffle=True)
class GPT(nn.Module):
def __init__(self, vocab=None, width=96, heads=4, blocks=2,
block_size=BLOCK_SIZE, dropout=0.1):
super().__init__()
vocab = vocab or VOCAB
self.block_size = block_size
self.tok = nn.Embedding(vocab, width)
self.pos = nn.Embedding(block_size, width)
self.drop = nn.Dropout(dropout)
layer = nn.TransformerEncoderLayer(
d_model=width, nhead=heads, dim_feedforward=width * 4,
dropout=dropout, batch_first=True, norm_first=True,
activation='gelu')
self.blocks = nn.TransformerEncoder(layer, num_layers=blocks)
self.norm = nn.LayerNorm(width)
self.head = nn.Linear(width, vocab, bias=False)
self.head.weight = self.tok.weight

def forward(self, idx):
T = idx.shape[1]
mask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
h = self.drop(self.tok(idx) + self.pos(torch.arange(T)))
return self.head(self.norm(self.blocks(h, mask=mask)))

@torch.no_grad()
def generate(self, idx, steps, temperature=1.0, top_k=None):
self.eval()
for _ in range(steps):
window = idx[:, -self.block_size:]
logits = self(window)[:, -1, :] / temperature
if top_k:
kth = logits.topk(min(top_k, logits.shape[-1])).values[:, -1:]
logits = logits.masked_fill(logits < kth, float('-inf'))
nxt = torch.multinomial(logits.softmax(-1), 1)
idx = torch.cat([idx, nxt], dim=1)
return idx

def train(model, epochs=5, lr=3e-3, warmup=100):
torch.manual_seed(0)
opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
loss_fn = nn.CrossEntropyLoss()
step = 0
for _ in range(epochs):
model.train()
for xb, yb in train_loader:
step += 1
for group in opt.param_groups:
group['lr'] = lr * min(1.0, step / warmup)
opt.zero_grad()
loss = loss_fn(model(xb).reshape(-1, VOCAB), yb.reshape(-1))
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
model.eval()
with torch.no_grad():
val = loss_fn(model(X_va).reshape(-1, VOCAB),
Y_va.reshape(-1)).item()
return loss.item(), val
import math, time

# baseline 1: character frequencies
from collections import Counter
counts = Counter(text)
total = sum(counts.values())
unigram = -sum((n / total) * math.log(n / total) for n in counts.values())

# baseline 2: bigram counts
bi = torch.zeros(VOCAB, VOCAB)
for a, b in zip(data[:split][:-1], data[:split][1:]):
bi[a, b] += 1
probs = (bi + 1) / (bi + 1).sum(1, keepdim=True)
val = data[split:]
bigram = -probs[val[:-1], val[1:]].log().mean().item()

print('%-28s %12s %12s' % ('', 'val loss', 'perplexity'))
print('%-28s %12.4f %12.1f'
% ('uniform guessing', math.log(VOCAB), VOCAB))
print('%-28s %12.4f %12.1f' % ('character frequencies', unigram,
math.exp(unigram)))
print('%-28s %12.4f %12.1f' % ('bigram counts', bigram,
math.exp(bigram)))

for blocks in [1, 2]:
torch.manual_seed(0)
model = GPT(blocks=blocks)
start = time.time()
tr, va = train(model, epochs=5)
print('%-28s %12.4f %12.1f'
% ('transformer, %d blocks' % blocks, va, math.exp(va)))

import torch

def build_corpus(seed=0, sentences=500):
"""A small language with real grammar, generated from a seed.

Written into the page rather than downloaded, so the week runs
offline and everybody trains on identical characters. It is
regular enough that a model can genuinely learn it, and long
enough that counting pairs of characters cannot."""
g = torch.Generator().manual_seed(seed)
def pick(xs):
return xs[int(torch.randint(len(xs), (1,), generator=g))]
subject = ['the model', 'the network', 'the optimiser',
'the gradient', 'the loss', 'attention',
'the encoder', 'the decoder', 'each layer',
'the batch', 'the learning rate', 'dropout']
verb = ['reads', 'predicts', 'updates', 'normalises',
'averages', 'scores', 'reduces', 'controls',
'stabilises', 'copies']
obj = ['the sequence', 'the next token', 'every position',
'the weights', 'the hidden state', 'the input',
'the residual path', 'the attention weights',
'the training data', 'the validation set']
tail = ['during training', 'at every step', 'in a single pass',
'before the softmax', 'after the residual connection',
'without a mask', 'across the batch', 'once per epoch']
joiner = [', and ', ', so ', ', because ', '. ', '. ', '. ']
out = []
for _ in range(sentences):
out.append('%s %s %s %s%s'
% (pick(subject), pick(verb), pick(obj),
pick(tail), pick(joiner)))
return ''.join(out)

CORPUS = build_corpus()
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset

text = CORPUS.strip()
chars = sorted(set(text))
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for c, i in stoi.items()}
VOCAB = len(chars)

def encode(s):
return torch.tensor([stoi[c] for c in s], dtype=torch.long)

def decode(t):
return ''.join(itos[int(i)] for i in t)

data = encode(text)
BLOCK_SIZE = 32

def windows(seq, block):
"""Every position is a training example: x is block characters,
y is the same block shifted one to the right."""
starts = range(0, len(seq) - block - 1, 4) # stride 4
x = torch.stack([seq[i:i + block] for i in starts])
y = torch.stack([seq[i + 1:i + block + 1] for i in starts])
return x, y

split = int(len(data) * 0.9)
X_tr, Y_tr = windows(data[:split], BLOCK_SIZE)
X_va, Y_va = windows(data[split:], BLOCK_SIZE)
train_loader = DataLoader(TensorDataset(X_tr, Y_tr), batch_size=64,
shuffle=True)
class GPT(nn.Module):
def __init__(self, vocab=None, width=96, heads=4, blocks=2,
block_size=BLOCK_SIZE, dropout=0.1):
super().__init__()
vocab = vocab or VOCAB
self.block_size = block_size
self.tok = nn.Embedding(vocab, width)
self.pos = nn.Embedding(block_size, width)
self.drop = nn.Dropout(dropout)
layer = nn.TransformerEncoderLayer(
d_model=width, nhead=heads, dim_feedforward=width * 4,
dropout=dropout, batch_first=True, norm_first=True,
activation='gelu')
self.blocks = nn.TransformerEncoder(layer, num_layers=blocks)
self.norm = nn.LayerNorm(width)
self.head = nn.Linear(width, vocab, bias=False)
self.head.weight = self.tok.weight

def forward(self, idx):
T = idx.shape[1]
mask = torch.triu(torch.ones(T, T, dtype=torch.bool), diagonal=1)
h = self.drop(self.tok(idx) + self.pos(torch.arange(T)))
return self.head(self.norm(self.blocks(h, mask=mask)))

@torch.no_grad()
def generate(self, idx, steps, temperature=1.0, top_k=None):
self.eval()
for _ in range(steps):
window = idx[:, -self.block_size:]
logits = self(window)[:, -1, :] / temperature
if top_k:
kth = logits.topk(min(top_k, logits.shape[-1])).values[:, -1:]
logits = logits.masked_fill(logits < kth, float('-inf'))
nxt = torch.multinomial(logits.softmax(-1), 1)
idx = torch.cat([idx, nxt], dim=1)
return idx

def train(model, epochs=5, lr=3e-3, warmup=100):
torch.manual_seed(0)
opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
loss_fn = nn.CrossEntropyLoss()
step = 0
for _ in range(epochs):
model.train()
for xb, yb in train_loader:
step += 1
for group in opt.param_groups:
group['lr'] = lr * min(1.0, step / warmup)
opt.zero_grad()
loss = loss_fn(model(xb).reshape(-1, VOCAB), yb.reshape(-1))
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
model.eval()
with torch.no_grad():
val = loss_fn(model(X_va).reshape(-1, VOCAB),
Y_va.reshape(-1)).item()
return loss.item(), val
torch.manual_seed(0)
model = GPT(blocks=2)
tr, va = train(model, epochs=5)
print('validation loss %.4f\n' % va)

for prompt in ['the model ', 'attention ', 'if the loss ']:
torch.manual_seed(0)
out = model.generate(encode(prompt).unsqueeze(0), 110,
temperature=0.7, top_k=8)
print('%r ->' % prompt)
print(' ' + repr(decode(out[0])))
print()