Words as numbers: a tiny next-word model

Count which word follows which, turn the counts into chances, and roll a die. That's a language model. It's tiny and a bit silly, and it's the great-great-grandparent of ChatGPT. Then you'll give it a real job on the robot: fixing typos in your commands, the way your phone keyboard does.

Going deeper: everything below is the full lesson: the idea built from nothing, two demos, your homework, and a challenge away from the screen.

1. The problem

Here's the only text our model will ever read. It's the STORY in lesson08_words.py:

the robot picks up the red block . the robot puts the red block in the bin .
the robot sees a cat . the robot stops . the cat walks away . the robot moves on .
the robot picks up the blue block . the robot puts the blue block in the bin .
the robot is happy . the room is tidy .

We want a program that can write new sentences in the same style. How?

Obvious idea 1: pick random words from the story. You get "block happy the in sees away". Every word is real, but the order is nonsense. Word order is the whole point of language.

Obvious idea 2: write grammar rules by hand. "A sentence is the + a noun + a verb…" That works for 5 sentences, then you need thousands of rules, and English breaks them all the time. People tried this for decades. It never scaled.

The idea that does work: don't write rules. Let the text itself tell you what tends to come next.

2. First principles

Step 1: chop text into tokens

Computers can't count "words" until we decide what a word is. Our rule: lowercase everything, keep runs of letters (and apostrophes) as words, and keep the full stop as its own token, because "end of sentence" is useful to know. Everything else (commas, question marks, numbers) is thrown away.

"The cat. Sat"                      → ["the", "cat", ".", "sat"]
"Where's the cat? It's 5 o'clock!"  → ["where's", "the", "cat", "it's", "o'clock"]

Notice "where's" stays as one token. Remember that. It causes a real bug in Lesson 9.

The whole story becomes 66 tokens, using 23 different ones.

Step 2: count what follows what

Slide along the tokens, two at a time, and keep a tally: "after this word, which word came next?"

the robot picks up the red block .
─────────
    ──────────
          ────────
             ────── ...
pairs: (the, robot) (robot, picks) (picks, up) (up, the) (the, red) ...

Do this for the whole story and here's what follows "robot":

after "robot":   picks 2   puts 2   sees 1   stops 1   moves 1   is 1     (8 in total)

In Python we store this as a dictionary of dictionaries: counts["robot"]["picks"] == 2. That's the whole "brain" of our model. Just counts.

Step 3: turn counts into probabilities

Counts depend on how long the text is. We want chances. Divide each count by the total:

picks  2 / 8 = 0.25
puts   2 / 8 = 0.25
sees   1 / 8 = 0.125
stops  1 / 8 = 0.125
moves  1 / 8 = 0.125
is     1 / 8 = 0.125
              ─────
total         1.0    ← probabilities always add up to 1

So after "robot", there's a 1-in-4 chance of "picks". That's what next_word_probs returns. For "the", it's robot 0.5, red 0.125, bin 0.125, blue 0.125, cat 0.0625, room 0.0625.

Step 4: sampling, or how to roll a weighted die

Now we write. We could always pick the most likely word, but then "the" is always followed by "robot", and "robot" always by "picks"… the model would repeat the same sentence forever. Instead we pick randomly, but fairly: likely words come up more often.

Picture a ruler from 0 to 1, cut into pieces as long as each probability:

after "the":
0             0.5    0.625  0.75  0.8125 0.9375  1
|    robot     | red  | bin  | cat |  blue | room |

(The pieces go in the order the words were first seen after "the": robot, red, bin, cat, blue, room.)

Pick a random number between 0 and 1 and see which piece it lands in. Python's random.Random(4) gives 0.236 first. That lands in "robot", so the next word is "robot". Python does this for you with rng.choices(words, weights=probs).

Step 5: repeat, and you've built a language model

start with ["the"]
look at the LAST word → get its next-word probabilities → sample one → add it
look at the LAST word → ...
stop after n new words, or if the last word never had anything after it

Here's real output from the solution (seed 4, 12 new words each):

-> the robot picks up the robot puts the cat walks away . the
-> the robot picks up the blue block in the robot stops . the
-> the blue block in the red block . the blue block in the

