Build a next-token prediction baseline

Problem

When you need this. A falling training loss does not tell you whether a transformer is better than a simple counting model. Build that reference and use the same scoring rule for both.

The next builds train a transformer on the story corpus. Before that, you need two things: a way to score any next-token model, and a simple model to compare against. Without them, a falling training loss tells you nothing about whether the transformer predicts unseen stories better than counting would.

This build writes three count-based baselines and one scorer. The scorer takes a model and a list of tokenized stories and returns the mean cross-entropy of each next token given the token before it. All three models are scored on the same 2,109 token positions of 60 test-story rows held out by position. Some rows repeat text from training; the split check below makes that limitation visible.

The build also measures one prediction that matters for this corpus. Every story ends with a sentence like “now leo has the red ball.”, repeating the color and object from its first sentence. A good model should predict the color after “now leo has the”. This build records how often the best baseline does.

Success has three checks, all reproducible with the supplied tests. Each model returns probabilities that add up to 1. The losses order as bigram < unigram < uniform on the test stories. The smoothing setting is chosen on the development stories only, and the test stories are scored once.

Concepts

This build holds out data and splits by whole story, in file order: the first 280 stories train the counts, the next 60 are development stories for choosing a setting, and the last 60 are test stories, scored once at the end. The generator in make_corpus.py wrote all 400 stories before any model existed, so no test story was written or picked with a model’s results in view. The stories are original synthetic text written for this series; no third-party text is included. Two limits matter. The tokenizer was trained on the whole corpus, including development and test text. Also, the story generator can emit identical stories in different row positions: 5 of the 60 development rows duplicate training text, and 5 of the 60 test rows duplicate training text. Eight test rows appear in training or development. The count models fit only the training rows, but that does not make every test story new. Keep the recorded split for the comparisons in this course; a new experiment should deduplicate or group identical stories before splitting and fit its tokenizer on training text only.

Cross-entropy is the negative log of the probability given to the correct answer. Here the answer is the next token, and the score is the mean over every scored position. Perplexity is e raised to that mean. Textbooks often write it with base-2 logarithms and 2 raised to the mean; the perplexity is the same. It reads as the number of equally likely choices the model is, on average, choosing between: a model that gives every one of the 96 tokens equal probability has perplexity 96.

The three baselines differ in what they look at:

Uniform

Every token has probability 1 / 96. It looks at nothing, and its loss is ln 96 = 4.5643.

Unigram

Each token’s share of all training tokens, with one added to every count (add-one smoothing) so that no token gets probability 0. It knows which tokens are common but ignores context.

Bigram

How often each token followed the previous token in training. It looks exactly one token back.

A bigram estimated from raw counts gives probability 0 to any pair it never saw, and one such pair in the test stories makes the loss impossible to compute, because the log of 0 is undefined. Add-k smoothing adds a small count k to every possible pair before dividing, so unseen pairs get a little probability. The value of k is a setting chosen on the development stories. Add-k is the simplest smoothing method, not the best one for language models; it is enough for a baseline.

Architecture

The figure shows what the bigram uses for one prediction. The context “now leo has the” ends in the token “ the”. Bigram.probs(prev) reads only that last token. It selects the count row for “ the”, which holds the counts of every token that followed “ the” in training, and argmax(probs) turns the most frequent one into the next-token guess. The dashed box marks everything before the last token: the bigram never reads it.

baseline.py holds the three models, the scorer, a recall check and a sampler. Each model has one method, probs(prev), that returns a probability for every token in the vocabulary. The scorer calls only that method. Later builds apply the same mean next-token loss to the transformer, which reads the whole allowed prefix rather than only prev. trace_baseline.py loads tokenizer.json, builds the splits, chooses k, scores the test stories and writes baseline-trace.json. The corpus and tokenizer files are copies of those from the previous build.

Working code

Runtime: Python 3.10 or newer. Dependencies: Python standard library only.

Download and extract the complete lab using Download the complete lab above. Open build-a-next-token-prediction-baseline/; keep its Python modules, tests and fixture files together. The bundled README includes setup commands for your platform.

