Local LLM Field Guide
Search

Intermediate14 min read

Train a tiny next-token model from counts

Count token pairs, turn them into probabilities, sample a sentence, and score predictions on text the model never saw.

Learn a probability table

A bigram model predicts the next token using the previous one. This example splits short English sentences into words, adds start and end markers, and counts each pair. Given the word “at,” the model assigns probability using the counts of words that followed “at.”

The script adds a small count to every possible outcome. This is additive smoothing. It gives unseen pairs a positive probability. A unigram baseline uses the same vocabulary but uses overall word counts instead of the preceding word.

Download train_bigram.py
View train_bigram.py
import argparse
from collections import Counter
import math
import random
import re

TRAIN = [
    "the library opens at nine",
    "the library closes at five",
    "the reading room opens at nine",
    "the reading room closes at five",
    "the media desk opens at ten",
    "the media desk closes at four",
    "ask the media desk for a scanner",
    "ask the front desk for a card",
    "return the book at the front desk",
]
TEST = [
    "the library opens at ten",
    "the reading room closes at four",
    "ask the media desk for a card",
    "return the magazine at the front desk",
]
START, STOP, UNKNOWN = "<start>", "<end>", "<unk>"


def tokenize(text):
    return re.findall(r"[a-z]+", text.casefold())


def train(documents):
    unigram = Counter()
    bigram = {}
    for text in documents:
        tokens = [START, *tokenize(text), STOP]
        for current, following in zip(tokens, tokens[1:]):
            unigram[following] += 1
            bigram.setdefault(current, Counter())[following] += 1
    vocabulary = sorted({*unigram, UNKNOWN})
    return unigram, bigram, vocabulary


def distribution(counts, vocabulary, alpha):
    denominator = counts.total() + alpha * len(vocabulary)
    return {token: (counts[token] + alpha) / denominator for token in vocabulary}


def perplexity(documents, rows, vocabulary):
    loss = count = 0
    for text in documents:
        tokens = [START, *(token if token in vocabulary else UNKNOWN for token in tokenize(text)), STOP]
        for current, following in zip(tokens, tokens[1:]):
            loss -= math.log(rows[current][following])
            count += 1
    return math.exp(loss / count)


def sample(rows, vocabulary, seed):
    rng = random.Random(seed)
    current = START
    words = []
    for _ in range(18):
        current = rng.choices(vocabulary, weights=list(rows[current].values()))[0]
        if current == STOP:
            break
        words.append(current)
    return " ".join(words) or "(end at first step)"


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Train a count-based next-token model on fictional notes.")
    parser.add_argument("--alpha", type=float, default=0.25, help="Smoothing count added to each option")
    parser.add_argument("--seed", type=int, default=7)
    args = parser.parse_args()
    if not math.isfinite(args.alpha) or not 0 < args.alpha <= 1_000:
        parser.error("Choose a positive smoothing count up to 1000.")
    unigram, bigram, vocabulary = train(TRAIN)
    baseline = distribution(unigram, vocabulary, args.alpha)
    contexts = [START, *(token for token in vocabulary if token != STOP)]
    unigram_rows = {token: baseline for token in contexts}
    bigram_rows = {token: distribution(bigram.get(token, Counter()), vocabulary, args.alpha)
                   for token in contexts}
    probabilities = bigram_rows["at"]
    print(f"Training vocabulary: {len(vocabulary)} possible next tokens")
    print("After 'at':", ", ".join(f"{token} {score:.3f}" for token, score in
                                 sorted(probabilities.items(), key=lambda item: -item[1])[:5]))
    print(f"Unigram test perplexity: {perplexity(TEST, unigram_rows, vocabulary):.2f}")
    print(f"Bigram test perplexity: {perplexity(TEST, bigram_rows, vocabulary):.2f}")
    print("Bigram sample:", sample(bigram_rows, vocabulary, args.seed))
python train_bigram.py
You’re ready when

The program prints a probability distribution summary after “at,” one perplexity for each model, and one sampled continuation. With the supplied defaults, the unigram scores about 18.48 and the bigram about 5.92 on the same four test sentences.

Check predictions on new text

The script scores how much probability the model gives to each observed next token. It averages the negative log probabilities, then exponentiates that average to obtain perplexity. Lower perplexity here means higher probability for the test tokens under the same vocabulary and tokenization. Words missing from training, such as “magazine,” use one shared unknown token.

python train_bigram.py --alpha 1
python train_bigram.py --seed 23

Changing smoothing can change both the scores and the sample. Changing the seed changes the sample while keeping the scores fixed. Add a training sentence with a word that follows “at,” then compare the probabilities while leaving the test set untouched.

Connect to language model training

Modern decoder language models learn weights that predict the next token using more context. This lab uses a table of counts and one preceding token. Both can be scored on held-out predictions, and both can generate by adding a sampled token and predicting again. Perplexity on this tiny set cannot establish that a model follows instructions or answers accurately.

Check your understanding

Why use the same tokenizer and test sentences for both scores?

Perplexity depends on which tokens you ask each model to predict. Keeping those tokens fixed makes this comparison meaningful.

Sources

Search the guide