Every pair of neighbouring words appeared in the story. But the sentences wander: "picks up the robot puts the cat". Why? Our model only remembers one word. After "the" it has no idea it was in the middle of "picks up the …".

Step 6: so how is ChatGPT different?

The loop is the same: give the chances of each next token, pick one, add it, repeat. When an AI chatbot writes to you, it's doing your generate loop, one token at a time. What modern LLMs add:

• Much longer memory. Ours looks at 1 previous word. Modern models look at thousands, even hundreds of thousands, of previous tokens: your whole question and the conversation so far.
• A neural network instead of a count table. Why not just count longer word groups? Try it: our story has 23 tokens, so a table for "last word → next" has 23 × 23 = 529 slots. English has about 50,000 common words. For "last 10 words → next" you'd need 50,000¹¹ slots, more than the number of atoms in the Earth, and almost all of them would be zero because nobody ever wrote those exact 10 words. A neural network (the neurons from Lesson 4, stacked very deep, using a trick called attention) learns patterns instead, so it can handle word combinations it has never seen.
• Huge amounts of text. Ours read 66 tokens. Big models read trillions, from books, websites and code.
• Smarter tokens. Instead of whole words, they use word pieces, so new words like "unboxable" become "un" + "box" + "able".
• Extra training to be helpful. After learning to predict text, they're trained further to follow instructions and to refuse harmful requests.

But under all that, the job is still: predict the next token. You've just built the smallest possible version.

Step 7: the robot's first job, fixing typos

People type fast and make mistakes. Next lesson's word counter (Lesson 9) has never seen "bleu", so for "go to the bleu block" it finds no object and replies Sorry, I can't do that: 'go' needs an object. For the correct spelling it drives to the blue block. A next-word model can fix that. Two parts.

Part 1: list every word one typo away. There are four kinds of typo. For "cat":

delete one letter          cat → at
swap two neighbours        cat → act
replace one letter         cat → bat   (try all 26 letters in each place)
insert one letter          cat → cart  (try all 26 letters in each gap)

A 3-letter word has 3 deletes, 2 swaps, 3 × 26 = 78 replaces and 4 × 26 = 104 inserts: 187 in all. Some come out the same (replacing "c" with "c" gives "cat" back), so the set has 182 different strings. Almost all are nonsense like "zat". We keep only the known words: words the model has counted.

Part 2: let the word before decide. For the robot, the known words come from Lesson 9's 34 example commands (53 different tokens). Here's the same typo with two different words before it:

"go te ..."   known words one typo from "te": me, the, to
              seen after "go": to (0.25)                → to
"put te ..."  known words one typo from "te": me, the, to
              seen after "put": the (0.67)              → the

Same typo, different fix, because of the word before. That's the bigram model at work. Two fallbacks: if none of the candidates was ever seen after the previous word, pick the one that starts the most pairs (the most common word); if there are no known candidates at all, leave the word alone. A word at the start of a command uses "." as its previous word, because commands are joined with full stops.

Real output on the 9 TYPO_COMMANDS in the file:

'put teh red blok in the bin'    -> 'put the red block in the bin'
'go to the bleu block'           -> 'go to the blue block'
'whre is the green block'        -> 'where is the green block'
'fnd the red block'              -> 'find the red block'
'stpo'                           -> 'stop'
'put the gren block away'        -> 'put the green block away'
'thorw away the blue block'      -> 'throw away the blue block'
'head to the red blokc'          -> 'head to the red block'
'put the red bclok away'         -> 'put the red bclok away'   (wanted 'put the red block away')
fixed 8 of 9

The last one fails on purpose. "bclok" has two typos (a swap and a swap), and we only look one edit away. You could look two edits away, but "bclok" has 286 one-edit strings and 36,860 two-edit strings, so there are many more chances of a wrong guess.

3. Play with it

–

4. Homework

Open homework/lesson08_words.py. STORY, LETTERS, TYPO_COMMANDS and known_words are written for you. Write these seven functions, in order. The first four are the language model; the last three are the typo fixer.

4.1 tokenize(text)

Lowercase the text and return a list of tokens: words (letters and apostrophes) plus "." as its own token.

Example: tokenize("The cat. Sat") → ["the", "cat", ".", "sat"].

