Step 1 — A Bigram Model from Counts

Released with Lecture 1 · ~3 hours · Prerequisites: Setup complete.

In what follows, there are blocks of code provided for you. It’s important that you can read the code. If you are not proficient with Python, then I propose the following workflow. For each block of code, keywords are provided for the python syntax which is being used, so you can look up python syntax and functions you don’t recognise. I’ve put a “Python background” foldout tutorial on topics you might need, right below the code blocks. You can also ask the copilot in Colab. Then I suggest that you first experiment with changing the code or printing out intermediate variables, until you have figured out what you think it does. Then point AI toward the page and step number, provide your precise explanation of what you believe each line of the code is doing, and ask for corrections. In this way you will quickly become a code reader (if not a code writer). — Kate

Goal

Build the simplest possible language model, namely a bigram model $p(x_{t+1} \mid x_t)$ estimated by counting character pairs in the corpus. Sample text from it, and measure its quality with average log loss (cross-entropy) in nats per character, in bits per character, and with perplexity. These are three reports of the same underlying score. You will measure it twice: first on the text the model was fit on, then on a held-out 10% the counts never saw.

Concept Review

1. Let $V$ be a finite alphabet and suppose we model a text $x_1,\dots,x_T$ by a Markov chain: $p(x_1,\dots,x_T) = p(x_1)\prod_{t=1}^{T-1} p(x_{t+1}\mid x_t)$, with parameters $\theta_{ab} = p(b \mid a)$. Show that the maximum likelihood estimator for a given text is the count ratio $\hat\theta_{ab} = N_{ab} / \sum_{c} N_{ac}$, where $N_{ab}$ is the number of times character $b$ follows character $a$ in the text.
2. The model’s average log loss on the text is $\mathcal L = -\frac{1}{T-1}\sum_{t=1}^{T-1} \log p(x_{t+1}\mid x_t)$. This is the empirical cross-entropy, measured in nats. Divide by $\ln 2$ for bits per character; perplexity is $e^{\mathcal L}$. Show that perplexity is $|V|$ exactly when the model is uniform (Lecture 1’s Proposition 4.2).

Tasks

1.1 Load the data

If on Colab, download the data file first: !wget -q https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt

with open('input.txt', 'r') as f:   # with statement (context manager)
    text = f.read()
print(len(text))        # 1,115,394 characters
print(text[:250])       # slicing — First Citizen: Before we proceed any further...

# Create a small file to practice on ('w' means open for writing):
with open('demo.txt', 'w') as f:
    f.write("To be, or not to be,\n")
    f.write("that is the question.\n")
# Read it back as one string ('r' means open for reading):
with open('demo.txt', 'r') as f:
    text = f.read()
print(len(text))
print(text)

text = "First Citizen: Before we proceed any further, hear me speak."
print(len(text))
print(text[0])        # indexing starts at 0
print(text[0:5])      # positions 0,1,2,3,4 — the stop index is excluded
print(text[:5])       # same thing: missing start = from the beginning
print(text[15:])      # missing stop = to the end
print(text[-6:])      # negative = count from the end

1.2 Tokenize

Our tokens, for the next five steps, are single characters.

vocab = sorted(set(text))                   # set; sorted
V = len(vocab)                              # 65
stoi = {ch: i for i, ch in enumerate(vocab)}  # dictionary (dict comprehension); enumerate
itos = {i: ch for i, ch in enumerate(vocab)}
encode = lambda s: [stoi[c] for c in s]     # lambda function; list comprehension
decode = lambda ids: ''.join(itos[i] for i in ids)  # str.join; generator expression
assert decode(encode("Fear no more")) == "Fear no more"  # assert statement

text = "to be or not to be"
distinct = set(text)
print(distinct)            # each character once, in no promised order
vocab = sorted(distinct)   # a list, in a fixed, reproducible order
print(vocab)
print(len(vocab))

vocab = ['a', 'b', 'n']
for i, ch in enumerate(vocab):
    print(i, ch)
# Build character-to-integer and integer-to-character tables:
stoi = {ch: i for i, ch in enumerate(vocab)}
itos = {i: ch for i, ch in enumerate(vocab)}
print(stoi)
print(itos)
print(stoi['b'])           # look up a key
print(itos[2])

squares = [n * n for n in [1, 2, 3, 4]]
print(squares)
double = lambda x: 2 * x       # a one-line function...
print(double(7))               # ...called like any other
# Combined, exactly as in the code above:
stoi = {'a': 0, 'b': 1, 'n': 2}
encode = lambda s: [stoi[c] for c in s]
print(encode("banana"))

