A Bigram Text Generator

Build a Tiny Language Model

Example

A bigram model predicts the next word from only the previous word. We "train" it on a few ShopEase support lines, then generate new sentences by sampling (picking randomly, weighted by how often each word followed). A fixed seed makes the output repeatable.

import random
from collections import defaultdict

training_text = """your order is on the way .
your order is delayed by one day .
your refund is on the way .
your refund is processed to your bank .
the delivery partner is on the way ."""

next_words = defaultdict(list)
words = training_text.split()
for a, b in zip(words, words[1:]):
    if a != ".":
        next_words[a].append(b)       # duplicates = higher probability

rng = random.Random(2)
def generate(start, max_words=12):
    out = [start]
    while out[-1] != "." and len(out) < max_words:
        out.append(rng.choice(next_words[out[-1]]))
    return " ".join(out)

for _ in range(5):
    print(generate("your"))

Output

your order is on the way .
your refund is on the way .
your order is processed to your bank .
your order is on the delivery partner is on the way .
your order is delayed by one day .

What This Shows

  • Some outputs are new sentences that never appeared in training, such as "your order is processed to your bank" (the model generalizes by recombining patterns).
  • Some are odd, like "your order is on the delivery partner is on the way": it only sees one word of context, so it cannot keep track of meaning.
  • Real LLMs replace "count the previous word" with a transformer that looks at thousands of previous tokens, which is why their text stays coherent.

Exercise

Exercise: Change the generator into a trigram model that uses the previous two words as the key. What happens to the output?

Show answer
from collections import defaultdict
import random

words = training_text.split()
nxt = defaultdict(list)
for a, b, c in zip(words, words[1:], words[2:]):
    nxt[(a, b)].append(c)

rng = random.Random(42)
out = ["your", "order"]
while out[-1] != "." and (out[-2], out[-1]) in nxt and len(out) < 15:
    out.append(rng.choice(nxt[(out[-2], out[-1])]))
print(" ".join(out))
# More context -> sentences copy the training text more closely
# (more coherent, less creative).