4.2 count_bigrams(tokens)

Return a dictionary of dictionaries: counts[word][next_word] = how many times next_word came right after word.

Example: count_bigrams(["a", "b", "a", "b", "c"]) → {"a": {"b": 2}, "b": {"a": 1, "c": 1}}. (The pairs are a-b, b-a, a-b, b-c.)

4.3 next_word_probs(counts, word)

Turn counts[word] into probabilities that add up to 1. If word was never seen, return {}.

Example: next_word_probs({"b": {"a": 1, "c": 3}}, "b") → {"a": 0.25, "c": 0.75}.

4.4 generate(counts, start, n, rng)

Start with [start]. Repeat up to n times: get the probabilities after the last word, sample one with rng.choices, add it. Stop early if there are no options. Return the words joined with spaces.

Example: generate(counts, "zzz", 5, rng) → "zzz" (no options after "zzz", so it stops at once). With a real start word and n = 5, you get 6 words.

4.5 edits1(word)

Every string one typo away from word: delete one letter, swap two neighbouring letters, replace one letter (with any of LETTERS), or insert one letter (anywhere, including the very start and end). Return a set.

Example: edits1("cat") contains "at", "act", "bat" and "cart" (and "cat" itself, from replacing a letter with the same letter). len(edits1("ab")) is 130.

4.6 fix_word(word, prev, counts, known)

Fix one word. If word is in known, keep it. Otherwise look at the known words in edits1(word):

• if any of them was seen right after prev, pick the one with the highest next_word_probs(counts, prev);
• otherwise pick the one that starts the most pairs, sum(counts.get(w, {}).values()), with ties going to the alphabetically first;
• if there are none, keep word as it is.

Example: fix_word("bok", "red", {"red": {"box": 1, "block": 3}}, {"red", "box", "block"}) → "box". "block" is likelier after "red", but it's two typos from "bok", so it isn't a candidate.

4.7 fix_typos(text, counts)

Fix each token of tokenize(text) with fix_word, using the previously fixed word as prev. Start with prev = ".". Return the fixed tokens joined with spaces.

Example: with counts from "put the red block in the bin . put the blue block in the bin", fix_typos("put teh red blok", counts) → "put the red block".

Run it

cd homework
python3 lesson08_words.py

What you should see (real output of the reference solution):

after 'robot': {'picks': 0.25, 'puts': 0.25, 'sees': 0.12, 'stops': 0.12, 'moves': 0.12, 'is': 0.12}
-> the robot picks up the robot puts the cat walks away . the
-> the robot picks up the blue block in the robot stops . the
-> the blue block in the red block . the blue block in the

Fixing typos with the robot's command words (from Lesson 9's examples):
  'put teh red blok in the bin'    -> 'put the red block in the bin'
  'go to the bleu block'           -> 'go to the blue block'
  'whre is the green block'        -> 'where is the green block'
  'fnd the red block'              -> 'find the red block'
  'stpo'                           -> 'stop'
  'put the gren block away'        -> 'put the green block away'
  'thorw away the blue block'      -> 'throw away the blue block'
  'head to the red blokc'          -> 'head to the red block'
  'put the red bclok away'         -> 'put the red bclok away'   (wanted 'put the red block away')
fixed 8 of 9

The typo part needs your Lesson 9 file, because it borrows TRAIN from lesson09_commands.py. If you haven't done Lesson 9 yet, the first part still runs; come back to the typo part next week. Why 0.12 and not 0.125? The program rounds to 2 decimal places for printing, and Python rounds 0.125 to the even neighbour, 0.12.

Check it

cd ..
python3 check.py 8

Seven checks: tokenize, count_bigrams, next_word_probs, generate, edits1 and fix_typos (which also tests your fix_word). Green ✓ per check, red ✗ with a hint, yellow … for "not written yet". Real hints you might see:

