Skip to main content

Predicting the next token

Count context, choose on validation data, and test once

Learn how a language model assigns probabilities to the next token and how to evaluate an inspectable count model.

Nia is building an autocomplete demonstration for an archive. After a reader types we, the tool must decide which token could come next.

A language model assigns probabilities to token sequences. In next-token prediction, the input is the context already seen and the output is a probability for each possible next token. This is one language-model training objective. It is not a definition of every language model or every text-generation system.

TipWhat you will learn

This lesson shows how to:

• keep training, validation, and test speeches separate;
• build next-token rows without crossing paragraph boundaries;
• estimate unigram, bigram, and trigram count models;
• compare greedy choice with probability-weighted sampling;
• select context length and interpolation strength with validation perplexity; and
• explain what held-out perplexity does not measure.

Load the split and study results

The inaugural corpus contains 60 public-domain speeches. The committed study uses 37 speeches for training, 10 for model selection, and 13 for the final test. A speech belongs to one split only. The split is stratified by era so both eras appear in every part. The split holds out speeches, not speakers: presidents with more than one inaugural address can appear in multiple splits.

library(digest)
library(dplyr)
library(tidyr)
library(tibble)
library(stringr)
library(purrr)
library(readr)
library(knitr)

source("R/inaugural-corpus.R")

split_path <- "data/inaugural/next-token-splits.csv"
validation_path <- "data/inaugural/next-token-validation.csv"
test_path <- "data/inaugural/next-token-test.csv"
metadata_path <- "data/inaugural/next-token-study-metadata.csv"

split_assignments <- read_csv(
  split_path,
  na = character(),
  col_types = cols(
    speech_id = col_character(),
    era = col_character(),
    split = col_character()
  )
)
validation_results <- read_csv(
  validation_path,
  na = c("", "NA"),
  col_types = cols(
    split = col_character(),
    model = col_character(),
    context_tokens = col_integer(),
    interpolation_strength = col_double(),
    next_token_rows = col_integer(),
    oov_target_share = col_double(),
    seen_context_share = col_double(),
    mean_log_loss = col_double(),
    perplexity = col_double(),
    selected = col_logical(),
    selection_reason = col_character()
  )
)
test_results <- read_csv(
  test_path,
  na = c("", "NA"),
  col_types = cols(
    split = col_character(),
    model = col_character(),
    context_tokens = col_integer(),
    interpolation_strength = col_double(),
    next_token_rows = col_integer(),
    oov_target_share = col_double(),
    seen_context_share = col_double(),
    mean_log_loss = col_double(),
    perplexity = col_double(),
    top_1_accuracy = col_double(),
    role = col_character()
  )
)
test_by_speech <- read_csv(
  "data/inaugural/next-token-test-by-speech.csv",
  na = character(),
  col_types = cols(
    speech_id = col_character(),
    next_token_rows = col_integer(),
    model = col_character(),
    perplexity = col_double()
  )
)
study_metadata <- read_csv(
  metadata_path,
  na = character(),
  col_types = cols(
    artifact = col_character(),
    description = col_character(),
    settings = col_character(),
    source = col_character(),
    license = col_character(),
    built_on = col_character(),
    fingerprint = col_character()
  )
)

split_counts <- split_assignments |>
  mutate(
    split = factor(
      split,
      levels = c("training", "validation", "test")
    ),
    era = factor(
      era,
      levels = c("before 1900", "1900 or later")
    )
  ) |>
  count(split, era, name = "speeches") |>
  arrange(split, era)

kable(
  split_counts,
  col.names = c("Split", "Era", "Speeches"),
  caption = "Speech-level assignments used for next-token modeling",
  row.names = FALSE
)

Speech-level assignments used for next-token modeling

Split	Era	Speeches
training	before 1900	17
training	1900 or later	20
validation	before 1900	5
validation	1900 or later	5
test	before 1900	6
test	1900 or later	7

The validation set is used to choose among candidate models. The test speeches stay closed until that choice is fixed. Calling both sets “held out” would hide their different jobs.

Build paragraph-bounded next-token rows

The study uses an explicit word-like pattern so token counts do not depend on the rendering machine’s word-boundary library. It then keeps lowercase ASCII alphabetic tokens. The 2,499 most common training token types get their own vocabulary entries; every other token maps to <unk>, giving a 2,500-token model vocabulary.

Each prediction row keeps its paragraph and speech ID. previous_1 is the token immediately before the target. previous_2 is one more step back. The first two tokens in each paragraph have no full two-token context, so they are not evaluation rows. No transition joins the end of one paragraph to the beginning of another, and this model does not learn how to start a paragraph.

