import plotly.io as pio; pio.renderers.default = "notebook_connected"

import math
import os
import re
import time
import urllib.request
from collections import Counter

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots
from IPython.display import HTML, display

DEVICE = 'cuda' if torch.cuda.is_available() else ('mps' if torch.backends.mps.is_available() else 'cpu')
SEED = 189
torch.manual_seed(SEED)
print('device:', DEVICE)

device: mps

BLUE, GOLD = '#002676', '#FDB515'
TOKEN_COLORS = ['#cfe2f3', '#fce5cd', '#d9ead3', '#ead1dc', '#fff2cc', '#d0e0e3']


def show_tokens(tokens, title=None):
    """Display a list of token strings with alternating background colors (like the tiktoken app).

    A space is drawn as a middle dot and a newline as a return arrow so that whitespace is visible.
    """
    spans = []
    for i, tok in enumerate(tokens):
        shown = (tok.replace('&', '&amp;').replace('<', '&lt;').replace('>', '&gt;')
                    .replace(' ', '&middot;').replace('\n', '&#8629;'))
        spans.append(f'<span style="background:{TOKEN_COLORS[i % len(TOKEN_COLORS)]};'
                     f'color:black;padding:1px 0px;margin:0px 1px;border-radius:3px">{shown}</span>')
    head = f'<b>{title}</b> ' if title else ''
    display(HTML(f'<div style="font-family:monospace;font-size:15px;line-height:2">{head}'
                 f'({len(tokens)} tokens)<br>{"".join(spans)}</div>'))


def plot_matrix(Z, x_labels, y_labels, title, text=None, colorscale='Blues', zmin=None, zmax=None,
                x_title='', y_title='', width=520, height=460):
    """Heatmap of a matrix with labeled rows and columns (row 0 at the top)."""
    fig = go.Figure(go.Heatmap(z=Z, x=x_labels, y=y_labels, text=text,
                               texttemplate='%{text}' if text is not None else None,
                               colorscale=colorscale, zmin=zmin, zmax=zmax, showscale=False,
                               xgap=1, ygap=1))
    fig.update_yaxes(autorange='reversed', title=y_title, type='category')
    fig.update_xaxes(side='top', title=x_title, type='category')
    fig.update_layout(title=title, width=width, height=height, margin=dict(l=60, r=20, t=110, b=20))
    return fig


def plot_lookup(onehot, table, result, row_labels, vocab_labels):
    """Three heatmaps: one-hot rows (N x K), the embedding table (K x D), and their product (N x D)."""
    fig = make_subplots(rows=1, cols=3, column_widths=[0.5, 0.25, 0.25], horizontal_spacing=0.08,
                        subplot_titles=['one-hot inputs (N x K)', 'embedding table (K x D)',
                                        'product = selected rows (N x D)'])
    dims = [f'd{j}' for j in range(table.shape[1])]
    fig.add_trace(go.Heatmap(z=onehot, x=vocab_labels, y=row_labels, colorscale='Greys',
                             showscale=False, xgap=1, ygap=1), 1, 1)
    fig.add_trace(go.Heatmap(z=table, x=dims, y=vocab_labels, colorscale='RdBu', zmid=0,
                             showscale=False, xgap=1, ygap=1), 1, 2)
    fig.add_trace(go.Heatmap(z=result, x=dims, y=row_labels, colorscale='RdBu', zmid=0,
                             showscale=False, xgap=1, ygap=1), 1, 3)
    for c in (1, 2, 3):
        fig.update_yaxes(autorange='reversed', type='category', row=1, col=c)
        fig.update_xaxes(type='category', row=1, col=c)
    fig.update_layout(width=1000, height=520, margin=dict(l=60, r=20, t=60, b=40))
    return fig


def plot_loss_curve(history, ln_K):
    """Training loss at every step, validation loss at the evaluation steps, and the line ln K."""
    fig = go.Figure()
    fig.add_scatter(x=history['step'], y=history['train'], mode='lines', name='training loss (one batch)',
                    line=dict(color=BLUE, width=1), opacity=0.6)
    fig.add_scatter(x=history['val_step'], y=history['val'], mode='lines+markers', name='validation loss',
                    line=dict(color=GOLD, width=3))
    fig.add_hline(y=ln_K, line_dash='dash', line_color='gray',
                  annotation_text=f'ln K = {ln_K:.2f} (uniform prediction)', annotation_position='top right')
    fig.update_layout(xaxis_title='step', yaxis_title='cross-entropy (nats per token)',
                      width=800, height=420, margin=dict(l=60, r=20, t=40, b=50))
    return fig