• ✗ edits1: 'ab' has exactly 130 different one-typo versions, got 133. You returned a list with repeats. Turn it into a set.
• ✗ fix_typos: 'bok' is one typo from 'box' but two from 'block' -> 'box'. Your fix_word picked from all the words seen after prev, not just the candidates one typo away.
• ✗ fix_typos: known words and hopeless words stay as they are. "purple" has no known word one typo away, so it must come back unchanged.
• ✗ count_bigrams: count_bigrams(['a','b','a','b','c']) should be {'a': {'b': 2}, 'b': {'a': 1, 'c': 1}}. Compare with your pairs on paper. A missing pair or an extra one at the end is the usual cause.
• ✗ generate: stop when there are no options. Your loop crashed or kept going after a word with nothing after it.

Stretch

• Make the fixer reach two typos: if no known word is one edit away, try every edit of every edit. Does it fix "bclok"? How long does edits1 of every string in edits1("bclok") take, and does it start "fixing" words it shouldn't?
• Replace STORY with a page of your own writing (an essay, a diary entry). Does the output "sound like you"?
• Build a trigram model: count which word follows each pair of words (key ("the", "robot")). Are the sentences less wandering? What goes wrong with only 66 tokens of training text?
• Add a greedy option that always picks the most likely word. What happens to the output?

5. Real-world challenge: a language model in a cup

You'll train and run a bigram model by hand on your own writing.

1. Pick about 20 messages you sent (not other people's), or a page of your own writing. Copy them onto paper. Leave out anything private.
2. Choose your 5 most common words (often "i", "the", "to", "you", "a"). For each, make a tally table of the word that came right after it.
3. Pick one word, say "i". For each follower, write it on scraps of paper: one scrap per tally mark. If "i" was followed by "am" 3 times and "think" once, that's 3 "am" scraps and 1 "think" scrap. Put them in a cup. Now the cup is the probabilities.
4. Generate: draw a scrap, write it down, then use the cup for that word (make cups for your other common words too; if you land on a word with no cup, stop). Make 3 "sentences".
5. Then type the same messages into the STORY in your Python file (or into the demo box above) and compare the computer's sentences with yours.

Record: your tally tables, your 3 hand-made sentences, and one probability you worked out (e.g. P(am after i) = 3/4 = 0.75). What it should teach you: the "model" is just your own habits, counted. It sounds a bit like you, but it has no idea what it's saying. A model can only be as good as the text it learned from.

6. Your Robo-Helper so far

You've added a baby language model (tokenize, count_bigrams, next_word_probs, generate) and its first robot job, a typo fixer (edits1, fix_word, fix_typos).

The robot won't write stories in the final project. Two pieces live on: the typo fixer, and tokenize. In Lesson 9, lesson09_commands.py starts with from lesson08_words import tokenize and uses it to chop up commands like "put the red block in the bin" before counting words to guess what you want. Then Lesson 11's capstone uses that command reader, so your tokenizer ends up inside the finished Robo-Helper. And in Lesson 10 you'll plug in a real LLM: the grown-up version of what you just made.

7. Common mistakes

• Forgetting to lowercase. Then "The" and "the" are counted as different words. Checker hint: 'tokenize("The cat. Sat") -> ["the", "cat", ".", "sat"]'.
• Off-by-one in the pairs. Looping i all the way to len(tokens) makes tokens[i + 1] crash on the last word.
• Dividing by the wrong number. Divide by the total count (4 in the check), not by how many different next words there are (2). Checker hint: "1 and 3 -> 0.25 and 0.75".
• Looking at the first word instead of the last. out[0] is always the start word. You need out[-1].
• Returning a list from edits1. Some edits come out the same, so a list has repeats. Checker hint: "'ab' has exactly 130 different one-typo versions".
• Using the typo as prev. In fix_typos, the next word must be chosen by the fixed word before it. "teh" was never seen before "red"; "the" was.
• Counting the start word as one of the n. n = 5 means 5 new words, 6 in total. Checker hint: "start word + 5 new words".

8. Check yourself

9. Go deeper (optional)

• 30u30 Day 01: The Unreasonable Effectiveness of RNNs: the next step up from counting. A neural network that predicts the next character learns spelling and grammar by itself. Has a temperature demo.
• 30u30 Day 13: Attention Is All You Need: the paper behind the "T" in GPT. How a model looks back at all earlier tokens, not just the last one.
• Post-training: how a raw next-word predictor is turned into a helpful assistant.
• Inference engineering: how the generate loop is run fast for millions of people.