words = ['To', 'be', 'or', 'not']
print(' '.join(words))     # the separator goes in front
print(''.join(words))
itos = {0: 'a', 1: 'b', 2: 'n'}
ids = [1, 0, 2, 0, 2, 0]
print(''.join(itos[i] for i in ids))   # decode, as in the code above

assert 2 + 2 == 4                # true: nothing happens
print("still running")
assert 2 + 2 == 5, "arithmetic is broken"   # false: stops with an error
print("this line is never reached")

Look at vocab. What are tokens 0 and 1? Which characters made the cut?

1.3 Count

PyTorch (import torch) is the deep-learning library used throughout this course. At its core it is a Python library for computing with tensors — multi-dimensional numerical arrays, generalizing vectors and matrices — with two extras that matter later: the same code can run on a GPU, and it can differentiate through computations automatically (autograd, which you will build a miniature of yourself in Step 3 and use ever after). This week we use none of that: it is purely an array library here, playing the role NumPy plays elsewhere in scientific Python.

Build the $65 \times 65$ count matrix $N$ with $N_{ab}$ = occurrences of character pair $(a, b)$, as a PyTorch integer tensor. Then visualize it:

import torch
import matplotlib.pyplot as plt

N = torch.zeros((V, V), dtype=torch.int64)  # keyword argument (dtype=)
ids = encode(text)
for a, b in zip(ids, ids[1:]):    # zip; tuple unpacking; list slicing
    N[a, b] += 1                  # augmented assignment (+=)

plt.figure(figsize=(10, 10))
plt.imshow(N.log1p(), cmap='Blues')   # log(1+N): raw counts span 4 orders of magnitude

import torch
N = torch.zeros((3, 2), dtype=torch.int64)  # a 3x2 grid of integer zeros
N[0, 1] += 1                     # index by row, column
print(N.shape)                   # torch.Size([3, 2])
t = torch.tensor([1., 2., 3.])   # build a tensor from a Python list
print(t * 10)                    # tensor([10., 20., 30.]) — elementwise
print(t.sum())                   # tensor(6.) — still a (0-dimensional) tensor...
print(t.sum().item())            # 6.0 — .item() unwraps it to a plain Python number

def repeat(word, times=2, sep=' '):     # times and sep have defaults
    return sep.join([word] * times)
print(repeat('fear'))                   # use both defaults
print(repeat('fear', times=4))          # override one default, by name
print(repeat('fear', sep='-', times=3)) # named arguments go in any order

ids = [5, 3, 8, 3, 5]
print(ids[1:])               # the same list shifted left by one
for a, b in zip(ids, ids[1:]):
    print(a, b)              # each element with its successor

count = 0
count += 1
count += 10
print(count)
tallies = {'a': 0, 'b': 0}
tallies['a'] += 1
tallies['a'] += 1
print(tallies)

Stare at the heatmap and find three things (each is checkable from N, and the answers are in the checkpoints):

1. the row for q — what single column dominates, and how completely?
2. the column for space — is every character ever followed by a space?
3. the capital-letter block — is it brighter into lowercase or into itself? Predict before you look, then explain what you actually see.

1.4 Normalize — carefully

Turn counts into the MLE transition matrix $P$ with rows summing to 1. Use add-one (Laplace) smoothing, $P_{ab} \propto N_{ab} + 1$: some pairs never occur, and a single zero probability makes the log-likelihood of any text containing that pair $-\infty$. (Question worth 30 seconds of thought: what is Laplace smoothing the Bayesian posterior mean for?)

P = (N + 1).float()                 # method call on an expression
P = P / P.sum(dim=1, keepdim=True)  # broadcasting (row-wise division)
assert torch.allclose(P.sum(dim=1), torch.ones(V))

s = "fear no more"
print(s.upper())              # a method on a string
print(s.split())              # split on spaces, into a list
print(s.upper().split())      # chained: left to right
print(("no " * 3).strip())    # a method on a parenthesized expression

import numpy as np
M = np.array([[1., 2., 3.],
              [4., 5., 6.]])
sums = M.sum(axis=1, keepdims=True)
print(sums)                 # shape (2, 1): one sum per row, kept as a column
print(M / sums)             # broadcasting: each row divided by its own sum
print(M.sum(axis=1))        # without keepdims: shape (2,) — a flat vector

1.5 Sample

Generate text by running the Markov chain: start from the newline character, repeatedly draw the next character from row $P_{x_t}$ (torch.multinomial), for 500 steps. Decode and read your model’s Shakespeare aloud.

1.6 Evaluate on the training text

