Build text generation with sampling controls

Problem

When you need this. The trained model predicts one next token. Turn those predictions into a text continuation by repeatedly choosing a token, appending it and checking when to stop.

The model so far only scores the next token. A language model is used to write: given an opening, it adds one token, reads its own output and adds the next. This build writes that loop for the model saved in best.json by the previous build, and the controls that decide which token to add.

The contract is generate(params, tok, prompt, max_new_tokens, temperature, top_k, seed). It returns the text, each chosen token with the probability it had, and the reason generation stopped. A temperature of 0 means greedy choice. The function rejects an empty prompt, a prompt longer than the model’s position table, a negative or non-finite temperature or one above 0 but below 0.01, and a top-k outside 1 to the vocabulary size.

Success has three checks. The same seed gives the same text, and a different seed can give a different one. Each control does what its definition says, which the tests check on real model scores. And each setting is measured on the 60 development stories: given a story’s first sentence, does the generated story have the right shape, and does its last sentence keep the first sentence’s color and object?

Concepts

Inference uses trained weights to compute predictions without updating them. The forward pass produces next-token scores; generation repeats that inference while choosing tokens and deciding when to stop. Training uses a backward pass to change weights; this loop does not. Each control below is a few lines of code applied to the last position’s scores. Generation settings and output checks provides optional background on these choices:

Greedy choice

Take the token with the highest score. The same prompt always gives the same text, and no seed is needed.

Temperature

Divide every score by a number T before softmax, then draw a token at random from the resulting probabilities. T below 1 sharpens the distribution toward the top token; T above 1 flattens it. A very small T can make the divided scores overflow to infinity, so generate() refuses any T above 0 but below 0.01, a chosen floor; T = 0 means greedy choice.

Top-k

Before softmax, keep only the k highest scores and set the rest to −∞, so only those k tokens can be drawn. When scores tie, the lower token ID is kept, so exactly k tokens remain. Removing the long tail of unlikely tokens is the point: drawing from the full distribution can pick tokens that make no sense in context.

Stop condition

Generation ends at a newline, the end of a story in this corpus; at max_new_tokens; or when the text fills all of the model’s positions.

Random draws use Python’s random.Random(seed), so a run can be repeated exactly. The model has no memory between steps: each step runs it again on the whole text so far. The next build, Build a KV cache for LLM generation, measures what that repeated work costs.

Architecture

The prompt, settings and seed enter Encode and validate. At each step, Positions full? checks whether the model can read the current prefix. On no, Score the prefix runs model.forward() and Choose one token applies greedy selection or temperature and top-k sampling. Append token extends the prefix. A newline or token limit returns the generated text; otherwise the loop returns to the position check. Invalid inputs take the separate ValueError branch. The following run compares the sampling settings on the development stories.

generate.py holds next_token_probs(), which turns one row of scores into the probabilities the sampler draws from, generate(), the loop, and follows_story(), which checks a generated story against its first sentence. It loads model.py and the verified best.json through checkpoint.load(), both copied from the previous build. trace_generate.py runs every setting on the development openings and writes generation-trace.json. It does not use the test stories.

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-text-generation-with-sampling-controls/; keep its Python modules, tests and fixture files together. The bundled README includes setup commands for your platform.

"""Generate text from the trained one-block model with greedy, temperature and top-k selection.

Each step runs the whole model on the text so far, takes the scores at the last
position and chooses one token. Generation stops at a newline (the end of a
story), at max_new_tokens, or when the model's position table is full.
Python 3.10+, standard library only.
"""
import math
import random

import model as m

MIN_TEMPERATURE = 0.01  # a chosen floor: a tiny T can overflow the divided scores; use 0 for greedy


def check_settings(temperature, top_k, vocab_size):
    if isinstance(temperature, bool) or not isinstance(temperature, (int, float)) \
            or not math.isfinite(temperature) or not (temperature == 0 or temperature >= MIN_TEMPERATURE):
        raise ValueError(f"temperature must be 0 or a finite number >= {MIN_TEMPERATURE}, got {temperature!r}")
    if top_k is not None and (isinstance(top_k, bool) or not isinstance(top_k, int) or not 1 <= top_k <= vocab_size):
        raise ValueError(f"top_k must be an integer from 1 to {vocab_size}, got {top_k!r}")