from transformers import AutoConfig, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained('Qwen/Qwen3-0.6B')

text = 'Hello, how are you?'
ids = tokenizer.encode(text)
print('token ids:', ids)
print('tokens:   ', [tokenizer.decode([i]) for i in ids])
show_tokens([tokenizer.decode([i]) for i in ids], 'Qwen3')

token ids: [9707, 11, 1246, 525, 498, 30]
tokens:    ['Hello', ',', ' how', ' are', ' you', '?']

config = AutoConfig.from_pretrained('Qwen/Qwen3-0.6B')
print(f'tokenizer.vocab_size = {tokenizer.vocab_size:,}')
print(f'len(tokenizer)       = {len(tokenizer):,}   ({len(tokenizer) - tokenizer.vocab_size} special tokens)')
print(f'embedding rows       = {config.vocab_size:,}   (= 128 x {config.vocab_size // 128})')
print('some special tokens: ', [t.content for t in list(tokenizer.added_tokens_decoder.values())[:6]])

tokenizer.vocab_size = 151,643
len(tokenizer)       = 151,669   (26 special tokens)
embedding rows       = 151,936   (= 128 x 1187)
some special tokens:  ['<|endoftext|>', '<|im_start|>', '<|im_end|>', '<|object_ref_start|>', '<|object_ref_end|>', '<|box_start|>']

for s in ['The transformer is unbelievably tokenizable.',
          'CS 189 meets at 15:30 in 2050 VLSB.',
          'Ich bin ein Berliner. 我喜欢机器学习。',
          'def f(x):\n    return x ** 2  # 🙂']:
    show_tokens([tokenizer.decode([i]) for i in tokenizer.encode(s)])

def split_words(text):
    """Split text into words that keep one leading whitespace character, e.g. 'to be' -> ['to', ' be'].

    Every character belongs to exactly one word, so ''.join(split_words(text)) == text.
    """
    return re.findall(r'\s?\S+|\s', text)


def count_pairs(words):
    """Count adjacent token pairs. `words` maps a tuple of tokens to the number of times it occurs."""
    counts = Counter()
    for tokens, n in words.items():
        for pair in zip(tokens, tokens[1:]):
            counts[pair] += n
    return counts


def merge_pair(tokens, pair):
    """Replace every occurrence of the adjacent pair (a, b) in a tuple of tokens by the token a + b."""
    a, b = pair
    out, i = [], 0
    while i < len(tokens):
        if i + 1 < len(tokens) and tokens[i] == a and tokens[i + 1] == b:
            out.append(a + b)
            i += 2
        else:
            out.append(tokens[i])
            i += 1
    return tuple(out)


def learn_bpe(text, num_merges, max_words=None):
    """Learn `num_merges` BPE merges from text. Returns the list of merges and their counts.

    With max_words, only the most frequent distinct words are used (faster on a large corpus).
    """
    word_counts = Counter(split_words(text))
    if max_words is not None:
        word_counts = dict(word_counts.most_common(max_words))
    words = {tuple(w): n for w, n in word_counts.items()}      # each word starts as its characters
    merges, counts = [], []
    for _ in range(num_merges):
        pair_counts = count_pairs(words)
        if not pair_counts:
            break
        best = max(pair_counts, key=pair_counts.get)             # most frequent pair (first one on ties)
        merges.append(best)
        counts.append(pair_counts[best])
        words = {merge_pair(w, best): n for w, n in words.items()}
    return merges, counts

peter = 'Peter Piper picked a peck of pickled peppers'
print('words:', split_words(peter))
merges, counts = learn_bpe(peter, 8)
for (a, b), n in zip(merges, counts):
    print(f'merge {a!r} + {b!r} -> {a + b!r}  ({n} times)')

words: ['Peter', ' Piper', ' picked', ' a', ' peck', ' of', ' pickled', ' peppers']
merge 'p' + 'e' -> 'pe'  (4 times)
merge 'c' + 'k' -> 'ck'  (3 times)
merge 'pe' + 'r' -> 'per'  (2 times)
merge ' ' + 'p' -> ' p'  (2 times)
merge ' p' + 'i' -> ' pi'  (2 times)
merge ' pi' + 'ck' -> ' pick'  (2 times)
merge 'e' + 'd' -> 'ed'  (2 times)
merge ' ' + 'pe' -> ' pe'  (2 times)

