Step 7 — Tokenizer and Sampler¶

Solution notebook. Replaces the two toy ends of the pipeline: a byte-pair-encoding tokenizer going in, and a proper sampler coming out. BPE training on 1.1M characters is a few minutes of pure Python; everything else is seconds. Requires step6_model.pt from the Step 6 notebook.

Setup: data and tokenizer (Step 1's solution, reproduced)¶

Every notebook in this series is self-contained: it re-creates what it needs from earlier steps in one compact cell, so you can run it top to bottom without opening the others. On Colab, uncomment the download.

In [ ]:
# !wget -q https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
import torch

with open('input.txt') as f:
    text = f.read()
vocab = sorted(set(text))
V = len(vocab)
stoi = {ch: i for i, ch in enumerate(vocab)}
itos = {i: ch for i, ch in enumerate(vocab)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: ''.join(itos[i] for i in ids)
ids = torch.tensor(encode(text), dtype=torch.long)

assert len(text) == 1_115_394 and V == 65
print(f'{len(text):,} characters, vocab {V}')

7.1 BPE trainer¶

Greedy compression (Lecture 7, Prop. 1.2): repeatedly merge the most frequent adjacent pair. Two helpers — count pairs, and replace every occurrence of one pair in a single left-to-right pass.

In [ ]:
from collections import Counter
import time

def pair_counts(seq):
    return Counter(zip(seq, seq[1:]))

def merge(seq, pair, new_id):
    out, i = [], 0
    while i < len(seq):
        if i < len(seq) - 1 and (seq[i], seq[i+1]) == pair:
            out.append(new_id); i += 2
        else:
            out.append(seq[i]); i += 1
    return out

NUM_MERGES = 256
base_ids = encode(text)
seq = list(base_ids)
merges = {}                         # (id, id) -> new_id, in training order
vocab_s = {i: ch for i, ch in enumerate(vocab)}
sizes = {}

t0 = time.time()
for m in range(NUM_MERGES):
    counts = pair_counts(seq)
    pair, cnt = counts.most_common(1)[0]
    new_id = V + m
    merges[pair] = new_id
    vocab_s[new_id] = vocab_s[pair[0]] + vocab_s[pair[1]]
    seq = merge(seq, pair, new_id)
    if m < 10:
        print(f'  merge {m:>3}: {vocab_s[new_id]!r:<8} ({cnt:,} occurrences)')
    sizes[m+1] = len(seq)
print(f'\ntrained {NUM_MERGES} merges in {time.time()-t0:.0f}s')
print(f'{len(base_ids):,} chars -> {len(seq):,} tokens '
      f'= {len(base_ids)/len(seq):.3f} chars/token')

Reference first ten merges: 'e ', 'th', 't ', 's ', 'd ', ', ', 'ou', 'er', 'in', 'y ' — then an, :\n, or, o , en, \n\n. Three things worth a sentence each:

  1. Word-final patterns dominate the top. The space is the most predictable character in English, so pairs ending in it are the most frequent. Frequency alone finds the word boundary.
  2. :\n and \n\n are the document format — the SPEAKER: convention and the blank line between speeches, discovered with no notion of what a play is.
  3. By merge 28 whole words appear ('and '). BPE crosses from morphology into vocabulary on its own.

Reference compression: 1.500 / 1.963 / 2.286 / 2.692 chars per token at 64 / 256 / 512 / 1024 merges — strongly diminishing returns.

7.2 Encoder and decoder¶

Encoding applies merges in training order; doing them in any other order can give a different tokenization, which is why merges is an ordered dict and not a set.

In [ ]:
def bpe_encode(s):
    out = [stoi[c] for c in s]
    for pair, new_id in merges.items():      # insertion order = training order
        if len(out) < 2:
            break
        out = merge(out, pair, new_id)
    return out

def bpe_decode(toks):
    return ''.join(vocab_s[t] for t in toks)

unseen = ('Shall I compare thee to a summer\'s day?\n'
          'Thou art more lovely and more temperate:')
toks = bpe_encode(unseen)
assert bpe_decode(toks) == unseen                  # exact round-trip
print(f'{len(unseen)} chars -> {len(toks)} tokens')
print([vocab_s[t] for t in toks[:14]])
print('round-trip on unseen text: exact')

7.3 The unit problem¶

Per-token loss is not comparable across tokenizers. Lecture 7's Proposition 2.1 gives the invariant. Our Step 6 character model scored 1.5178 nats/token, but each token was one character; a BPE model's tokens carry ~1.96 characters each, so its per-token loss must be higher to break even.

In [ ]:
import math
def bits_per_char(nats_per_tok, n_tok, n_chars):
    return nats_per_tok / math.log(2) * n_tok / n_chars

print(f"{'model':<26}{'nats/tok':>10}{'bits/char':>11}")
print(f"{'Step 1 bigram (chars)':<26}{2.4549:>10.4f}"
      f"{bits_per_char(2.4549, len(base_ids), len(text)):>11.4f}")
print(f"{'Step 6 GPT (chars)':<26}{1.5178:>10.4f}"
      f"{bits_per_char(1.5178, len(base_ids), len(text)):>11.4f}")
print(f"\nA BPE model would need {1.5178 * len(base_ids)/len(seq):.4f} "
      f"nats/token to match Step 6's bits/char.")

Retraining on BPE tokens is the same Step 6 loop with $V' = 321$ and the retokenized stream — about an hour of CPU, so it is left as an exercise rather than executed here. Convert both runs to bits/char before declaring a winner. A second effect partly offsets the cost: with the same $T_{\max}$, a BPE model's window spans roughly twice as much text.

7.4 The sampler¶

Three knobs. Temperature is entropy-regularized soft argmax (Lecture 7, Prop. 3.1); top-$k$ truncates to a fixed count; nucleus truncates to an adaptive count — the smallest set carrying probability $p$.

In [ ]:
import math
import torch.nn as nn
import torch.nn.functional as F

class Head(nn.Module):
    def __init__(self, d, dh, T):
        super().__init__()
        self.key   = nn.Linear(d, dh, bias=False)
        self.query = nn.Linear(d, dh, bias=False)
        self.value = nn.Linear(d, dh, bias=False)
        self.register_buffer('tril', torch.tril(torch.ones(T, T)))

    def forward(self, x):
        B, T, _ = x.shape
        q, k, v = self.query(x), self.key(x), self.value(x)
        s = q @ k.transpose(-2, -1) / k.shape[-1]**0.5
        s = s.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
        A = F.softmax(s, dim=-1)
        self.A = A.detach()          # stashed for Step 8 interpretability
        return A @ v

class MHA(nn.Module):
    def __init__(self, d, H, T):
        super().__init__()
        self.heads = nn.ModuleList(Head(d, d // H, T) for _ in range(H))
        self.proj = nn.Linear(d, d)

    def forward(self, x):
        return self.proj(torch.cat([h(x) for h in self.heads], dim=-1))

class MLP(nn.Module):
    def __init__(self, d):
        super().__init__()
        self.fc, self.proj = nn.Linear(d, 4 * d), nn.Linear(4 * d, d)

    def forward(self, x):
        return self.proj(F.gelu(self.fc(x)))

class Block(nn.Module):
    def __init__(self, d, H, T):
        super().__init__()
        self.ln1, self.attn = nn.LayerNorm(d), MHA(d, H, T)
        self.ln2, self.mlp  = nn.LayerNorm(d), MLP(d)

    def forward(self, x):
        x = x + self.attn(self.ln1(x))
        return x + self.mlp(self.ln2(x))

class GPT(nn.Module):
    def __init__(self, V, T, d, H, L):
        super().__init__()
        self.T = T
        self.tok = nn.Embedding(V, d)
        self.pos = nn.Embedding(T, d)
        self.blocks = nn.ModuleList(Block(d, H, T) for _ in range(L))
        self.lnf = nn.LayerNorm(d)
        self.head = nn.Linear(d, V, bias=False)
        for blk in self.blocks:   # GPT-2 residual scaling (Lecture 6 §1)
            nn.init.normal_(blk.attn.proj.weight, std=0.02 / math.sqrt(2 * L))
            nn.init.normal_(blk.mlp.proj.weight,  std=0.02 / math.sqrt(2 * L))

    def forward(self, idx, targets=None):
        B, T = idx.shape
        x = self.tok(idx) + self.pos(torch.arange(T))
        for blk in self.blocks:
            x = blk(x)
        logits = self.head(self.lnf(x))
        loss = None if targets is None else F.cross_entropy(
            logits.view(-1, logits.size(-1)), targets.reshape(-1))
        return logits, loss

    @torch.no_grad()
    def generate(self, idx, n_new, temperature=1.0):
        for _ in range(n_new):
            logits, _ = self(idx[:, -self.T:])
            probs = F.softmax(logits[:, -1, :] / temperature, dim=-1)
            idx = torch.cat([idx, torch.multinomial(probs, 1)], dim=1)
        return idx
In [ ]:
torch.manual_seed(1337)
model = GPT(V=65, T=64, d=128, H=4, L=4)
model.load_state_dict(torch.load('step6_model.pt'))
model.eval()
print(f'loaded trained model, {sum(p.numel() for p in model.parameters()):,} parameters')
In [ ]:
import torch.nn.functional as F

def sample_next(logits, temperature=1.0, top_k=None, top_p=None):
    logits = logits.clone()
    if temperature <= 0:                      # greedy = tau -> 0 limit
        return logits.argmax(-1, keepdim=True)
    logits = logits / temperature
    if top_k is not None:
        kth = logits.topk(top_k, dim=-1).values[..., -1:]
        logits = logits.masked_fill(logits < kth, float('-inf'))
    if top_p is not None:
        srt, idx = logits.sort(dim=-1, descending=True)
        cum = srt.softmax(-1).cumsum(-1)
        drop = cum - srt.softmax(-1) >= top_p   # keep the token that crosses p
        srt = srt.masked_fill(drop, float('-inf'))
        logits = torch.full_like(logits, float('-inf')).scatter(-1, idx, srt)
    return torch.multinomial(logits.softmax(-1), 1)

@torch.no_grad()
def generate(prompt, n_new=250, **kw):
    idx = torch.tensor([encode(prompt)])
    for _ in range(n_new):
        logits, _ = model(idx[:, -model.T:])
        idx = torch.cat([idx, sample_next(logits[:, -1, :], **kw)], dim=1)
    return decode(idx[0].tolist())

The decoding gallery¶

Same prompt, same model, seven decoding rules. Watch greedy fall into a loop and $\tau=1.5$ lose English.

In [ ]:
settings = [('greedy (tau->0)', dict(temperature=0.0)),
            ('tau = 0.3',       dict(temperature=0.3)),
            ('tau = 0.8',       dict(temperature=0.8)),
            ('tau = 1.0',       dict(temperature=1.0)),
            ('tau = 1.5',       dict(temperature=1.5)),
            ('top-k = 5',       dict(temperature=1.0, top_k=5)),
            ('nucleus p = 0.9', dict(temperature=1.0, top_p=0.9))]
for name, kw in settings:
    torch.manual_seed(0)
    print('=' * 70); print(name); print('=' * 70)
    print(generate('ROMEO:', 200, **kw)); print()

Greedy degenerates — it locks into a repeated phrase and stays there. This is Holtzman et al.'s neural text degeneration, on your own model. The explanation (Lecture 7 §3) is that the mode of a high-dimensional distribution is wildly unrepresentative of it: natural text is typical, not most-probable, and maximizing likelihood lands you outside the region where real text lives.

The objective we train and the objective we want at decode time genuinely differ. That seam is what Lecture 8's alignment methods work on.

Entropy of the decoding distributions¶

Quantifying what each knob does: the mean entropy of the next-token distribution, and how many tokens the nucleus actually keeps.

In [ ]:
idx = torch.tensor([encode('ROMEO:\nWhat light through yonder')])
with torch.no_grad():
    logits = model(idx)[0][0, -1, :]
print(f"{'setting':<20}{'entropy (nats)':>16}{'effective support':>20}")
for name, kw in settings[1:]:
    lg = logits / kw.get('temperature', 1.0)
    if kw.get('top_k'):
        kth = lg.topk(kw['top_k']).values[-1]
        lg = lg.masked_fill(lg < kth, float('-inf'))
    if kw.get('top_p'):
        srt, i2 = lg.sort(descending=True)
        cum = srt.softmax(-1).cumsum(-1)
        srt = srt.masked_fill(cum - srt.softmax(-1) >= kw['top_p'], float('-inf'))
        lg = torch.full_like(lg, float('-inf')).scatter(-1, i2, srt)
    p = lg.softmax(-1)
    H = -(p[p > 0] * p[p > 0].log()).sum()
    print(f'{name:<20}{H:>16.4f}{int((p > 1e-6).sum()):>20}')

→ Continue with Step 8: the capstone.