def next_token_probs(logits, temperature, top_k):
    """Turn one row of scores into the probabilities the sampler draws from."""
    if temperature == 0:  # greedy: all probability on the highest score, ties to the lowest ID
        best = max(range(len(logits)), key=lambda i: (logits[i], -i))
        return [1.0 if i == best else 0.0 for i in range(len(logits))]
    scaled = [v / temperature for v in logits]
    if top_k is not None:
        # The k highest scores; among equal scores the lower token ID is kept.
        kept = set(sorted(range(len(scaled)), key=lambda i: (-scaled[i], i))[:top_k])
        scaled = [v if i in kept else -math.inf for i, v in enumerate(scaled)]
    return m.softmax_row(scaled)


def generate(p, tok, prompt, max_new_tokens=60, temperature=0.0, top_k=None, seed=0):
    """Return the generated text, the chosen tokens with their probabilities, and why it stopped."""
    check_settings(temperature, top_k, len(p["tok_emb"]))
    ids = tok.encode(prompt)
    if not ids:
        raise ValueError("empty prompt: the model needs at least one token to continue")
    max_len = len(p["pos_emb"])
    if len(ids) > max_len:
        raise ValueError(f"prompt of {len(ids)} tokens is longer than the {max_len} positions the model has")
    rng = random.Random(seed)
    steps, reason = [], "max_new_tokens"
    newline = tok.ids["\n"]
    for _ in range(max_new_tokens):
        if len(ids) == max_len:
            reason = "context full"
            break
        logits, _ = m.forward(p, ids)
        probs = next_token_probs(logits[-1], temperature, top_k)
        choice = rng.choices(range(len(probs)), weights=probs)[0] if temperature > 0 else probs.index(1.0)
        ids.append(choice)
        steps.append({"token": tok.vocab[choice], "probability": round(probs[choice], 4)})
        if choice == newline:
            reason = "end of story"
            break
    return {"text": tok.decode(ids), "steps": steps, "stopped": reason}


def follows_story(prompt_sentence, generated):
    """Check a generated story against its first sentence: same giver, color and object at the end."""
    words = prompt_sentence.rstrip(".").split()  # name has a color object
    giver, color, thing = words[0], words[3], words[4]
    lines = generated.split("\n")[0].split(". ")
    if len(lines) != 4 or not lines[3].startswith("now "):
        return {"well_formed": False, "keeps_color_and_object": False}
    last = lines[3].rstrip(".").split()
    well_formed = (lines[1].startswith(f"{giver} goes to the ") and lines[2].startswith("at the ")
                   and len(last) == 6 and last[2:4] == ["has", "the"])
    return {"well_formed": well_formed, "keeps_color_and_object": well_formed and last[4:6] == [color, thing]}

next_token_probs() breaks greedy ties toward the lowest token ID, so greedy output never depends on the seed. Top-k sorts by score, then by token ID, and keeps the first k, so it keeps exactly k tokens and never drops a higher score when lower ones tie. generate() checks the stop conditions before and after each step, so it never asks the model for a position it does not have.

"""Generate from the best checkpoint with several settings and record the results.

Run: python3 trace_generate.py          (print the trace)
     python3 trace_generate.py --write  (rewrite generation-trace.json and generation-time.json)
     python3 trace_generate.py --check  (exit 1 if generation-trace.json differs)
generation-time.json holds the time of the --write run and is not compared.
The settings are compared on the development stories; the test stories are not used here.
"""
import json
import platform
import sys
import time
from pathlib import Path

import checkpoint as ck
import generate as g
import model as m
import train as t
from bpe import Tokenizer

HERE = Path(__file__).parent
SETTINGS = [
    ("greedy", {"temperature": 0.0}),
    ("temperature 0.5", {"temperature": 0.5}),
    ("temperature 1.0", {"temperature": 1.0}),
    ("temperature 1.5", {"temperature": 1.5}),
    ("temperature 1.0, top-k 3", {"temperature": 1.0, "top_k": 3}),
]
SEED_PROMPT = "zoe has a blue kite."
COLOR_PROMPT = "mia has a red ball. mia goes to the park. at the park, mia gives the ball to leo. now leo has the"


