Bigram Language Model

The problem

A bigram model predicts the next token from the previous one only. It is a table P with $V$ rows and $V$ columns, where $V$ is the vocabulary size: row $i$ is the distribution of the token that comes after token $i$, so P[i, j] is the probability that $j$ follows $i$. Training it is just counting. You will build the table, measure how well it predicts a text, and turn that into a perplexity.

The text is one long array of token ids, tokens. Its consecutive pairs (previous, next) are what the model learns from.

Part 1: the table

Count every pair, add alpha to every cell (smoothing, so that no probability is 0), and divide each row by its total:

where:

• $\mathrm{count}(i \to j)$ is how many times token $j$ comes right after token $i$;
• $\alpha$ (alpha) is alpha, the pretend extra count of every pair;
• $V$ is vocab_size. Adding $\alpha$ to all $V$ cells of a row adds $\alpha V$ to its total, so the row still adds up to 1.

Example. tokens = [0, 1, 0, 1, 0, 2], vocab_size = 3, alpha = 1. The pairs are 0→1, 1→0, 0→1, 1→0, 0→2.

row	counts	plus alpha	divided by the total
0	$(0, 2, 1)$	$(1, 3, 2)$	$(1/6, 1/2, 1/3)$
1	$(2, 0, 0)$	$(3, 1, 1)$	$(0.6, 0.2, 0.2)$
2	$(0, 0, 0)$	$(1, 1, 1)$	$(1/3, 1/3, 1/3)$

Part 2: how well the table predicts a text

The mean negative log-likelihood of a text is the average of $-\log$ P[i, j] over its consecutive pairs $(i, j)$: a text of $T$ tokens has $T - 1$ pairs. The smaller it is, the better the table predicts the text.

Example. The same tokens under the table above. The five pairs get probabilities P[0, 1] = 0.5, P[1, 0] = 0.6, P[0, 1] = 0.5, P[1, 0] = 0.6, P[0, 2] = 0.33. Their $-\log$ are $0.69, 0.51, 0.69, 0.51, 1.10$, and the mean is about $0.70$.

Part 3: perplexity

The perplexity is $e^{\text{nll}}$, where nll is the mean negative log-likelihood. It reads as the number of tokens the model is hesitating between, on average.

Example. $e^{0.70} \approx 2.0$: on this text the model hesitates between about 2 of the 3 tokens. A uniform table, with every probability $1/V$, has a mean negative log-likelihood of $\log V$ and so a perplexity of exactly $V$: a good sanity check.

What to implement

• bigram_probs(tokens, vocab_size, alpha=1.0): tokens is a 1-D integer array. Return the $(V, V)$ table of part 1.
• bigram_nll(P, tokens): the mean negative log-likelihood of part 2, a number.
• perplexity(nll): $e^{\text{nll}}$, a Python float.

Requirements

• No Python loops over tokens: count with np.add.at, and read the probabilities of the pairs with integer-array indexing.
• With alpha = 0, a row of a token that is never followed by anything would be $0/0$ (row 2 in the example). You will not be tested on that case, but notice why smoothing exists.

Hints

• The previous tokens of all pairs are tokens[:-1] and the next tokens are tokens[1:]: for the example, [0, 1, 0, 1, 0] and [1, 0, 1, 0, 2].
• np.add.at(counts, (prev, next), 1) adds 1 to counts[prev[k], next[k]] for every k, and counts repeated pairs correctly (see the embedding exercise for why counts[prev, next] += 1 does not).