"""Count-based next-token baselines for the story corpus, standard library only.

Three models predict the next token ID:
  uniform  - every vocabulary entry equally likely
  unigram  - how often each token appears in the training stories
  bigram   - how often each token follows the previous token, with add-k smoothing

Each model returns a probability for every token in the vocabulary, so all three
are scored the same way: mean cross-entropy (natural log) on held-out tokens.
"""
import math
import random
from collections import Counter

from bpe import Tokenizer


def split_stories(text, train=280, dev=60):
    """Split by story, in file order: training, development, then test."""
    stories = [s + "\n" for s in text.split("\n") if s]
    if train + dev >= len(stories):
        raise ValueError(f"need more than {train + dev} stories, got {len(stories)}")
    return stories[:train], stories[train:train + dev], stories[train + dev:]


def encode_stories(tok, stories):
    return [tok.encode(s) for s in stories]


class Uniform:
    def __init__(self, vocab_size):
        self.vocab_size = vocab_size

    def probs(self, prev):
        return [1.0 / self.vocab_size] * self.vocab_size


class Unigram:
    def __init__(self, sequences, vocab_size):
        counts = Counter(t for seq in sequences for t in seq)
        total = sum(counts.values())
        # add-one smoothing so a token unseen in training still gets a probability
        self.p = [(counts[i] + 1) / (total + vocab_size) for i in range(vocab_size)]

    def probs(self, prev):
        return self.p


class Bigram:
    def __init__(self, sequences, vocab_size, k):
        if not (k > 0 and math.isfinite(k)):
            raise ValueError(f"smoothing k must be a positive finite number, got {k!r}")
        self.vocab_size, self.k = vocab_size, k
        self.counts = [Counter() for _ in range(vocab_size)]
        for seq in sequences:
            for a, b in zip(seq, seq[1:]):
                self.counts[a][b] += 1
        self.totals = [sum(c.values()) for c in self.counts]

    def probs(self, prev):
        c, total = self.counts[prev], self.totals[prev]
        denom = total + self.k * self.vocab_size
        return [(c[i] + self.k) / denom for i in range(self.vocab_size)]


def cross_entropy(model, sequences):
    """Mean negative log probability of each token given the token before it.

    The first token of each story has no previous token in the story, so it is
    not scored; every model is scored on exactly the same positions.
    """
    losses = []
    for seq in sequences:
        for prev, nxt in zip(seq, seq[1:]):
            p = model.probs(prev)[nxt]
            if p <= 0:
                raise ValueError(f"model gave probability {p} to an observed token")
            losses.append(-math.log(p))
    if not losses:
        raise ValueError("no token pairs to score")
    return math.fsum(losses) / len(losses), len(losses)


def argmax(probs):
    return max(range(len(probs)), key=lambda i: (probs[i], -i))


def recall_check(tok, model, stories, field):
    """Predict the color or object in each story's last sentence from the token before it.

    field="color": context ends with "now <name> has the"; the answer is " <color>".
    field="object": context ends with "now <name> has the <color>"; the answer is " <object>".
    """
    correct, guesses = 0, Counter()
    for story in stories:
        last = story[story.index("now "):].rstrip(".\n")
        words = last.split()  # now, name, has, the, color, object
        color, thing = words[4], words[5]
        cut = story.rindex(" " + color) if field == "color" else story.rindex(" " + thing)
        # The answer word can be several tokens (" kite" is " kit" + "e"); score its first token.
        answer = tok.tokens(" " + (color if field == "color" else thing))[0]
        context = tok.encode(story[:cut])
        guess = tok.vocab[argmax(model.probs(context[-1]))]
        guesses[guess] += 1
        correct += guess == answer
    return {"correct": correct, "total": len(stories), "guesses": dict(sorted(guesses.items()))}


def sample(model, tok, start, length, rng, greedy=False):
    ids = tok.encode(start)
    for _ in range(length):
        p = model.probs(ids[-1])
        ids.append(argmax(p) if greedy else rng.choices(range(len(p)), weights=p)[0])
    return tok.decode(ids)