def top_choices(probs, tok, n=5):
    order = sorted(range(len(probs)), key=lambda i: (-probs[i], i))[:n]
    return [{"token": tok.vocab[i], "probability": round(probs[i], 4)} for i in order]


def trace():
    tok = Tokenizer.load(HERE / "tokenizer.json")
    best = ck.load(HERE / "best.json", t.CONFIG, ck.input_signature(tok, (HERE / "corpus.txt").read_text()))
    p = best["params"]
    _, dev, _ = t.split_stories((HERE / "corpus.txt").read_text())
    table = []
    for name, kw in SETTINGS:
        counts = {"well_formed": 0, "keeps_color_and_object": 0}
        stops = {"end of story": 0, "max_new_tokens": 0, "context full": 0}
        for i, story in enumerate(dev):
            first = story.split(". ")[0] + "."
            result = g.generate(p, tok, first, seed=i, **kw)
            check = g.follows_story(first, result["text"])
            counts = {k: counts[k] + check[k] for k in counts}
            stops[result["stopped"]] += 1
        table.append({"setting": name, **kw, "stories": len(dev), **counts, "stopped": stops})
    logits, _ = m.forward(p, tok.encode(COLOR_PROMPT))
    distributions = {str(temp): top_choices(g.next_token_probs(logits[-1], temp, None), tok)
                     for temp in (0.5, 1.0, 2.0)}
    # Two stories cut to 62 tokens: only two positions remain in the 64-row table.
    long_prompt = tok.decode(tok.encode(dev[0] + dev[1])[:62])
    seeds = {seed: g.generate(p, tok, SEED_PROMPT, temperature=1.0, seed=seed) for seed in (1, 2)}
    step = next(i for i, (a, b) in enumerate(zip(seeds[1]["steps"], seeds[2]["steps"])) if a != b)
    return {
        "checkpoint": {"step": best["step"], "dev_loss": best["dev_loss"]},
        "greedy_example": g.generate(p, tok, "mia has a red ball."),
        "unseen_words_example": g.generate(p, tok, "maya has a purple umbrella."),
        "seed_prompt": SEED_PROMPT,
        "seeds": [{"seed": seed, "text": g.generate(p, tok, SEED_PROMPT, temperature=1.0, seed=seed)["text"]}
                  for seed in (1, 2, 1)],
        "seed_divergence": {"first_different_step": step + 1,
                            "same_tokens_before": [s["token"] for s in seeds[1]["steps"][:step]],
                            **{f"seed_{k}": seeds[k]["steps"][step] for k in (1, 2)}},
        "color_prompt": COLOR_PROMPT,
        "next_token_at_color": distributions,
        "context_full_example": {"prompt_tokens": len(tok.encode(long_prompt)),
                                 **g.generate(p, tok, long_prompt, max_new_tokens=60)},
        "dev_story_table": table,
    }


if __name__ == "__main__":
    path = HERE / "generation-trace.json"
    start = time.perf_counter()
    text = json.dumps(trace(), indent=1, ensure_ascii=False) + "\n"
    timing = {"seconds_for_trace": round(time.perf_counter() - start, 1), "python": platform.python_version(),
              "machine": platform.machine(), "system": platform.system()}
    if "--write" in sys.argv:
        path.write_text(text)
        (HERE / "generation-time.json").write_text(json.dumps(timing, indent=1) + "\n")
    elif "--check" in sys.argv:
        sys.exit(0 if path.read_text() == text else "generation-trace.json is out of date")
    else:
        print(text, end="")

Run the trace:

python3 trace_generate.py

python3 trace_generate.py --check exits with an error if your output differs from the recorded generation-trace.json. The recorded run took 32.7 seconds, measured on one machine; your numbers will differ. That is one measurement, kept in generation-time.json, which --check does not compare.

Generate one story

Greedy generation from mia has a red ball. gives:

mia has a red ball. mia goes to the kitchen. at the kitchen, mia gives the ball to ben. now ben has the yellow ball.