words = [tuple(w) for w in split_words(peter)]
show_tokens([t for w in words for t in w], 'characters:')
for step, ((a, b), n) in enumerate(zip(merges, counts), start=1):
    words = [merge_pair(w, (a, b)) for w in words]
    show_tokens([t for w in words for t in w], f'merge {step}: {a!r} + {b!r} -> {a + b!r} ({n} times)')

class BPETokenizer:
    """A character-level BPE tokenizer with Bishop's whitespace rule."""

    def __init__(self, text, num_merges, max_words=None, alphabet=None):
        """Learn merges from text. The base tokens are the characters of `alphabet` (default: of text)."""
        self.merges, self.counts = learn_bpe(text, num_merges, max_words)
        self.rank = {pair: r for r, pair in enumerate(self.merges)}      # earlier merge = lower rank
        chars = sorted(set(text if alphabet is None else alphabet))
        self.vocab = chars + [a + b for a, b in self.merges]               # id -> token string
        self.token_id = {tok: i for i, tok in enumerate(self.vocab)}
        self._cache = {}

    def __len__(self):
        return len(self.vocab)

    def encode_word(self, word):
        tokens = tuple(word)
        while len(tokens) > 1:
            pairs = list(zip(tokens, tokens[1:]))
            best = min(pairs, key=lambda p: self.rank.get(p, math.inf))
            if best not in self.rank:
                break                                   # no learned merge applies
            tokens = merge_pair(tokens, best)
        return [self.token_id[t] for t in tokens]

    def encode(self, text):
        ids = []
        for w in split_words(text):
            if w not in self._cache:
                self._cache[w] = self.encode_word(w)
            ids.extend(self._cache[w])
        return ids

    def decode(self, ids):
        return ''.join(self.vocab[i] for i in ids)

    def tokens(self, text):
        return [self.vocab[i] for i in self.encode(text)]


toy = BPETokenizer(peter, 8)
ids = toy.encode('a pickled pepper')
print('K =', len(toy), 'tokens:', toy.vocab)
print('ids:', ids, ' tokens:', toy.tokens('a pickled pepper'))
print('round trip:', toy.decode(ids) == 'a pickled pepper')

K = 23 tokens: [' ', 'P', 'a', 'c', 'd', 'e', 'f', 'i', 'k', 'l', 'o', 'p', 'r', 's', 't', 'pe', 'ck', 'per', ' p', ' pi', ' pick', 'ed', ' pe']
ids: [2, 20, 9, 21, 22, 11, 17]  tokens: ['a', ' pick', 'l', 'ed', ' pe', 'p', 'per']
round trip: True

SHAKESPEARE_URL = 'https://raw.githubusercontent.com/BerkeleyML/fa26-student/main/lecture/lec18/shakespeare.txt'


def load_shakespeare(path='shakespeare.txt'):
    if not os.path.exists(path):   # e.g., on Colab
        urllib.request.urlretrieve(SHAKESPEARE_URL, path)
    raw = open(path, encoding='utf-8').read()
    start = raw.index('\n', raw.index('*** START OF'))
    end = raw.index('*** END OF')
    lines = [line.strip() for line in raw[start:end].splitlines()]
    return re.sub(r'\n{3,}', '\n\n', '\n'.join(lines)).strip()


text = load_shakespeare()
split = int(0.9 * len(text))
train_text, val_text = text[:split], text[split:]
print(f'{len(text):,} characters, {len(set(text))} distinct')

t0 = time.perf_counter()
bpe = BPETokenizer(train_text, num_merges=400, max_words=10_000, alphabet=text)  # all characters are tokens
print(f'learned {len(bpe.merges)} merges in {time.perf_counter() - t0:.1f} s, K = {len(bpe)}')
print('first 20 merges:', [a + b for a, b in bpe.merges[:20]])
print('last 20 merges: ', [a + b for a, b in bpe.merges[-20:]])

5,341,108 characters, 100 distinct