cross_entropy() scores each token given the one before it and skips the first token of each story, which has no previous token. Every model is scored on exactly the same positions, so their losses are comparable on this split. It raises an error if a model gives probability 0 to a token that occurs, instead of returning an infinite loss. Bigram stores one Counter per previous token. recall_check() cuts each test story just before the color or object in its last sentence and asks the model for its single most likely next token.

Run the trace:

python3 trace_baseline.py

python3 trace_baseline.py --check exits with an error if your output differs from the recorded baseline-trace.json.

Inspect the split before fitting

The bundled checker reads the same whole-story split and counts rows whose exact text occurs in an earlier split. It does not fit a model or change the corpus. Distinct row ranges prevent one row from entering two splits; they do not prevent duplicate text.

"""Inspect exact story-text overlap without fitting or changing any model."""
import json
from pathlib import Path
from baseline import split_stories

HERE = Path(__file__).parent


def overlap_counts(text):
    train, dev, test = split_stories(text)
    training_text = set(train)
    earlier_text = training_text | set(dev)
    return {
        "rows": {"train": len(train), "dev": len(dev), "test": len(test)},
        "dev_rows_matching_training_text": sum(row in training_text for row in dev),
        "test_rows_matching_training_text": sum(row in training_text for row in test),
        "test_rows_matching_train_or_dev_text": sum(row in earlier_text for row in test),
    }


if __name__ == "__main__":
    print(json.dumps(overlap_counts((HERE / "corpus.txt").read_text()), indent=2))

python3 inspect_split.py

{
  "rows": {
    "train": 280,
    "dev": 60,
    "test": 60
  },
  "dev_rows_matching_training_text": 5,
  "test_rows_matching_training_text": 5,
  "test_rows_matching_train_or_dev_text": 8
}

Choose k on the development stories

The trace scores the bigram on the 60 development stories for 4 values of k:

The lowest loss is at k = 0.001, the smallest value tried. In this corpus, 0 of the 2,121 token pairs in the development stories are missing from training, so probability moved to unseen pairs is wasted, and less smoothing helps. The best k may be smaller still; the grid did not test below 0.001.

Score the test stories

With k fixed, each model is scored once on the test stories:

Knowing which tokens are common cuts the perplexity from 96 to 31.89. Knowing the previous token cuts it to 3.22, because most of each story is a fixed template. The bigram’s test loss, 1.1702, is the number the transformer has to beat.

Generate text

sample() starts from “mia has” and repeatedly draws the next token from the bigram’s probabilities, with seed 1:

mia has the kitchen. at the book. mia.
 y s sch� bo sam has a red cup.
ow cue.
trez cup.

Every adjacent pair occurred in training, but the text does not hold together beyond two tokens. The � is the <unk> token: smoothing gave it a small probability, and the sampler drew it. Taking the most likely token every time instead, which is greedy decoding, falls into a loop:

mia has the school, sam gives the school, sam gives the school, sam gives the school, sam gives the school, sam gives

Check your understanding

The recall check asks the bigram to predict the color after “now leo has the” in each of the 60 test stories. There are four colors. Before you open the comparison, predict how many it gets right, and what single token it guesses.