It stopped at "end of story". The story has the corpus’s shape, and the object stays the same, but the last sentence names a different color from the first. The best checkpoint makes that mistake on some stories, as the previous build’s recall check showed.

A prompt with words the corpus never contained, maya has a purple umbrella., gives text that does not follow the story pattern:

maya has a purple umbrella. at the cup to mia has the blue cup.

Watch temperature reshape one choice

The trace runs the model on mia has a red ball. mia goes to the park. at the park, mia gives the ball to leo. now leo has the, where the right next token is the color from the first sentence, “ red”. The five most likely tokens at three temperatures:

On this prompt the model puts more probability on " yellow" than on the right color. A lower temperature makes the wrong color more likely, 0.6052 at T = 0.5; a higher one spreads probability toward tokens that are not colors at all. Temperature changes how sure the choice is. It cannot change which token the model ranks first.

Compare the settings on development openings

For each of the 60 development stories, the trace gives the model the first sentence and generates until the story ends. follows_story() then checks two things: the story has the corpus’s four-sentence shape with the same child, and its last sentence repeats the first sentence’s color and object.

The clear effect is at T = 1.5: only 26 of 60 stories are well-formed, and 1 ran until the text filled every position. Between the other settings the differences are a few stories. Top-k with k = 3 at T = 1.0 gives 57 of 60 well-formed stories against 56 of 60 without it, one draw per story with one seed, which is too small a difference to call an improvement. Greedy choice ended 60 of 60 stories at a newline, and none reached max_new_tokens, where a story caught in a repeating loop would stop. This table compares settings; it does not select one, so the test stories stay unused.

Check your understanding

The trace generates from zoe has a blue kite. at T = 1.0 with seeds 1, 2, 1. Before you open the comparison, predict which of the three outputs are identical.