paragraphs <- inaugural_paragraphs()
token_pattern <- paste0(
  "(?:[a-z]\\.){2,}|",
  "[a-z0-9]+(?:['\\x{2019}][a-z0-9]+)*"
)

training_tokens <- paragraphs |>
  select(paragraph_id, speech_id, paragraph) |>
  inner_join(
    split_assignments |>
      filter(split == "training") |>
      select(speech_id),
    by = join_by(speech_id)
  ) |>
  mutate(
    token = str_extract_all(
      str_to_lower(paragraph, locale = "en"),
      regex(token_pattern)
    )
  ) |>
  select(-paragraph) |>
  unnest_longer(token) |>
  filter(str_detect(token, "^[a-z]+$")) |>
  mutate(position = row_number(), .by = paragraph_id)

training_vocabulary <- training_tokens |>
  count(token, sort = TRUE) |>
  slice_head(n = 2499L) |>
  pull(token)
model_vocabulary <- c(training_vocabulary, "<unk>")

training_next_tokens <- training_tokens |>
  mutate(
    model_token = if_else(
      token %in% training_vocabulary,
      token,
      "<unk>"
    )
  ) |>
  arrange(paragraph_id, position) |>
  mutate(
    previous_1 = lag(model_token, 1L),
    previous_2 = lag(model_token, 2L),
    .by = paragraph_id
  ) |>
  filter(position >= 3L) |>
  transmute(
    paragraph_id,
    speech_id,
    previous_2,
    previous_1,
    next_token = model_token
  )

training_summary <- tibble(
  measure = c(
    "training speeches",
    "training paragraphs",
    "training tokens",
    "next-token rows",
    "model vocabulary"
  ),
  value = c(
    n_distinct(training_next_tokens$speech_id),
    n_distinct(training_tokens$paragraph_id),
    nrow(training_tokens),
    nrow(training_next_tokens),
    length(model_vocabulary)
  )
)

kable(
  training_summary,
  col.names = c("Training item", "Value"),
  caption = "Training data for the inspectable count model",
  row.names = FALSE
)

Training data for the inspectable count model

Training item	Value
training speeches	37
training paragraphs	867
training tokens	83458
next-token rows	81724
model vocabulary	2500

Count what followed each token

A unigram ignores context and uses overall token frequency. A bigram uses one previous token. A trigram uses two. More context can help, but each longer context appears less often.

Raw counts assign probability zero to a next token that never followed the context in training. The study avoids zero probabilities through interpolation: every estimate mixes the longer-context counts with the shorter model. The interpolation strength controls how much weight the shorter model receives, so a larger value leans more on that model. Six values are compared on validation speeches; the selected trigram uses 100.

selected_validation <- validation_results |>
  filter(selected)
interpolation_strength <- selected_validation$interpolation_strength

unigram_counts <- training_tokens |>
  mutate(
    model_token = if_else(
      token %in% training_vocabulary,
      token,
      "<unk>"
    )
  ) |>
  count(model_token, name = "unigram_count")
training_token_total <- sum(unigram_counts$unigram_count)

unigram_probabilities <- tibble(
  model_token = model_vocabulary
) |>
  left_join(unigram_counts, by = join_by(model_token)) |>
  mutate(
    unigram_count = replace_na(unigram_count, 0L),
    unigram_probability = (unigram_count + 1) /
      (training_token_total + length(model_vocabulary))
  )

bigram_counts <- training_next_tokens |>
  count(previous_1, next_token, name = "bigram_count")
trigram_counts <- training_next_tokens |>
  count(
    previous_2,
    previous_1,
    next_token,
    name = "trigram_count"
  )

predict_after_one <- function(previous_token) {
  context_counts <- bigram_counts |>
    filter(previous_1 == previous_token)

  tibble(next_token = model_vocabulary) |>
    left_join(
      unigram_probabilities,
      by = join_by(next_token == model_token)
    ) |>
    left_join(context_counts, by = join_by(next_token)) |>
    mutate(
      bigram_count = replace_na(bigram_count, 0L),
      context_count = sum(bigram_count),
      bigram_probability = (
        bigram_count +
          interpolation_strength * unigram_probability
      ) / (
        context_count +
          interpolation_strength
      )
    ) |>
    arrange(desc(bigram_probability), next_token)
}

predict_after_two <- function(previous_2_token, previous_1_token) {
  context_counts <- trigram_counts |>
    filter(
      previous_2 == previous_2_token,
      previous_1 == previous_1_token
    )

  predict_after_one(previous_1_token) |>
    left_join(context_counts, by = join_by(next_token)) |>
    mutate(
      trigram_count = replace_na(trigram_count, 0L),
      context_count = sum(trigram_count),
      probability = (
        trigram_count +
          interpolation_strength * bigram_probability
      ) / (
        context_count +
          interpolation_strength
      ),
      context = str_c(
        previous_2_token,
        previous_1_token,
        sep = " "
      )
    ) |>
    arrange(desc(probability), next_token)
}