{
  "vocab_size": 96,
  "split": {
    "train": 280,
    "dev": 60,
    "test": 60,
    "rule": "stories in file order: first 280 train, next 60 dev, last 60 test"
  },
  "uniform_loss_formula": 4.5643,
  "dev_pairs": {
    "total": 2121,
    "unseen_in_training": 0
  },
  "dev_bigram_by_k": {
    "0.001": {
      "loss": 1.1683,
      "perplexity": 3.22,
      "tokens": 2121
    },
    "0.01": {
      "loss": 1.1722,
      "perplexity": 3.23,
      "tokens": 2121
    },
    "0.1": {
      "loss": 1.2105,
      "perplexity": 3.36,
      "tokens": 2121
    },
    "1": {
      "loss": 1.5009,
      "perplexity": 4.49,
      "tokens": 2121
    }
  },
  "chosen_k": 0.001,
  "test": {
    "uniform": {
      "loss": 4.5643,
      "perplexity": 96.0,
      "tokens": 2109
    },
    "unigram": {
      "loss": 3.4623,
      "perplexity": 31.89,
      "tokens": 2109
    },
    "bigram": {
      "loss": 1.1702,
      "perplexity": 3.22,
      "tokens": 2109
    }
  },
  "recall": {
    "color": {
      "correct": 0,
      "total": 60,
      "guesses": {
        " school": 60
      }
    },
    "object": {
      "correct": 6,
      "total": 60,
      "guesses": {
        " book": 18,
        " box": 10,
        " kit": 15,
        ".": 17
      }
    }
  },
  "top_after_the": [
    " school",
    " kitchen",
    " park",
    " g"
  ],
  "samples": {
    "seed": 1,
    "sampled": "mia has the kitchen. at the book. mia.\n y s sch� bo sam has a red cup.\now cue.\ntrez cup.",
    "greedy": "mia has the school, sam gives the school, sam gives the school, sam gives the school, sam gives the school, sam gives"
  }
}

Failure cases

Each case has a named test in test_baseline.py.

The complete lab includes these tests and their fixtures. Run the 15 standard-library tests:

python3 -m unittest -v test_baseline.py

"""Tests for baseline.py. Run: python3 -m unittest -v test_baseline.py"""
import math
import json
import random
import subprocess
import sys
import unittest
from pathlib import Path

import baseline as b
from bpe import Tokenizer
from inspect_split import overlap_counts

HERE = Path(__file__).parent
TOK = Tokenizer.load(HERE / "tokenizer.json")
V = len(TOK.vocab)
TRAIN, DEV, TEST = b.split_stories((HERE / "corpus.txt").read_text())
TR = b.encode_stories(TOK, TRAIN)


class Recorded(unittest.TestCase):
    def test_recorded_trace(self):
        result = subprocess.run([sys.executable, "trace_baseline.py", "--check"], cwd=HERE)
        self.assertEqual(result.returncode, 0)

    def test_tokenizer_matches_its_corpus(self):
        self.assertEqual(Tokenizer.train((HERE / "corpus.txt").read_text(), 96).vocab, TOK.vocab)


class Mechanism(unittest.TestCase):
    def test_every_model_gives_a_distribution(self):
        for model in (b.Uniform(V), b.Unigram(TR, V), b.Bigram(TR, V, 0.01)):
            for prev in (0, TOK.ids[" the"], TOK.ids["."]):
                with self.subTest(model=type(model).__name__, prev=prev):
                    p = model.probs(prev)
                    self.assertEqual(len(p), V)
                    self.assertAlmostEqual(math.fsum(p), 1.0, places=12)

    def test_uniform_loss_is_log_vocab(self):
        loss, _ = b.cross_entropy(b.Uniform(V), b.encode_stories(TOK, DEV))
        self.assertAlmostEqual(loss, math.log(V), places=12)

    def test_bigram_beats_unigram_beats_uniform_on_test(self):
        te = b.encode_stories(TOK, TEST)
        losses = [b.cross_entropy(m, te)[0] for m in (b.Bigram(TR, V, 0.001), b.Unigram(TR, V), b.Uniform(V))]
        self.assertEqual(losses, sorted(losses))

    def test_split_rows_reconstruct_the_corpus(self):
        self.assertEqual((len(TRAIN), len(DEV), len(TEST)), (280, 60, 60))
        self.assertEqual(TRAIN + DEV + TEST, [s + "\n" for s in (HERE / "corpus.txt").read_text().split("\n") if s])

    def test_split_inspection_counts_rows_not_distinct_texts(self):
        stories = [f"story {i}" for i in range(400)]
        stories[280] = stories[281] = stories[0]
        stories[340] = stories[342] = stories[0]
        stories[341] = stories[282]
        counts = overlap_counts("\n".join(stories))
        self.assertEqual(counts["dev_rows_matching_training_text"], 2)
        self.assertEqual(counts["test_rows_matching_training_text"], 2)
        self.assertEqual(counts["test_rows_matching_train_or_dev_text"], 3)

    def test_bundled_split_overlap_is_reported(self):
        counts = overlap_counts((HERE / "corpus.txt").read_text())
        self.assertEqual(counts, json.loads((HERE / "split-overlap.json").read_text()))
        self.assertEqual(counts["dev_rows_matching_training_text"], 5)
        self.assertEqual(counts["test_rows_matching_training_text"], 5)
        self.assertEqual(counts["test_rows_matching_train_or_dev_text"], 8)

    def test_unseen_pair_gets_smoothed_probability(self):
        model = b.Bigram(TR, V, 0.01)
        prev = TOK.ids[" the"]
        unseen = TOK.ids["<unk>"]
        self.assertEqual(model.counts[prev][unseen], 0)
        self.assertAlmostEqual(model.probs(prev)[unseen], 0.01 / (model.totals[prev] + 0.01 * V))