{
 "checkpoint": {
  "step": 450,
  "dev_loss": 0.2941
 },
 "greedy_example": {
  "text": "mia has a red ball. mia goes to the kitchen. at the kitchen, mia gives the ball to ben. now ben has the yellow ball.\n",
  "steps": [
   {
    "token": " mia",
    "probability": 1.0
   },
   {
    "token": " goes",
    "probability": 1.0
   },
   {
    "token": " to",
    "probability": 1.0
   },
   {
    "token": " the",
    "probability": 1.0
   },
   {
    "token": " kitchen",
    "probability": 1.0
   },
   {
    "token": ".",
    "probability": 1.0
   },
   {
    "token": " at",
    "probability": 1.0
   },
   {
    "token": " the",
    "probability": 1.0
   },
   {
    "token": " kitchen",
    "probability": 1.0
   },
   {
    "token": ",",
    "probability": 1.0
   },
   {
    "token": " mia",
    "probability": 1.0
   },
   {
    "token": " gives",
    "probability": 1.0
   },
   {
    "token": " the",
    "probability": 1.0
   },
   {
    "token": " ball",
    "probability": 1.0
   },
   {
    "token": " to",
    "probability": 1.0
   },
   {
    "token": " ben",
    "probability": 1.0
   },
   {
    "token": ".",
    "probability": 1.0
   },
   {
    "token": " now",
    "probability": 1.0
   },
   {
    "token": " ben",
    "probability": 1.0
   },
   {
    "token": " has",
    "probability": 1.0
   },
   {
    "token": " the",
    "probability": 1.0
   },
   {
    "token": " yellow",
    "probability": 1.0
   },
   {
    "token": " ball",
    "probability": 1.0
   },
   {
    "token": ".",
    "probability": 1.0
   },
   {
    "token": "\n",
    "probability": 1.0
   }
  ],
  "stopped": "end of story"
 },
 "unseen_words_example": {
  "text": "maya has a purple umbrella. at the cup to mia has the blue cup.\n",
  "steps": [
   {
    "token": " at",
    "probability": 1.0
   },
   {
    "token": " the",
    "probability": 1.0
   },
   {
    "token": " ",
    "probability": 1.0
   },
   {
    "token": "c",
    "probability": 1.0
   },
   {
    "token": "u",
    "probability": 1.0
   },
   {
    "token": "p",
    "probability": 1.0
   },
   {
    "token": " to",
    "probability": 1.0
   },
   {
    "token": " mia",
    "probability": 1.0
   },
   {
    "token": " has",
    "probability": 1.0
   },
   {
    "token": " the",
    "probability": 1.0
   },
   {
    "token": " b",
    "probability": 1.0
   },
   {
    "token": "l",
    "probability": 1.0
   },
   {
    "token": "u",
    "probability": 1.0
   },
   {
    "token": "e",
    "probability": 1.0
   },
   {
    "token": " ",
    "probability": 1.0
   },
   {
    "token": "c",
    "probability": 1.0
   },
   {
    "token": "u",
    "probability": 1.0
   },
   {
    "token": "p",
    "probability": 1.0
   },
   {
    "token": ".",
    "probability": 1.0
   },
   {
    "token": "\n",
    "probability": 1.0
   }
  ],
  "stopped": "end of story"
 },
 "seed_prompt": "zoe has a blue kite.",
 "seeds": [
  {
   "seed": 1,
   "text": "zoe has a blue kite. zoe goes to the park. at the park, zoe gives the bl mia. now leo.\n"
  },
  {
   "seed": 2,
   "text": "zoe has a blue kite. zoe goes to the park. at the park, zoe gives the kite to leo. now ava has the blue kite.\n"
  },
  {
   "seed": 1,
   "text": "zoe has a blue kite. zoe goes to the park. at the park, zoe gives the bl mia. now leo.\n"
  }
 ],
 "seed_divergence": {
  "first_different_step": 14,
  "same_tokens_before": [
   " zoe",
   " goes",
   " to",
   " the",
   " park",
   ".",
   " at",
   " the",
   " park",
   ",",
   " zoe",
   " gives",
   " the"
  ],
  "seed_1": {
   "token": " b",
   "probability": 0.0034
  },
  "seed_2": {
   "token": " kit",
   "probability": 0.9934
  }
 },
 "color_prompt": "mia has a red ball. mia goes to the park. at the park, mia gives the ball to leo. now leo has the",
 "next_token_at_color": {
  "0.5": [
   {
    "token": " yellow",
    "probability": 0.6052
   },
   {
    "token": " red",
    "probability": 0.2274
   },
   {
    "token": " green",
    "probability": 0.1673
   },
   {
    "token": " b",
    "probability": 0.0001
   },
   {
    "token": " book",
    "probability": 0.0
   }
  ],
  "1.0": [
   {
    "token": " yellow",
    "probability": 0.4628
   },
   {
    "token": " red",
    "probability": 0.2837
   },
   {
    "token": " green",
    "probability": 0.2433
   },
   {
    "token": " b",
    "probability": 0.0062
   },
   {
    "token": " book",
    "probability": 0.0015
   }
  ],
  "2.0": [
   {
    "token": " yellow",
    "probability": 0.3146
   },
   {
    "token": " red",
    "probability": 0.2463
   },
   {
    "token": " green",
    "probability": 0.2281
   },
   {
    "token": " b",
    "probability": 0.0364
   },
   {
    "token": " book",
    "probability": 0.0177
   }
  ]
 },
 "context_full_example": {
  "prompt_tokens": 62,
  "text": "mia has a red book. mia goes to the kitchen. at the kitchen, mia gives the book to ben. now ben has the red book.\nzoe has a green kite. zoe goes to the kitchen. at the kitchen, zoe gives the kite to leo. now leo has the green kite to",
  "steps": [
   {
    "token": "e",
    "probability": 1.0
   },
   {
    "token": " to",
    "probability": 1.0
   }
  ],
  "stopped": "context full"
 },
 "dev_story_table": [
  {
   "setting": "greedy",
   "temperature": 0.0,
   "stories": 60,
   "well_formed": 60,
   "keeps_color_and_object": 27,
   "stopped": {
    "end of story": 60,
    "max_new_tokens": 0,
    "context full": 0
   }
  },
  {
   "setting": "temperature 0.5",
   "temperature": 0.5,
   "stories": 60,
   "well_formed": 59,
   "keeps_color_and_object": 31,
   "stopped": {
    "end of story": 60,
    "max_new_tokens": 0,
    "context full": 0
   }
  },
  {
   "setting": "temperature 1.0",
   "temperature": 1.0,
   "stories": 60,
   "well_formed": 56,
   "keeps_color_and_object": 30,
   "stopped": {
    "end of story": 60,
    "max_new_tokens": 0,
    "context full": 0
   }
  },
  {
   "setting": "temperature 1.5",
   "temperature": 1.5,
   "stories": 60,
   "well_formed": 26,
   "keeps_color_and_object": 11,
   "stopped": {
    "end of story": 59,
    "max_new_tokens": 0,
    "context full": 1
   }
  },
  {
   "setting": "temperature 1.0, top-k 3",
   "temperature": 1.0,
   "top_k": 3,
   "stories": 60,
   "well_formed": 57,
   "keeps_color_and_object": 30,
   "stopped": {
    "end of story": 60,
    "max_new_tokens": 0,
    "context full": 0
   }
  }
 ]
}

