# !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}')

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

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')

import torch.nn.functional as F
import matplotlib.pyplot as plt

T = 64
@torch.no_grad()
def head_patterns(n_seq=32, seed=0):
    torch.manual_seed(seed)
    acc = torch.zeros(len(model.blocks), len(model.blocks[0].attn.heads), T, T)
    for _ in range(n_seq):
        i = torch.randint(0, len(ids) - T - 1, (1,)).item()
        x = model.tok(ids[i:i+T][None]) + model.pos(torch.arange(T))
        for li, blk in enumerate(model.blocks):
            xn = blk.ln1(x)
            for hi, h in enumerate(blk.attn.heads):
                h(xn)
                acc[li, hi] += h.A[0]
            x = blk(x)
    return acc / n_seq

pat = head_patterns()
L, H = pat.shape[:2]
fig, axes = plt.subplots(L, H, figsize=(2.1*H, 2.1*L))
for li in range(L):
    for hi in range(H):
        ax = axes[li, hi]
        ax.imshow(pat[li, hi], cmap='Blues')
        ax.set_xticks([]); ax.set_yticks([])
        if hi == 0: ax.set_ylabel(f'layer {li}', fontsize=9)
        if li == 0: ax.set_title(f'head {hi}', fontsize=9)
plt.suptitle('mean attention pattern per head, averaged over 32 real sequences')
plt.tight_layout()

# Quantify: how much mass sits exactly one position back, on the diagonal,
# or on the very first token (a common 'attention sink')?
print(f"{'head':<10}{'prev-token':>12}{'self':>8}{'first-tok':>11}  guess")
for li in range(L):
    for hi in range(H):
        A = pat[li, hi]
        rows = torch.arange(2, T)
        prev = A[rows, rows-1].mean().item()
        self_ = A[rows, rows].mean().item()
        first = A[rows, 0].mean().item()
        guess = ('previous-token' if prev > 0.35 else
                 'self/current'   if self_ > 0.35 else
                 'first-token sink' if first > 0.35 else 'diffuse')
        print(f'L{li}H{hi:<7}{prev:>12.3f}{self_:>8.3f}{first:>11.3f}  {guess}')

import copy

n_split = int(0.9 * len(ids))
val_ids = ids[n_split:]

@torch.no_grad()
def val_loss(m, iters=60, B=64, T=64, seed=0):
    torch.manual_seed(seed)
    tot = 0.0
    for _ in range(iters):
        ix = torch.randint(len(val_ids) - T - 1, (B,))
        xb = torch.stack([val_ids[i:i+T] for i in ix])
        yb = torch.stack([val_ids[i+1:i+T+1] for i in ix])
        tot += m(xb, yb)[1].item()
    return tot / iters

def ablate_head(layer, head_idx):
    m = copy.deepcopy(model)
    with torch.no_grad():
        m.blocks[layer].attn.heads[head_idx].value.weight.zero_()
    return m

base = val_loss(model)
print(f'{"intact":<34}val {base:.4f}')
for l, h, name in [(0, 3, 'L0H3 (previous-token head)'),
                   (0, 1, 'L0H1 (diffuse, control)')]:
    v = val_loss(ablate_head(l, h))
    print(f'ablate {name:<27}val {v:.4f}   (+{v-base:.4f})')

@torch.no_grad()
def copy_gain(half=32, trials=300, seed=0):
    torch.manual_seed(seed)
    first, second = 0.0, 0.0
    for _ in range(trials):
        base = torch.randint(0, 65, (half,))
        seq = torch.cat([base, base])[None]
        lp = F.cross_entropy(model(seq[:, :-1])[0][0], seq[0, 1:],
                             reduction='none')
        first  += lp[:half-1].mean().item()
        second += lp[half:].mean().item()
    return first/trials, second/trials

f, s = copy_gain()
print(f'loss on first  (novel)  copy : {f:.4f} nats')
print(f'loss on second (repeat) copy : {s:.4f} nats')
print(f'in-context copying gain      : {f - s:+.4f} nats')
print(f'(for scale, ln V = {math.log(65):.4f}; these are >> that because\n'
      f' random uniform text is wildly off this model\'s distribution)')

# Per-position loss across the repeated half: an induction head would make
# this DROP sharply after the first couple of tokens. It stays flat.
@torch.no_grad()
def per_position(half=32, trials=400, seed=0):
    torch.manual_seed(seed)
    acc = torch.zeros(half)
    for _ in range(trials):
        base = torch.randint(0, 65, (half,))
        seq = torch.cat([base, base])[None]
        lp = F.cross_entropy(model(seq[:, :-1])[0][0], seq[0, 1:],
                             reduction='none')
        acc += lp[half-1:2*half-1]
    return acc / trials

pp = per_position()
print('per-position loss in the repeated half:')
print('  ', '  '.join(f'{v:.1f}' for v in pp[:8].tolist()), '...')
print(f'  flat at ~{pp.mean():.1f} nats — no drop, so no copying')

class LoRALinear(nn.Module):
    def __init__(self, base, r):
        super().__init__()
        self.base = base
        for p in self.base.parameters():
            p.requires_grad = False
        self.A = nn.Parameter(torch.randn(r, base.in_features) * 0.01)
        self.B = nn.Parameter(torch.zeros(base.out_features, r))
    def forward(self, x):
        return self.base(x) + (x @ self.A.T) @ self.B.T

step	model	bits/char
1	bigram counts	3.54
3	Bengio MLP, $k=3$	2.86
6	GPT, 816,640 params	2.19
—	Shannon's estimate for English	~1

Step 8 — Capstone: Track C worked example¶

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

Track C.1 — The attention atlas¶

Track C.2 — Ablate and confirm the previous-token head¶

Track C.3 — A rigorous negative: no induction head¶

Tracks A and B — the mathematics, unexecuted¶

Track A: LoRA¶

Track B: KL-regularized preference optimization¶

Your bits/char table¶