context_predictions <- bind_rows(
  predict_after_two("we", "must") |>
    slice_head(n = 5),
  predict_after_two("the", "united") |>
    slice_head(n = 5)
) |>
  select(context, next_token, trigram_count, probability)

kable(
  context_predictions |>
    mutate(probability = round(probability, 4)),
  col.names = c("Previous two tokens", "Possible next token", "Training count", "Probability"),
  caption = "Highest-probability next tokens from the selected trigram",
  row.names = FALSE
)

Highest-probability next tokens from the selected trigram

Previous two tokens	Possible next token	Training count	Probability
we must	be	9	0.1725
we must		8	0.0864
we must	not	5	0.0474
we must	do	5	0.0380
we must	show	4	0.0288
the united	states	112	0.7189
the united	nations	6	0.0386
the united		2	0.0275
the united	and	1	0.0172
the united	the	0	0.0139

After the united, the model assigns about 0.72 probability to states. That reflects a repeated phrase in US inaugural addresses. It does not show a general English rule. After we must, several continuations remain plausible.

The <unk> row stands for any token outside the training vocabulary. Predicting that bucket does not identify the original word.

Choose or sample a next token

The probability table does not decide how text is generated. Greedy decoding takes the highest-probability token. Sampling draws according to the full distribution. The seed below makes five draws reproducible.

greedy_next <- we_must_distribution$next_token[[1]]

set.seed(6303)
sampled_next <- sample(
  we_must_distribution$next_token,
  size = 5,
  replace = TRUE,
  prob = we_must_distribution$probability
)

choice_table <- tibble(
  method = c("greedy", "five seeded samples"),
  output = c(
    greedy_next,
    str_c(sampled_next, collapse = ", ")
  )
)

kable(
  choice_table,
  col.names = c("Choice rule", "Next-token output after 'we must'"),
  caption = "One probability distribution can support different choice rules",
  row.names = FALSE
)

One probability distribution can support different choice rules

Choice rule	Next-token output after ‘we must’
greedy	be
five seeded samples	be, races, maintain, strive, change

Greedy choice repeats the most common continuation. Sampling can select less common tokens and will change when the seed changes. Neither rule checks whether a continuation is factual, useful, or appropriate for a reader.

Select context length and smoothing before testing

The builder compares six interpolation strengths for both the bigram and trigram, plus the unigram baseline, on validation speeches. Log loss penalizes probability assigned away from the observed next token. Perplexity is the exponentiated mean log loss; lower is better when the tokenizer, vocabulary, and evaluation rows are identical.

validation_display <- validation_results |>
  filter(
    model == "unigram" |
      (model == "bigram" & interpolation_strength == 100) |
      (
        model == "trigram" &
          interpolation_strength %in% c(10, 100, 1000)
      )
  ) |>
  arrange(context_tokens, interpolation_strength) |>
  transmute(
    model,
    context_tokens,
    interpolation_strength = if_else(
      is.na(interpolation_strength),
      "not applicable",
      as.character(interpolation_strength)
    ),
    seen_context = if_else(
      is.na(seen_context_share),
      "not applicable",
      sprintf("%.1f%%", 100 * seen_context_share)
    ),
    perplexity = round(perplexity, 1),
    selected = if_else(selected, "yes", "no")
  )

kable(
  validation_display,
  col.names = c(
    "Candidate",
    "Context tokens",
    "Interpolation strength",
    "Validation contexts seen in training",
    "Validation perplexity",
    "Selected"
  ),
  caption = "Validation settings around the selected trigram",
  row.names = FALSE
)

Validation settings around the selected trigram

Candidate	Context tokens	Interpolation strength	Validation contexts seen in training	Validation perplexity	Selected
unigram	0	not applicable	not applicable	238.0	no
bigram	1	100	100.0%	134.8	no
trigram	2	10	74.9%	194.4	no
trigram	2	100	74.9%	133.3	yes
trigram	2	1000	74.9%	143.6	no

The trigram with interpolation strength 100 has the lowest validation perplexity, 133.3. The best bigram is close at 134.8. The tested strengths on both sides of 100 score worse, so the selected value is not a grid boundary. Only 74.9 percent of the trigram’s two-token validation contexts appeared in training. When a context is unseen, its estimate relies entirely on the bigram and then the unigram. The bigram’s 100 percent seen-context rate follows from mapping every token into the training vocabulary, including <unk>; it is not evidence that the model generalizes to every word.

Open the test speeches once