Failure cases

Each case has a named test in test_generate.py.

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

python3 -m unittest -v test_generate.py

"""Tests for generate.py. Run: python3 -m unittest -v test_generate.py"""
import math
import subprocess
import sys
import unittest
from pathlib import Path

import checkpoint as ck
import generate as g
import model as m
import train as t
from bpe import Tokenizer

HERE = Path(__file__).parent
TOK = Tokenizer.load(HERE / "tokenizer.json")
P = ck.load(HERE / "best.json", t.CONFIG, ck.input_signature(TOK, (HERE / "corpus.txt").read_text()))["params"]
LOGITS = m.forward(P, TOK.encode("mia has a red ball. mia goes to the park. now leo has the"))[0][-1]


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


class Selection(unittest.TestCase):
    def test_greedy_puts_all_probability_on_the_top_score(self):
        probs = g.next_token_probs(LOGITS, 0.0, None)
        self.assertEqual(probs.index(1.0), max(range(len(LOGITS)), key=lambda i: LOGITS[i]))
        self.assertEqual(sum(probs), 1.0)

    def test_greedy_ties_go_to_the_lowest_token_id(self):
        self.assertEqual(g.next_token_probs([1.0, 5.0, 5.0, 2.0], 0.0, None), [0.0, 1.0, 0.0, 0.0])

    def test_temperature_one_is_plain_softmax(self):
        self.assertEqual(g.next_token_probs(LOGITS, 1.0, None), m.softmax_row(LOGITS))

    def test_lower_temperature_sharpens(self):
        tops = [max(g.next_token_probs(LOGITS, temp, None)) for temp in (2.0, 1.0, 0.5)]
        self.assertEqual(tops, sorted(tops))
        self.assertLess(tops[0], tops[2])

    def test_top_k_keeps_exactly_k_tokens(self):
        probs = g.next_token_probs(LOGITS, 1.0, 3)
        kept = [i for i, v in enumerate(probs) if v > 0]
        self.assertEqual(len(kept), 3)
        self.assertEqual(sorted(kept), sorted(sorted(range(len(LOGITS)), key=lambda i: -LOGITS[i])[:3]))
        self.assertAlmostEqual(math.fsum(probs), 1.0, places=12)

    def test_top_k_ties_keep_the_highest_scores(self):
        probs = g.next_token_probs([1.0, 1.0, 5.0], 1.0, 2)
        self.assertEqual([v > 0 for v in probs], [True, False, True])
        self.assertEqual([v > 0 for v in g.next_token_probs([2.0, 2.0, 2.0, 0.0], 1.0, 2)], [True, True, False, False])

    def test_same_seed_same_text(self):
        a = g.generate(P, TOK, "zoe has a blue kite.", temperature=1.0, seed=1)
        b = g.generate(P, TOK, "zoe has a blue kite.", temperature=1.0, seed=1)
        c = g.generate(P, TOK, "zoe has a blue kite.", temperature=1.0, seed=2)
        self.assertEqual(a, b)
        self.assertNotEqual(a["text"], c["text"])


