HeyMSK
Chapter 6 of 10

Test it

I freeze the model and test it on examples that did not choose its weights. The improved base gets 501 of 502 answers right, including EOS. It gets none of the 497 outside-range examples right.

Try 1+1 without the equals sign. The model completes it as 1+15=16. It is predicting a continuation from learned patterns, so even a familiar prefix can produce an unexpected result.

Run 9+2= and compare it with 11. Try 100+1= against the reference answer 101. Switch the results table to the outside range, then compare 1+1 with 1+1=.

Frozen improved-base evaluation
MeasureObserved count
Exact answer and EOS501 / 502
Predicted EOS502 / 502
With carry351 / 352
Without carry150 / 150
Answer digits: 15 / 5
Answer digits: 2255 / 256
Answer digits: 3241 / 241

Four output tokens, including EOS. Weights were frozen before the exam.

Recorded mistake: 9+2= continued as 10. Expected 11.

Build 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. The test never updates the weights. Repeat with ood.jsonl to check unfamiliar operands.

Test the trained model

Try 9+2=. The reference answer is 11. Then try 100+1=, outside the training range. A familiar-looking task can still fail.

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

Improved base

Exact input: "9+2="

Continuation
No output yet
Input + continuation
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 base. Epoch 140. 1.74 MB of weights. The WASM runtime is downloaded on first use.

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