Compute your model’s average log loss on the text (average $-\log P_{x_t, x_{t+1}}$; vectorize it — P[ids[:-1], ids[1:]] indexes all the transitions at once) and report:

• log loss (cross-entropy) in nats and in bits per character,
• perplexity $e^{\mathcal L}$,
• the same reports for the uniform model $P_{ab} = 1/65$ as a baseline.

One caveat to keep in mind throughout: the counts came from this same text, so every number in this task is a training value — the model is being graded on the exam it studied from. Task 1.7 is the honest version.

Guidance. You are implementing Lecture 1 §3’s token-level log loss, specialized to a bigram model with $p_\theta(b\mid a)=P_{ab}$ and $T-1$ predicted tokens. Perplexity is defined in Lecture 1 §4; §6 gives the bits-per-character conversion.

You are computing one number per transition: for each of the $T-1$ adjacent pairs $(x_t, x_{t+1})$ in the text, look up the probability your model assigned to the character that actually came next, take $-\log$, and average. The vectorized lookup works because PyTorch accepts a pair of index tensors:

ids_t = torch.tensor(ids)              # list -> tensor
p_next = P[ids_t[:-1], ids_t[1:]]      # paired (fancy) indexing: entry t is P[ids[t], ids[t+1]]

import numpy as np
P = np.array([[0.1, 0.2, 0.7],
              [0.5, 0.5, 0.0],
              [0.3, 0.3, 0.4]])
print(P[1, 2])                # one entry: row 1, column 2
rows = np.array([0, 1, 2])
cols = np.array([2, 0, 1])
print(P[rows, cols])          # entries (0,2), (1,0), (2,1), all at once
ids = np.array([0, 2, 1, 1, 0])
print(P[ids[:-1], ids[1:]])   # one entry per adjacent pair of ids

p_next has length len(text) - 1, and entry $t$ is the model’s probability for character $t{+}1$ given character $t$ — the likelihood of the corpus, unrolled into a vector. From it: the mean of $-\log$ is $\mathcal L$ in nats; divide by $\ln 2$ for bits per character, and apply $e^{\mathcal L}$ to obtain perplexity.

Checks and common stumbles, in the order people hit them:

• Validate the vectorization against a loop. Compute the same average with an explicit for loop over the first 10,000 pairs and compare the two numbers.
• Run the uniform baseline through the same code, rather than only computing $\ln 65$ on paper. The uniform answer is insensitive to the data, so getting exactly 4.1744 is a test of your code.
• If your model scores above the uniform 4.1744 (this is bad), maybe you swapped the inputs in P[ids_t[1:], ids_t[:-1]]? The probability of the previous character given the next gives 4.6647 nats on this corpus, worse than knowing nothing. A model evaluated backwards underperforms ignorance!
• Keep logarithm bases straight. Work in natural logs until the final conversion: perplexity is $e^{\mathcal L_{\text{nats}}}$, equivalently $2^{\text{bits/char}}$, and mixing them ($e^{3.54} \approx 35$) produces a plausible-looking wrong number.
• Expected values are in the checkpoints below.

1.7 Evaluate on held-out text

Task 1.6 graded the model on the text its counts came from (teaching to the test). It’s better to test it on fresh data.

1. Split chronologically: the first 90% of the text is the training set, the last 10% the validation set.

   n_train = int(0.9 * len(ids))            # 1,003,854
   ids_tr = torch.tensor(ids[:n_train])
   ids_va = torch.tensor(ids[n_train:])     # 111,540 characters

   Keep the vocabulary from Task 1.2, built from the full text. (Conveniently, the last 10% contains no characters absent from the first 90%, so nothing else changes.)

2. Refit from training counts only. Rebuild the count matrix N_tr and the add-one matrix P_tr exactly as in Tasks 1.3–1.4, but from ids_tr alone.

3. Evaluate twice through one pipeline. Wrap the Task 1.6 computation in a function avg_log_loss(P, ids) and call it on ids_tr and on ids_va (score the transitions internal to each portion). Report log loss in nats, bits per character, and perplexity for both.

4. Count the transitions training never saw.

   train_count_of_val_pairs = N_tr[ids_va[:-1], ids_va[1:]]
   print((train_count_of_val_pairs == 0).sum())

5. Run the uniform model through the same validation pipeline. It must give exactly 4.1744 nats again: any change from Task 1.6 is a bug.

def avg(numbers):
    total = sum(numbers)
    return total / len(numbers)
print(avg([1.0, 2.0, 6.0]))
print(avg([10, 20]))          # the same code, reused on different input