learned 400 merges in 5.7 s, K = 500
first 20 merges: [' t', 'he', ' a', 'ou', ' s', ' m', ' w', 'nd', ' the', 'ha', 'in', 're', ' b', 'is', ' o', 'll', ' f', 'it', 'or', 've']
last 20 merges:  [' speak', 'ound', ' yet', 'ue', 'au', ' r', 'art', 'el', '\nWh', 'fe', 'ene', ' ey', 'ort', 'y.', '\nWith', ' these', ' who', 'ce.', 'eunt', ' some']

passage = text[text.index('To be, or not to be'):][:300]
show_tokens(bpe.tokens(passage), f'our BPE (K = {len(bpe)})')
show_tokens([tokenizer.decode([i]) for i in tokenizer.encode(passage)], f'Qwen3 (K = {len(tokenizer):,})')

t0 = time.perf_counter()
train_ids = torch.tensor(bpe.encode(train_text))
val_ids = torch.tensor(bpe.encode(val_text))
print(f'encoded in {time.perf_counter() - t0:.1f} s: {len(train_ids):,} training and {len(val_ids):,} validation tokens')
print(f'characters per token: ours {len(text) / (len(train_ids) + len(val_ids)):.2f}, '
      f'Qwen3 on the passage {len(passage) / len(tokenizer.encode(passage)):.2f}')
assert bpe.decode(bpe.encode(passage)) == passage

encoded in 1.0 s: 2,200,314 training and 244,087 validation tokens
characters per token: ours 2.19, Qwen3 on the passage 3.90

torch.manual_seed(SEED)
D_toy = 6
embedding = nn.Embedding(len(toy), D_toy)          # weight: (K, D), row k is the embedding of token k
ids = torch.tensor(toy.encode('a peck of pickled peppers'))

X = F.one_hot(ids, num_classes=len(toy)).float()   # (N, K) one-hot rows
lookup = embedding(ids)                            # (N, D) by indexing
product = X @ embedding.weight                     # (N, D) by matrix product
print('ids:', ids.tolist())
print('lookup == one-hot product:', torch.equal(lookup, product))
print('lookup == indexing the table:', torch.equal(lookup, embedding.weight[ids]))

labels = [f'{n}: {toy.vocab[i]!r}' for n, i in enumerate(ids.tolist())]
plot_lookup(X.numpy(), embedding.weight.detach().numpy(), product.detach().numpy(),
            labels, [repr(t) for t in toy.vocab])

ids: [2, 22, 16, 0, 10, 6, 20, 9, 21, 22, 11, 17, 13]
lookup == one-hot product: True
lookup == indexing the table: True

def causal_mask(N, device=None):
    """N x N mask M with M[n, i] = 0 for i <= n and -inf for i > n (strictly upper triangular)."""
    return torch.full((N, N), float('-inf'), device=device).triu(diagonal=1)


def masked_attention(Q, K, V):
    """Softmax_row[(Q K^T + M) / sqrt(D_k)] V for Q, K of shape (..., N, D_k) and V of shape (..., N, D_v)."""
    N, D_k = Q.shape[-2], Q.shape[-1]
    scores = (Q @ K.transpose(-2, -1) + causal_mask(N, Q.device)) / math.sqrt(D_k)
    return torch.softmax(scores, dim=-1) @ V


inputs = ['<start>', 'I', 'swam', 'across', 'the']
outputs = ['I', 'swam', 'across', 'the', 'river']      # the token each position learns to predict
N = len(inputs)
M_mask = causal_mask(N)
print(M_mask)
plot_matrix((M_mask == 0).float().numpy(), inputs, inputs, 'Mask M (row n: query, column i: key)',
            text=[['0' if v == 0 else '-inf' for v in row] for row in M_mask.tolist()],
            colorscale=[[0, '#f4a6a6'], [1, 'white']], x_title='key i', y_title='query n')

tensor([[0., -inf, -inf, -inf, -inf],
        [0., 0., -inf, -inf, -inf],
        [0., 0., 0., -inf, -inf],
        [0., 0., 0., 0., -inf],
        [0., 0., 0., 0., 0.]])

torch.manual_seed(SEED)
D_k = 4
Q, K, V = torch.randn(N, D_k), torch.randn(N, D_k), torch.eye(N)   # V = I, so the output is the weight matrix
alpha = masked_attention(Q, K, V)                                   # alpha @ I = alpha

print('row sums:', alpha.sum(dim=-1))
print('largest weight on a future token:', alpha.triu(diagonal=1).abs().max().item())