class FailureCases(unittest.TestCase):
    def test_bigram_cannot_recall_color(self):
        result = b.recall_check(TOK, b.Bigram(TR, V, 0.001), TEST, "color")
        self.assertEqual(result["correct"], 0)
        self.assertEqual(list(result["guesses"]), [" school"])

    def test_greedy_decoding_repeats(self):
        text = b.sample(b.Bigram(TR, V, 0.001), TOK, "mia has", 25, random.Random(1), greedy=True)
        self.assertGreaterEqual(text.count("sam gives the school,"), 4)

    def test_sampling_can_emit_unknown_token(self):
        text = b.sample(b.Bigram(TR, V, 0.001), TOK, "mia has", 40, random.Random(1))
        self.assertIn("�", text)

    def test_zero_smoothing_rejected(self):
        with self.assertRaisesRegex(ValueError, "positive finite"):
            b.Bigram(TR, V, 0)

    def test_zero_probability_is_reported(self):
        class Never:
            def probs(self, prev):
                return [0.0] * V
        with self.assertRaisesRegex(ValueError, "probability 0.0 to an observed token"):
            b.cross_entropy(Never(), TR[:1])

    def test_too_few_stories(self):
        with self.assertRaisesRegex(ValueError, "need more than 340 stories"):
            b.split_stories("a.\nb.\n")


if __name__ == "__main__":
    unittest.main()

test_every_model_gives_a_distribution fails if a smoothing denominator is wrong. test_bigram_beats_unigram_beats_uniform_on_test fails if the bigram counts pairs in the wrong direction.

Deployment

There is nothing to deploy. The bigram’s counts are a table of 96 × 96 numbers, rebuilt from the training stories on every run. Keep baseline-trace.json as the reference result that later builds compare against, together with the tokenizer and split it used.

Scaling

This section is unmeasured. The count table grows with the square of the vocabulary, and a model that looks n tokens back needs a table indexed by n tokens of context, most of which never occur. That growth is why count-based models look only a few tokens back, and why the next builds learn a fixed-size model that can read the whole story.

Key takeaways

• Score every next-token model on the same held-out positions with mean cross-entropy; perplexity is e raised to that mean.
• Choose settings such as k on development data and score the test stories once.
• An average loss can look strong while a specific, context-dependent prediction fails. Measure that prediction directly.

Best practices

• Keep a simple baseline’s test loss as the first number a new model must beat.
• Split by whole document and check for duplicate text across splits. A split by row position alone does not prevent leakage.
• Reject a probability of 0 for an observed token before you take its log.

References

Sources checked October 20, 2026.

• Jurafsky and Martin, Speech and Language Processing, 3rd edition draft of August 19, 2026, chapter 3: N-gram Language Models. It covers the bigram approximation (section 3.1), training, development and test sets (3.2), perplexity (3.3, 3.7), sampling (3.4) and add-one and add-k smoothing with k chosen on a development set (3.6).
• PyTorch 2.14: torch.nn.CrossEntropyLoss: the same mean cross-entropy, computed from logits by a library.