Guidance. Validation loss is the number that estimates performance on new text (Lecture 1 §6.3); training loss only tells you how well the counts memorized their own corpus. The gap you should see is real but tiny, about 0.03 nats. The bigram matrix has only 4,225 cells sharing a million characters of evidence, so almost every cell is estimated from abundant data. Overfitting is the tendency of a model to do better on the training data than on novel validation data. In general, overfitting is a function of the ratio of parameters to data, and here that ratio is tiny.

Checkpoints

• len(text) = 1,115,394; vocabulary size 65; count matrix sums to len(text) - 1 = 1,115,393. 2,822 of the 4,225 cells of $N$ are zero — two thirds of possible character pairs never occur in the complete works of Shakespeare, which is the concrete reason smoothing is not optional.
• Heatmap answers: (1) q occurs 609 times and is followed by u 609 times so probability exactly 1.000. (2) Only 46 of 65 characters are ever followed by a space. (3) The capital block is brightest into itself — 51% of characters following a capital are capitals against 33% lowercase, backwards from ordinary prose, because of the ALL-CAPS SPEAKER: convention of a printed play (a further 6.9% are :).
• Training-corpus values (Task 1.6 — counts and evaluation both use the full text): average log loss 2.4549 nats = 3.5417 bits/char; perplexity 11.646 (uniform baseline: $\ln 65 = 4.1744$ nats, perplexity exactly 65). Unsmoothed MLE gives 2.4526 — the smoothing costs you 0.002 nats and buys finiteness.
• Held-out values (Task 1.7): the split is 1,003,854 training / 111,540 validation characters, and the validation slice opens mid-Taming of the Shrew (GREMIO: Good morrow, neighbour Baptista). The add-one bigram fit on the training portion scores 2.4546 nats on train, 2.4819 nats on validation (3.5412 vs 3.5806 bits/char; perplexity 11.641 vs 11.964) — a gap of 0.027 nats. 187 validation transitions (23 distinct pairs) have zero training count. Look at which: the top three, S→P (63), E→B (42), and N→Z (37), are PROSPERO, SEBASTIAN, and GONZALO — The Tempest and its ALL-CAPS cast enter the corpus only in the final 10%. The unsmoothed MLE scores 2.4519 on train and infinite on validation; any one of the 187 suffices.
• Samples look like: Thas ppomat he s I fone hell, t athe pen. Nonsense — but English-shaped nonsense: mostly pronounceable, word lengths right.

For calibration: Shannon (1951) estimated English at ~1 bit/character, and your Step 6 transformer will reach roughly 1.5 bits/char on this dataset. Watch your bits/char drop as the course proceeds.

Troubleshooting

• KeyError in encode: you’re encoding a character not in Shakespeare’s vocabulary (e.g. curly quotes from copy-pasting).
• Rows of P don’t sum to 1: check dim= and keepdim= in the sum: did you mix rows and columns?
• Samples are all one repeated character: you probably normalized columns instead of rows.

Going further (optional)

• Trigrams: condition on two characters ($65^2$ contexts). With counts fit and evaluated on the full corpus, average log loss drops to 1.9532 nats (2.818 bits/char, perplexity 7.05) — but that is a training-corpus value, and Lecture 1 §6.4 proved that $n$-gram training loss can only improve as context grows, all the way to memorization. So comparing the bigram’s training optimum with the trigram’s training optimum is not evidence that the trigram generalizes better. The comparison that counts: refit the trigram on the same 90% training portion, with the same character tokenizer, and score the same validation portion as Task 1.7. Expected: 1.9526 nats train, 2.0684 nats validation (2.9841 bits/char, perplexity 7.91) against the bigram’s 2.4819 — so the trigram genuinely does generalize better here. But the warning signs are growing: its train/val gap is 0.116 nats to the bigram’s 0.027, and 1,553 validation transitions are unseen in training, up from 187. Now estimate how the number of parameters grows with context length $k$: the matrix has $65^{k+1}$ entries, so $k=4$ already exceeds a billion while the data stays at a million characters. This exponential blow-up against fixed data is precisely why Lecture 2 replaces matrices with functions.
• Vary smoothing: $P \propto N + \alpha$ for $\alpha \in \{0.01, 1, 100\}$. Explain qualitatively what large $\alpha$ does to samples and to log loss. Then choose $\alpha$ by validation loss on the Task 1.7 split — using validation data to select a hyperparameter is exactly the role Lecture 1 §6.3 assigns it.

Catch-up

None — this is the first step. If you’re joining after Lecture 2+, run solutions/step-01.ipynb and read its prose; it’s self-contained.