unmasked = torch.softmax(Q @ K.T / math.sqrt(D_k), dim=-1)
renormalized = unmasked.tril() / unmasked.tril().sum(dim=-1, keepdim=True)
print('same as zeroing the future and renormalizing:', torch.allclose(alpha, renormalized))
print('same as PyTorch is_causal=True:', torch.allclose(
    alpha, F.scaled_dot_product_attention(Q, K, V, is_causal=True), atol=1e-6))

plot_matrix(alpha.numpy(), inputs, [f'{a} -> {b}' for a, b in zip(inputs, outputs)],
            'Attention weights (each row sums to 1)', text=[[f'{v:.2f}' for v in row] for row in alpha.tolist()],
            zmin=0, zmax=1, x_title='key (input token)', y_title='query (input -> target)', width=620)

row sums: tensor([1.0000, 1.0000, 1.0000, 1.0000, 1.0000])
largest weight on a future token: 0.0
same as zeroing the future and renormalizing: True
same as PyTorch is_causal=True: True

class SwiGLU(nn.Module):
    """Gated MLP: (Swish(x W1) * (x V)) W2, applied to each token vector."""

    def __init__(self, D, d_ff):
        super().__init__()
        self.W1 = nn.Linear(D, d_ff, bias=False)
        self.V = nn.Linear(D, d_ff, bias=False)
        self.W2 = nn.Linear(d_ff, D, bias=False)

    def forward(self, x):
        return self.W2(F.silu(self.W1(x)) * self.V(x))     # F.silu(a) = a * sigmoid(a) = Swish(a)


def relu_mlp(D, d_ff):
    """The original transformer MLP, ReLU(x W1) W2, without biases."""
    return nn.Sequential(nn.Linear(D, d_ff, bias=False), nn.ReLU(), nn.Linear(d_ff, D, bias=False))


def num_params(module):
    return sum(p.numel() for p in module.parameters())


