Chapter 6 of 10

Test it

The model has learned from its training examples. Now I want to know how it handles sums it hasn't practised. I kept a separate set of examples for this check. This is the test set.

Testing doesn't change the weights. I give the model each question, let it write an answer and compare that answer with the one in the test set. It also has to stop in the right place.

I also try numbers above 99. These are outside the range it was trained on. Doing well on familiar sums doesn't tell us how it will do on these.

How did it do on new questions?

These are saved test results. Compare sums within the training range with sums using larger numbers.

Questions to check

Correct answers that stopped in the right place: 501 of 502

An answer counts as correct only if the digits match and the model sends its ending token, EOS. This test allows up to four output tokens, including that ending.

More details from the test

Check	Answers right
Needs a carry	351 of 352
No carry needed	150 of 150
Digits in the answer: 1	5 of 5
Digits in the answer: 2	255 of 256
Digits in the answer: 3	241 of 241

The model sent an ending token in 502 of 502 cases. Stopping doesn't tell us whether the answer was correct.

Try a test question

Run 9+2=. The answer should be 11. Then try 100+1=, which should give 101. You can also compare 1+1= with 1+1 to see why the exact text matters.

Runs on your device. First load is about 14 MB.

Improved model

Your input

Exact input: "9+2="

Model output

No output yet

Input + output

9+2=

Run to see a prediction.

Model and input details

Printable ASCII is passed verbatim. BOS is added once. All 98 tokens remain available. Context: 32 tokens. Run budget: 16 output tokens, capped by remaining context. EOS ends generation. Unexpected PAD or BOS tokens are visible.

Improved model. Epoch 140. 1.74 MB of weights. The WASM runtime is downloaded on first use.

Optional code

Code and commands for this step

Now load your saved weights and generate an answer. This code runs from the course folder:

import torch
from tiny_gpt.checkpoint import load_checkpoint
from tiny_gpt.generation import generate_answer

model, metadata = load_checkpoint("runs/my-model/best.pt")
answer = generate_answer(model, "1+1=", max_new_tokens=4)
print(answer.text, answer.ended)

Save it as try_model.py and run python try_model.py. ended is true when the model emits EOS. Use runs/my-first-run/best.pt instead to inspect the one-epoch model. The browser below runs my saved weights. Your local run may produce different answers.

Read the complete generation.py

"""Generate tokens using the network alone. No calculator or answer lookup."""

from dataclasses import dataclass

import torch

from .examples import encode_prompt
from .tokenizer import CharacterTokenizer


@dataclass(frozen=True)
class GeneratedAnswer:
    token_ids: tuple[int, ...]
    ended: bool

    @property
    def text(self):
        body = self.token_ids[:-1] if self.ended else self.token_ids
        # Preserve unexpected PAD/BOS markers rather than hiding malformed output.
        return CharacterTokenizer().decode(body)

    def matches(self, expected_answer):
        expected = tuple(CharacterTokenizer().encode(expected_answer, add_eos=True))
        return self.ended and self.token_ids == expected


@torch.inference_mode()
def greedy_generate_batch(model, prompts, *, max_new_tokens=4):
    """Prompts must have equal encoded lengths. Choices use all 98 tokens."""
    if type(max_new_tokens) is not int or max_new_tokens <= 0:
        raise ValueError("max_new_tokens must be a positive integer")
    prompts = list(prompts)
    if not prompts:
        raise ValueError("At least one prompt is required")
    tokenizer = CharacterTokenizer()
    encoded = [encode_prompt(prompt, tokenizer, context_length=model.config.context_length)
               for prompt in prompts]
    if len({len(ids) for ids in encoded}) != 1:
        raise ValueError("Batch prompts by equal length; do not left-pad them")
    if len(encoded[0]) + max_new_tokens > model.config.context_length:
        raise ValueError("Prompt plus output limit exceeds the context length")
    device = next(model.parameters()).device
    ids = torch.tensor(encoded, dtype=torch.long, device=device)
    answers = [[] for _ in prompts]
    ended = [False for _ in prompts]
    was_training = model.training
    model.eval()
    try:
        for _ in range(max_new_tokens):
            logits = model(ids)[:, -1, :]
            selected = logits.argmax(dim=-1)
            for index, token_id in enumerate(selected.cpu().tolist()):
                if not ended[index]:
                    answers[index].append(token_id)
                    ended[index] = token_id == tokenizer.EOS_ID
            if all(ended):
                break
            # Already-ended rows are ignored; their later choices cannot affect other rows.
            ids = torch.cat((ids, selected[:, None]), dim=1)
    finally:
        model.train(was_training)
    return [GeneratedAnswer(tuple(tokens), finished) for tokens, finished in zip(answers, ended)]


def generate_answer(model, prompt, *, max_new_tokens=4):
    return greedy_generate_batch(model, [prompt], max_new_tokens=max_new_tokens)[0]

Download generation.py

To test every reserved example, call the same generator with each prefix, then compare both the answer and ending token:

import json
from pathlib import Path

rows = [json.loads(line) for line in
        Path("data/addition-v1/test.jsonl").read_text().splitlines()]
correct = 0
for row in rows:
    result = generate_answer(model, row["prompt"], max_new_tokens=4)
    correct += result.matches(row["answer"])
print(f"{correct}/{len(rows)} correct")

Add this to try_model.py. Repeat with ood.jsonl to check numbers outside the training range.

Technical details

Test and out-of-range sets never select checkpoints. Evaluation groups prompts by equal encoded length without left padding. Exact correctness includes EOS. Recorded measurements identify their source checkpoint and dataset hashes.

Code and training steps