Next-word prediction: how a chatbot's core trick is just counting

Type "How are" into a phone and it offers "you". Ask a chatbot a question and it replies with whole paragraphs. Both are doing the same basic job: guess the next word. Real chatbots do it with enormous maths, but the idea is small enough to build yourself in ten lines, and building it makes the big version much less mysterious.

Prediction is just counting

Suppose you have read this text:

the cat sat on the mat. the cat ate the fish. the dog sat on the rug. the dog ate the bone.

Which word tends to come after "cat"? Look at every place "cat" appears: once it is followed by "sat", once by "ate". So after "cat", "sat" and "ate" are equally likely. After "the", the picture is wider: "cat" and "dog" appear twice each, while "mat", "fish", "rug" and "bone" appear once each. That is the whole trick. A model that looks only at the single previous word is called a bigram model ("bi" means two: the word you have, and the word you are guessing).

The same idea in Python

First, split the text into tokens (pieces the program treats as one unit; here, words and full stops). Then, for every pair of neighbouring tokens, add one to a counter.

import re
from collections import Counter, defaultdict

text = "the cat sat on the mat. the cat ate the fish. the dog sat on the rug. the dog ate the bone."
words = re.findall(r"[a-z']+|[.]", text.lower())

follows = defaultdict(Counter)
for a, b in zip(words, words[1:]):
    follows[a][b] += 1

print(follows["the"])
print(follows["cat"])

Running it prints:

Counter({'cat': 2, 'dog': 2, 'mat': 1, 'fish': 1, 'rug': 1, 'bone': 1})
Counter({'sat': 1, 'ate': 1})

That table of counts is the model. "Training" here means nothing more than filling it in by reading text.

Try it: train a model and make it talk

Some things to try. Press "Pick by chance" a dozen times: the sentences differ every time, but they only ever contain word pairs that were in the training text. Press "Add most likely word" repeatedly and watch it fall into a loop ("the cat ate the cat ate the cat...") because the single most common choice always leads back to the same place. Then replace the training text with a few sentences of your own and see what the model learns from them.

Choosing, not just ranking

You saw two ways to pick the next word. Always take the most likely is predictable and repeats itself. Pick at random, weighted by the counts gives variety: a word seen twice is twice as likely to be chosen as a word seen once. In Python that is one line:

import random

def next_word(word):
    options = follows[word]
    return random.choices(list(options), weights=list(options.values()))[0]

random.seed(4)
word, out = "the", ["the"]
for _ in range(8):
    word = next_word(word)
    out.append(word)
print(" ".join(out))

With that seed on Python 3.13 it printed the cat sat on the cat sat on the. Without the random.seed line, you get something different each run. Chatbots have a dial for exactly this choice, called temperature: low means "stick to the likeliest", high means "take more chances".

What this model cannot do

Our model has a short memory: it sees only one word. "the dog sat on the" and "the cat ate the" both end in "the", so it treats them identically and has no idea whether the sentence is about a dog or a cat. It also has a dead end: the last "." in the text is never followed by anything, so if the model lands there, it has nothing to say. The tool above tells you when that happens.

The fix for the short memory is to look at more words. A model that uses the previous two words (a trigram) is already better, and the idea can be pushed further, but the table gets huge: there are far too many possible word sequences to ever count each one. This is the problem that neural networks solve. Instead of a giant lookup table, they learn a compact set of numbers that can estimate the next-word odds for sequences they have never seen exactly. Modern language models use the same task, predict the next piece of text, with a far longer window and far better estimates, but they do not store the counts we used here.

So what, for you

When a chatbot says something fluent, it is because fluent text was what predicting the next token rewarded in training. That is a useful mental model: the output is a likely continuation, not a looked-up fact, which is one reason these tools can sound confident and still be wrong. Check anything that matters. And if you ever want to see the idea for yourself, you now have a working language model in about ten lines.