class Stopping(unittest.TestCase):
    def test_stops_at_the_end_of_a_story(self):
        result = g.generate(P, TOK, "mia has a red ball.")
        self.assertEqual(result["stopped"], "end of story")
        self.assertTrue(result["text"].endswith("\n"))

    def test_stops_at_max_new_tokens(self):
        result = g.generate(P, TOK, "mia has a red ball.", max_new_tokens=2)
        self.assertEqual((result["stopped"], len(result["steps"])), ("max_new_tokens", 2))

    def test_stops_when_the_context_is_full(self):
        prompt = TOK.decode(TOK.encode((HERE / "corpus.txt").read_text())[:63])
        result = g.generate(P, TOK, prompt)
        self.assertEqual(result["stopped"], "context full")
        self.assertEqual(len(TOK.encode(result["text"])), 64)


class FailureCases(unittest.TestCase):
    def test_invalid_temperature(self):
        for bad in (-0.1, 0.001, float("nan"), float("inf"), True, "1"):
            with self.subTest(temperature=bad), self.assertRaisesRegex(ValueError, "temperature must be"):
                g.generate(P, TOK, "mia", temperature=bad)

    def test_invalid_top_k(self):
        for bad in (0, 97, True, 2.0):
            with self.subTest(top_k=bad), self.assertRaisesRegex(ValueError, "top_k must be"):
                g.generate(P, TOK, "mia", temperature=1.0, top_k=bad)

    def test_empty_prompt(self):
        with self.assertRaisesRegex(ValueError, "empty prompt"):
            g.generate(P, TOK, "")

    def test_prompt_longer_than_the_context(self):
        prompt = TOK.decode(TOK.encode((HERE / "corpus.txt").read_text())[:65])
        self.assertEqual(len(TOK.encode(prompt)), 65)
        with self.assertRaisesRegex(ValueError, "prompt of 65 tokens is longer than the 64 positions"):
            g.generate(P, TOK, prompt)

    def test_story_check_rejects_a_broken_story(self):
        broken = g.generate(P, TOK, "maya has a purple umbrella.")["text"]
        self.assertEqual(g.follows_story("maya has a purple umbrella.", broken),
                         {"well_formed": False, "keeps_color_and_object": False})


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

test_lower_temperature_sharpens fails if temperature is not applied. test_top_k_keeps_exactly_k_tokens fails if top-k keeps one token too few or too many, and test_top_k_ties_keep_the_highest_scores fails if a tie at lower scores pushes out the highest one. test_stops_when_the_context_is_full fails if the loop asks for a position the model does not have.

Deployment

There is still nothing to serve; Build a server for your trained LLM adds an HTTP interface. Whoever calls generate() should record the prompt, settings, seed, checkpoint and stop reason with every output, as generation-trace.json does.

Scaling

Each new token runs the full model again on the growing prefix. For a fixed prompt and n additional tokens, the repeated row computations grow quadratically with n. Each full attention pass also compares pairs of prefix positions, so the total attention-score work grows cubically. These are operation counts at fixed model width, not a wall-clock measurement; the position-table limit bounds this small implementation. The next build caches the keys and values of earlier positions and measures the difference.

Key takeaways

• Generation is a loop: run the model, take the last position’s scores, choose a token, check the stop conditions, append, repeat.
• Temperature and top-k change how a token is drawn from the model’s scores. They cannot make the model rank the right token first.
• Record the seed and the stop reason with every output; without them a sampled result cannot be reproduced or interpreted.

Best practices

• Test each control on real model scores: greedy picks the top score, temperature 1 equals plain softmax, top-k leaves exactly k candidates.
• Compare settings on development prompts with a check that matches the task, here the story’s color and object, and keep the test prompts for a final check of the chosen setting.
• Treat any stop reason other than the natural end as an incomplete output.

References

Sources checked October 20, 2026.

• Holtzman et al., The Curious Case of Neural Text Degeneration (ICLR 2020), arXiv 1904.09751v2: temperature as division of the scores before softmax (section 3.3), top-k sampling from the k most probable tokens after renormalizing (3.2), maximization-based decoding producing repetitive text in open-ended generation, and the unreliable tail of full-distribution sampling.
• Hugging Face Transformers v5.20: generation strategies and GenerationConfig: greedy search is the default strategy; do_sample, temperature and top_k select and shape sampling, and a model’s own generation config can change the library’s fallback defaults.