Chapter 12 — Predicting the next word by counting (n-grams)

This is the most important chapter in Part 3. An LLM does exactly what you're about to build, with a neural network instead of a counting table. If you understand this chapter, Chapter 31 will feel obvious.

What a language model is

A language model assigns probabilities to text:

p("the cat sat on the mat")  =  high
p("the cat sat on the the")  =  low
p("colorless green ideas sleep furiously")  =  low (grammatical but weird)

Using the chain rule of probability, that whole-sentence probability breaks into a series of next-word predictions:

p(w1, w2, w3, ..., wT) = p(w1) x p(w2 | w1) x p(w3 | w1,w2) x ... x p(wT | w1...wT-1)
                                  \_______________________________________________/
                                        "given everything before, what's next?"

That conditional — "given the words so far, what comes next?" — is what every language model computes, from a 1990s n-gram to GPT. Everything else is implementation.

The n-gram shortcut

Conditioning on all previous words is impossible (you'd need counts for every sentence ever). So assume only the last few words matter:

Unigram  (n=1):  p(next)                            ignores context entirely
Bigram   (n=2):  p(next | previous word)
Trigram  (n=3):  p(next | previous TWO words)       <- the classic workhorse

Then just count:

p("mat" | "on", "the")  =  count("on the mat") / count("on the")

the cat sat on the mat
the cat sat on the rug
the dog sat on the mat

count("on the")      = 3        (appears in all three sentences)
count("on the mat")  = 2
count("on the rug")  = 1

p("mat" | "on the") = 2/3 = 0.67
p("rug" | "on the") = 1/3 = 0.33
p("cat" | "on the") = 0/3 = 0.00     <-- PROBLEM

count("the cat")     = 2
count("the cat sat") = 2
p("sat" | "the cat") = 2/2 = 1.00

The two fatal problems (and why they matter for LLMs)

Problem 1 — Zero counts. p("cat" | "on the") = 0. Any sentence containing an unseen trigram gets probability exactly zero, which is clearly wrong — you just never saw it. Fixes are called smoothing:

• Add-one (Laplace): pretend every possible trigram was seen once extra.
• Backoff: if the trigram is unseen, fall back to the bigram, then unigram.
• Interpolation: always blend all three: p = 0.6*trigram + 0.3*bigram + 0.1*unigram.
• Kneser–Ney: the sophisticated version; state of the art for n-grams.

Problem 2 — No sense of similarity. This is the killer.

Corpus contains:  "the cat sat on the mat"     many times
Corpus never has: "the dog sat on the rug"

n-gram model:  p("rug" | "on the") is low, and it has NO IDEA that
               "dog" is like "cat" or "rug" is like "mat".

Every word is an atomic symbol with no relationship to any other word.

To an n-gram model, "cat" and "dog" are as unrelated as "cat" and "democracy". Learning about one teaches it nothing about the other. It has no way to generalise.

This is the exact limitation that word2vec (Part 5) and neural language models remove. Represent words as vectors where similar words are nearby, and what the model learns about "cat" partly transfers to "dog".

Measuring a language model: perplexity

Perplexity = "on average, how many words is the model choosing between?" Lower is better.

perplexity = exp( average of -log(probability assigned to each true next word) )

(Perplexity is only comparable between models using the same vocabulary and test data — a common source of misleading comparisons.)

import random, math
from collections import defaultdict, Counter

def train_trigram(tokens):
    """Count trigrams and the contexts they follow."""
    tri = Counter()
    ctx = Counter()
    uni = Counter()
    toks = ["<s>", "<s>"] + tokens + ["</s>"]
    for i in range(2, len(toks)):
        w1, w2, w3 = toks[i-2], toks[i-1], toks[i]
        tri[(w1, w2, w3)] += 1
        ctx[(w1, w2)] += 1
        uni[w3] += 1
    return tri, ctx, uni

def make_prob_fn(tri, ctx, uni):
    N = sum(uni.values())
    V = len(uni)
    def prob(w1, w2, w3):
        # interpolated: trigram + unigram, with a floor so nothing is zero
        p_tri = tri[(w1,w2,w3)] / ctx[(w1,w2)] if ctx[(w1,w2)] else 0.0
        p_uni = uni[w3] / N
        return max(0.7 * p_tri + 0.3 * p_uni, 1e-9)
    return prob

def generate(tri, ctx, uni, n_words=40):
    """Walk the chain, sampling the next word from the counts."""
    out = ["<s>", "<s>"]
    for _ in range(n_words):
        w1, w2 = out[-2], out[-1]
        candidates = [(w3, c) for (a,b,w3), c in tri.items() if (a,b) == (w1,w2)]
        if not candidates:                    # unseen context -> restart context
            candidates = [(w, c) for w, c in uni.items()]
        words, counts = zip(*candidates)
        nxt = random.choices(words, weights=counts, k=1)[0]
        if nxt == "</s>":
            break
        out.append(nxt)
    return " ".join(out[2:])

def perplexity(prob, tokens):
    toks = ["<s>", "<s>"] + tokens + ["</s>"]
    total, n = 0.0, 0
    for i in range(2, len(toks)):
        total += math.log(prob(toks[i-2], toks[i-1], toks[i]))
        n += 1
    return math.exp(-total / n)

# ---- get a corpus: any plain-text book works ----
# Download one from Project Gutenberg, e.g.:
#   curl -o book.txt https://www.gutenberg.org/files/11/11-0.txt   (Alice in Wonderland)
text = open("book.txt", encoding="utf-8", errors="ignore").read().lower()
words = text.replace("\n", " ").split()

split = int(0.9 * len(words))
train_words, test_words = words[:split], words[split:]

tri, ctx, uni = train_trigram(train_words)
prob = make_prob_fn(tri, ctx, uni)

print("vocabulary size:", len(uni))
print("perplexity on held-out text:", round(perplexity(prob, test_words[:2000]), 1))
print("\n--- generated text ---")
for _ in range(3):
    print(" *", generate(tri, ctx, uni))