D = 192
print(f'ReLU MLP,  d_ff = 4D   = {4 * D}: {num_params(relu_mlp(D, 4 * D)):,} parameters (8 D^2 = {8 * D * D:,})')
print(f'SwiGLU,    d_ff = 8D/3 = {8 * D // 3}: {num_params(SwiGLU(D, 8 * D // 3)):,} parameters')
print('output shape:', SwiGLU(D, 8 * D // 3)(torch.randn(2, 10, D)).shape)

a = torch.linspace(-4, 4, 201)
fig = go.Figure()
fig.add_scatter(x=a, y=F.relu(a), name='ReLU(a)', line=dict(color='gray', dash='dash'))
fig.add_scatter(x=a, y=F.silu(a), name='Swish(a) = a sigmoid(a)', line=dict(color=BLUE, width=3))
fig.update_layout(xaxis_title='a', width=600, height=350, margin=dict(l=40, r=20, t=30, b=40))

ReLU MLP,  d_ff = 4D   = 768: 294,912 parameters (8 D^2 = 294,912)
SwiGLU,    d_ff = 8D/3 = 512: 294,912 parameters
output shape: torch.Size([2, 10, 192])

class MoE(nn.Module):
    """Top-k mixture of SwiGLU experts. Each token is sent to the k experts with the largest router scores."""

    def __init__(self, D, d_ff, n_experts=8, k=2):
        super().__init__()
        self.k = k
        self.router = nn.Linear(D, n_experts, bias=False)
        self.register_buffer('score_bias', torch.zeros(n_experts))   # 0 here; used below to unbalance the router
        self.experts = nn.ModuleList([SwiGLU(D, d_ff) for _ in range(n_experts)])

    def forward(self, x):
        shape = x.shape
        x = x.reshape(-1, shape[-1])                         # one row per token
        scores = self.router(x) + self.score_bias            # s = x W_r, (tokens, n_experts)
        top_scores, top_experts = scores.topk(self.k, dim=-1)
        gates = torch.softmax(top_scores, dim=-1)            # g_e for the k chosen experts
        y = torch.zeros_like(x)
        for e, expert in enumerate(self.experts):
            token, slot = (top_experts == e).nonzero(as_tuple=True)   # tokens routed to expert e
            if len(token) > 0:
                y.index_add_(0, token, gates[token, slot, None] * expert(x[token]))
        self.scores, self.top_experts = scores, top_experts  # kept for the routing statistics below
        return y.reshape(shape)


torch.manual_seed(SEED)
D, d_ff = 192, 512
moe = MoE(D, d_ff, n_experts=8, k=2)
expert_params = num_params(moe.experts[0])
total, active = num_params(moe), num_params(moe.router) + moe.k * expert_params
print(f'one expert: {expert_params:,}   total: {total:,}   active per token: {active:,} ({active / total:.0%})')

x = torch.randn(4096, D)
print('output shape:', moe(x).shape)

one expert: 294,912   total: 2,360,832   active per token: 591,360 (25%)

output shape: torch.Size([4096, 192])

def routing_stats(moe, x):
    """Fraction of routing slots f_e per expert and the balance loss N_e * sum_e f_e P_e."""
    with torch.no_grad():
        moe(x)
    n_experts = moe.scores.shape[-1]
    f = torch.bincount(moe.top_experts.flatten(), minlength=n_experts).float() / moe.top_experts.numel()
    P = torch.softmax(moe.scores, dim=-1).mean(dim=0)
    return f, n_experts * (f * P).sum().item()


f_balanced, aux_balanced = routing_stats(moe, x)
moe.score_bias[0] = 1.5                       # every token now prefers expert 1
f_skewed, aux_skewed = routing_stats(moe, x)
moe.score_bias[0] = 0.0

experts = [f'E{e + 1}' for e in range(8)]
fig = go.Figure()
fig.add_bar(x=experts, y=f_balanced, name=f'random router (aux. loss {aux_balanced:.2f})', marker_color=BLUE)
fig.add_bar(x=experts, y=f_skewed, name=f'router biased to E1 (aux. loss {aux_skewed:.2f})', marker_color=GOLD)
fig.add_hline(y=1 / 8, line_dash='dash', line_color='gray', annotation_text='uniform 1/8')
fig.update_layout(barmode='group', yaxis_title='share of routing slots', width=750, height=380,
                  margin=dict(l=50, r=20, t=30, b=40))

N_CTX = 128            # context length N
BATCH_SIZE = 32        # sequences per batch B


def get_batch(data, gen, batch_size=BATCH_SIZE):
    """Random windows of N_CTX + 1 tokens: inputs x (B, N) and targets t (B, N), t shifted by one."""
    starts = torch.randint(len(data) - N_CTX - 1, (batch_size,), generator=gen)
    x = torch.stack([data[s:s + N_CTX] for s in starts])
    t = torch.stack([data[s + 1:s + N_CTX + 1] for s in starts])
    return x.to(DEVICE), t.to(DEVICE)


x, t = get_batch(train_ids, torch.Generator().manual_seed(SEED))
print('x:', tuple(x.shape), ' t:', tuple(t.shape))
for n in range(8):
    print(f'{bpe.decode(x[0, :n + 1].tolist())!r:>45}  ->  {bpe.vocab[t[0, n]]!r}')

x: (32, 128)  t: (32, 128)
                                         ' w'  ->  'on'
                                       ' won'  ->  't'
                                      ' wont'  ->  ','
                                     ' wont,'  ->  ' or'
                                  ' wont, or'  ->  ' is'
                               ' wont, or is'  ->  '’t'
                             ' wont, or is’t'  ->  ' not'
                         ' wont, or is’t not'  ->  '\nT'

class CausalSelfAttention(nn.Module):
    """Multi-head masked self-attention with h heads of size D_k = D / h."""

    def __init__(self, D, n_heads):
        super().__init__()
        self.h = n_heads
        self.W_qkv = nn.Linear(D, 3 * D, bias=False)    # [W^(q) W^(k) W^(v)] for all heads
        self.W_o = nn.Linear(D, D, bias=False)

    def forward(self, X):
        B, N, D = X.shape
        Q, K, V = self.W_qkv(X).view(B, N, 3, self.h, D // self.h).permute(2, 0, 3, 1, 4)  # each (B, h, N, D_k)
        heads = masked_attention(Q, K, V)                                                # (B, h, N, D_k)
        return self.W_o(heads.transpose(1, 2).reshape(B, N, D))                          # concatenate heads


class DecoderBlock(nn.Module):
    """Pre-norm transformer layer: X + Attn(RMSNorm(X)), then H + SwiGLU(RMSNorm(H))."""

    def __init__(self, D, n_heads, d_ff):
        super().__init__()
        self.norm1, self.attn = nn.RMSNorm(D), CausalSelfAttention(D, n_heads)
        self.norm2, self.mlp = nn.RMSNorm(D), SwiGLU(D, d_ff)

    def forward(self, X):
        H = X + self.attn(self.norm1(X))
        return H + self.mlp(self.norm2(H))


class DecoderLM(nn.Module):
    """Decoder-only transformer language model with learned positions and tied input/output embeddings."""

    def __init__(self, vocab_size, D=192, n_layers=4, n_heads=6, max_len=N_CTX):
        super().__init__()
        self.tok_emb = nn.Embedding(vocab_size, D)            # E^T: (K, D)
        self.pos_emb = nn.Embedding(max_len, D)               # one learned vector per position
        self.blocks = nn.Sequential(*[DecoderBlock(D, n_heads, 8 * D // 3) for _ in range(n_layers)])
        self.norm = nn.RMSNorm(D)
        for p in self.parameters():                           # small random initial weights
            if p.dim() == 2:
                nn.init.normal_(p, std=0.02)

    def forward(self, ids):
        """ids: (B, N) token ids. Returns logits (B, N, K); row n scores the token at position n + 1."""
        positions = torch.arange(ids.shape[1], device=ids.device)
        X = self.tok_emb(ids) + self.pos_emb(positions)
        X_tilde = self.norm(self.blocks(X))
        return X_tilde @ self.tok_emb.weight.T                # weight tying: W^(p) = E


torch.manual_seed(SEED)
model = DecoderLM(len(bpe)).to(DEVICE)
emb = model.tok_emb.weight.numel()
print(f'{num_params(model):,} parameters ({emb:,} in the embedding table, shared with the output layer)')
print('logits:', tuple(model(x).shape), ' = (B, N, K)')

1,891,776 parameters (96,000 in the embedding table, shared with the output layer)

logits: (32, 128, 500)  = (B, N, K)

with torch.no_grad():
    x_changed = x.clone()
    x_changed[:, -1] = (x_changed[:, -1] + 1) % len(bpe)
    diff = (model(x) - model(x_changed)).abs().amax(dim=(0, 2))
print(f'largest change at positions 1..N-1: {diff[:-1].max().item():.1e},  at position N: {diff[-1].item():.2f}')

largest change at positions 1..N-1: 0.0e+00,  at position N: 1.16

def lm_loss(model, x, t):
    """Mean cross-entropy over all B x N positions (F.cross_entropy applies the softmax)."""
    logits = model(x)
    return F.cross_entropy(logits.reshape(-1, logits.shape[-1]), t.reshape(-1))


with torch.no_grad():
    loss0 = lm_loss(model, x, t).item()
print(f'initial loss {loss0:.3f},  ln K = ln {len(bpe)} = {math.log(len(bpe)):.3f}')

initial loss 6.239,  ln K = ln 500 = 6.215

@torch.no_grad()
def generate(model, prompt, n_new, seed=0):
    """Sample n_new tokens from the model (temperature 1), starting from a text prompt."""
    was_training = model.training
    model.eval()
    gen = torch.Generator().manual_seed(seed)
    ids = torch.tensor([bpe.encode(prompt)], device=DEVICE)
    for _ in range(n_new):
        logits = model(ids[:, -N_CTX:])[0, -1]                 # logits for the next token
        p = torch.softmax(logits.float().cpu(), dim=-1)
        next_id = torch.multinomial(p, 1, generator=gen)
        ids = torch.cat([ids, next_id.to(DEVICE)[None]], dim=1)
    model.train(was_training)
    return bpe.decode(ids[0].tolist())


PROMPT = 'ROMEO.\n'
print(generate(model, PROMPT, 60))

ROMEO.
 Doushere if comero beé’dightK
S
P
O haveiENrIR tim wellouldant
Youes kn
That doS bfêSu mustce,riR
AndordA.ame wouldetC theirUCELiltherSndordif oneEx you,ousy.w

TRAIN_STEPS, ETA, WARMUP = 1500, 2e-3, 100


def learning_rate(step):
    warmup = min(1.0, (step + 1) / WARMUP)
    return ETA * warmup * 0.5 * (1 + math.cos(math.pi * step / TRAIN_STEPS))


@torch.no_grad()
def validation_loss(model, batches):
    model.eval()
    loss = np.mean([lm_loss(model, xv, tv).item() for xv, tv in batches])
    model.train()
    return loss


val_gen = torch.Generator().manual_seed(0)
val_batches = [get_batch(val_ids, val_gen) for _ in range(10)]
train_gen = torch.Generator().manual_seed(SEED)
decay = [p for p in model.parameters() if p.dim() >= 2]      # weight matrices and the embedding table
no_decay = [p for p in model.parameters() if p.dim() < 2]    # RMSNorm gains
optimizer = torch.optim.AdamW([{'params': decay, 'weight_decay': 0.1},
                               {'params': no_decay, 'weight_decay': 0.0}], lr=ETA)
history = {'step': [], 'train': [], 'val_step': [], 'val': []}

model.train()
t0 = time.perf_counter()
for step in range(TRAIN_STEPS):
    if step % 100 == 0:
        history['val_step'].append(step)
        history['val'].append(validation_loss(model, val_batches))
    for group in optimizer.param_groups:
        group['lr'] = learning_rate(step)
    xb, tb = get_batch(train_ids, train_gen)
    loss = lm_loss(model, xb, tb)
    optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    optimizer.step()
    history['step'].append(step)
    history['train'].append(loss.item())
    if step % 250 == 0:
        print(f'step {step:5d}  training loss {loss.item():.3f}  ({time.perf_counter() - t0:.0f} s)')
history['val_step'].append(TRAIN_STEPS)
history['val'].append(validation_loss(model, val_batches))
train_gen_eval = torch.Generator().manual_seed(1)
train_loss = validation_loss(model, [get_batch(train_ids, train_gen_eval) for _ in range(10)])
print(f'done in {time.perf_counter() - t0:.0f} s; loss on 10 batches of training text {train_loss:.3f}, '
      f'of validation text {history["val"][-1]:.3f}')
plot_loss_curve(history, math.log(len(bpe)))

step     0  training loss 6.239  (1 s)

step   250  training loss 3.811  (24 s)

step   500  training loss 3.388  (45 s)

step   750  training loss 2.986  (66 s)

step  1000  training loss 2.943  (87 s)

step  1250  training loss 2.789  (107 s)

done in 128 s; loss on 10 batches of training text 2.778, of validation text 3.251

for seed in range(3):
    print(generate(model, PROMPT, 80, seed=seed))
    print('-' * 60)

ROMEO.
’Twas come not better than you know.

CROMWELL.
What’s that?

ALEXANT.
And must I gave thee forth; when I say.

SEBASTIAN.
That’s the matter?

ARCHBISHOP.
What is he? Br
------------------------------------------------------------

ROMEO.
’Gutor, the Emperor, lords, i’ three with his gifts;
Thou hast know what my bodies all most seigned shame,
I cannot hope to whom I have heard my hand.

[_Exit._]

TALBOT.
[_s
------------------------------------------------------------

ROMEO.
Queen, we are maid’s a knaves, yet my side
wrong many in. After playings, wish! Fie for glerit in thy tears I’ll reach
despury!

CLEOPATRA.
Alack, her worth sword
------------------------------------------------------------

Notebook section	Slide
(Lecture Demo 1) Tokens: the Qwen3 Tokenizer	Tiktoken App (Tokenizer Demo)
Byte Pair Encoding	Byte Pair Encoding
Language Modeling and Token Embeddings	Embedding Table Lookup; Predicting the Next Token (Transformer)
The Causal Mask	The Masked Attention; The Causal Mask
Gated MLPs (SwiGLU)	Gated MLPs (SwiGLU)
Mixture of Experts	Mixture of Experts (MoE); Total vs. Active Parameters; Keeping the Experts Balanced
(Lecture Demo 2) A Decoder-Only Transformer on Shakespeare	Decoder Transformer Notebook

Lecture 18: Transformers and LLMs (CS 189, Fall 2026)

Plotting Code

(Lecture Demo 1) Tokens: the Qwen3 Tokenizer

Byte Pair Encoding

BPE on Shakespeare

Language Modeling and Token Embeddings

The Causal Mask

Gated MLPs (SwiGLU)

Mixture of Experts

(Lecture Demo 2) A Decoder-Only Transformer on Shakespeare

Training data: shifted targets

The model

The loss at initialization

Training