After selecting the trigram and its interpolation strength, the study scores it, the predeclared unigram baseline, and a uniform 2,500-token floor on the 13 test speeches.

test_display <- test_results |>
  transmute(
    model,
    role,
    next_token_rows,
    oov_targets = sprintf("%.1f%%", 100 * oov_target_share),
    perplexity = round(perplexity, 1),
    top_1_accuracy = if_else(
      is.na(top_1_accuracy),
      "not applicable",
      sprintf("%.1f%%", 100 * top_1_accuracy)
    )
  )

perplexity_reduction <- 1 -
  test_results$perplexity[test_results$model == "trigram"] /
    test_results$perplexity[test_results$model == "unigram"]

test_comparison_by_speech <- test_by_speech |>
  select(speech_id, model, perplexity) |>
  pivot_wider(
    names_from = model,
    values_from = perplexity
  )

speech_people <- paragraphs |>
  distinct(speech_id, president)
training_presidents <- split_assignments |>
  filter(split == "training") |>
  inner_join(speech_people, by = join_by(speech_id)) |>
  pull(president)
test_speaker_overlap <- split_assignments |>
  filter(split == "test") |>
  inner_join(speech_people, by = join_by(speech_id)) |>
  mutate(
    president_seen_in_training = president %in% training_presidents
  )

paired_test <- test_comparison_by_speech |>
  summarise(
    test_speeches = n(),
    trigram_wins = sum(trigram < unigram),
    ties = sum(trigram == unigram),
    trigram_losses = sum(trigram > unigram)
  )

speaker_overlap_test <- test_speaker_overlap |>
  summarise(
    repeated_president_speeches = sum(president_seen_in_training),
    other_speeches = sum(!president_seen_in_training)
  )
non_overlap_test <- test_comparison_by_speech |>
  inner_join(
    test_speaker_overlap |>
      select(speech_id, president_seen_in_training),
    by = join_by(speech_id)
  ) |>
  filter(!president_seen_in_training) |>
  summarise(
    speeches = n(),
    trigram_wins = sum(trigram < unigram)
  )

kable(
  test_display,
  col.names = c(
    "Model",
    "Role",
    "Test next-token rows",
    "Targets mapped to <unk>",
    "Test perplexity",
    "Top-1 accuracy"
  ),
  caption = "Untouched-test result for the selected model and baseline",
  row.names = FALSE
)

Untouched-test result for the selected model and baseline

Model	Role	Test next-token rows	Targets mapped to	Test perplexity	Top-1 accuracy
uniform	uniform vocabulary floor	24337	11.7%	2500.0	not applicable
unigram	predeclared frequency baseline	24337	11.7%	240.7	11.7%
trigram	validation-selected model	24337	11.7%	145.3	19.2%

The uniform model has perplexity 2,500 because it gives every vocabulary item the same probability. The selected trigram reduces test perplexity by about 39.6% relative to the unigram baseline and wins the paired comparison in all 13 test speeches.

Its highest-probability guess is right on 19.2 percent of test rows. The unigram scores 11.7 percent because its most likely token is <unk>, the bucket that collects many rare training words. A correct unknown-bucket guess does not recover the original word. The trigram also leans on this bucket: roughly one third of its correct guesses are unknown-bucket matches rather than recovered words.

Eight of the 13 test speeches were delivered by a president who also appears in training. The test is speech-held-out, not speaker-held-out; repeated presidents can carry phrasing across splits. The trigram still beats the unigram on all five test speeches whose president is absent from training.

Perplexity would change under another vocabulary, tokenizer, boundary rule, or test set. The values here cannot be compared with the LDA perplexity in Lesson 55 or a published model score from another tokenization.

The result does not show that the model understands a speech, writes a good report, or produces facts. The next lesson uses these same measurements to show how a report can stay tied to evidence.

What to remember

• Next-token prediction assigns probabilities after a context.
• Paragraph boundaries prevent invented transitions between source records.
• Validation data select context length and interpolation strength; test data check the fixed choice.
• The trigram and interpolation strength are selected together on validation data.
• The selected trigram beats the unigram on all 13 test speeches here, but its top choice is still wrong on about four out of five rows.
• Perplexity is comparable only under the same tokens, vocabulary, and test rows.
• A likely next token is not necessarily factual or useful.

Sources

• Jurafsky and Martin, Speech and Language Processing, chapter on n-gram language models
• Bengio et al., A Neural Probabilistic Language Model
• Chen and Goodman, An Empirical Study of Smoothing Techniques for Language Modeling
• Hugging Face documentation on perplexity for fixed-length models
• Meister and Cotterell, Language Model Evaluation Beyond Perplexity
• quanteda inaugural-address corpus documentation