Train a tiny language model

What you will be able to do

Train a language model from scratch, on a laptop, and watch it learn to write. You build a small decoder from the parts of this course and train it on six Shakespeare plays. You measure it against simple counting baselines on text it never saw. Then you can say, with numbers, what it learned and how far it is from a real large language model.

Case studyPortfolio piece

Dataset

Six Shakespeare tragedies (Project Gutenberg eBook #100)

Licence

Public domain in the USA

Time

About 2 hours

Runs in

Colab, Kaggle or VS Code

Skills used

• Character tokens
• Contiguous split
• Counting baselines
• Decoder with a causal mask
• Samples during training
• Temperature
• Save and reload

Open in ColabOpen in KaggleView on GitHub

Colab and Kaggle open your own copy. To use VS Code, download the notebook from GitHub. Stuck? See the solution notebook.

The notebook runs every step of this page, with the same code and the same output, plus four tasks with checks. Its setup cell downloads the text (about 0.9 MB). Training takes about 3 minutes on the shared Mac mini where this page was made: 168 to 212 seconds in five runs on 2 CPU threads; the whole notebook ran in about 3 minutes. A slower computer, such as free Colab, can take two or three times as long.

Before you startWhat to know and to have readyAbout 120 minutes

What you need to know

This case study uses Modules 5 to 7 most. These lessons matter most:

• Tokens, token IDs and the vocabulary, from Tokens and embeddings.
• The transformer block, from Transformer blocks, and the causal mask, from Position and masking.
• Next-token prediction and the model that cheats without the mask, from Training objectives.
• The generation loop, greedy choice and temperature, from Training versus inference.
• Saving a state_dict and reloading it in a fresh process, from Save and inspect results.

What you need on your computer (only for the local route; the notebook needs none of it)

• The folder ticket-router with its .venv, from Set up PyTorch: Route 2. This case study needs only PyTorch; it does not use the tickets or ticketnet.py.
• The notebook, downloaded from the link above, in VS Code or JupyterLab. Its setup cell makes a new folder tiny-lm. To start again, delete that folder.

Tested with Python 3.12 and PyTorch 2.14.1 on the CPU (macOS, 2 threads), in Jupyter from a fresh start. Windows and Linux are not tested.

The situation

Nadia runs education at a city science museum. She is planning an exhibit called “How does a chatbot learn to write?”. Visitors often ask whether chatbots understand what they write. Nadia wants them to see the answer, not only read it on a sign.

Her idea: a small computer in the hall trains a language model. A screen shows what it writes at the start, after a few minutes and at the end. Ben, a data scientist who volunteers at the museum, builds it. Nadia has four conditions:

• It runs on an ordinary computer, with no graphics card and no internet, and trains in a few minutes, so that a visitor can watch.
• The text is free to use and show in public. The museum cannot use text that someone owns.
• The progress is real. The exhibit must not show a trick. Ben must measure the model against something simpler, on text that it did not train on.
• It is honest about the gap. A sign must say how far this model is from the chatbots that visitors use.

Ben chooses Shakespeare: everyone recognises the look of a play, and the text has been free for centuries.

What you will build

You will build a character-level language model: a decoder transformer that reads up to 64 characters and predicts the next one. It learns with the next-token objective of Training objectives. You compare it with counting baselines and read samples from four points in training. Then you try the temperature, and reload the model in a fresh process to write from your own prompt.

The data

The text is six tragedies by William Shakespeare: Hamlet, Julius Caesar, King Lear, Macbeth, Othello and Romeo and Juliet. Ben compared five candidate texts, including Pride and Prejudice, The Adventures of Sherlock Holmes and Grimms’ Fairy Tales. He chose the plays because their shape is easy to see: a speaker’s name in capitals on its own line, then a speech in short lines. A model that learns this shape shows its progress at a glance.

Source and status

Six plays from “The Complete Works of William Shakespeare”, Project Gutenberg eBook #100, file pg100.txt (updated August 24, 2025; downloaded 2026-10-08; SHA-256 3cf4b3d4…0a37). The plays are public domain in the USA. Shakespeare died in 1616, so they are free in other countries too. Still, the law of the country where you are applies, so check it. Nextia Learning kept only the six plays and removed the Project Gutenberg header, footer and licence. The licence asks for this when you share the text without them. Nothing inside a play was changed. The course hosts this copy, because Project Gutenberg updates its files and blocks automated downloads. The file shakespeare_tragedies.txt has SHA-256 59bd62e6…94b4; the dataset card next to the notebook has the full checksums.

Ben loads the text and prints the first lines of Hamlet:

text = Path("shakespeare_tragedies.txt").read_text(encoding="utf-8")
print(f"{len(text):,} characters, {len(text.splitlines()):,} lines")
start = text.index("BARNARDO.\nWho’s there?")
print(text[start:start + 196])

The output:

850,875 characters, 33,143 lines
BARNARDO.
Who’s there?

FRANCISCO.
Nay, answer me. Stand and unfold yourself.

BARNARDO.
Long live the King!

FRANCISCO.
Barnardo?

BARNARDO.
He.

FRANCISCO.
You come most carefully upon your hour

What this tells us. The text has about 850,000 characters. Every speech starts with a name in capitals and a full stop, on its own line. That is the shape to watch for in the samples.

First checks

A language model’s vocabulary is the list of symbols it can read and write, so Ben counts the characters. He also counts the speaker lines with a regular expression: a line of capital letters that ends with a full stop.

import re
from collections import Counter

counts = Counter(text)
print(len(counts), "different characters")
print("most common:", [(c, n) for c, n in counts.most_common(8)])
print("rarest:     ", [(c, n) for c, n in counts.most_common()[-5:]])
SPEAKER = re.compile(r"^([A-Z][A-Z’ ]+)\.$", flags=re.M)      # a line like "HAMLET."
print("speaker lines:", len(SPEAKER.findall(text)), "| different speakers:", len(set(SPEAKER.findall(text))))

The output:

75 different characters
most common: [(' ', 128959), ('e', 70034), ('t', 51382), ('o', 48881), ('a', 42744), ('s', 38314), ('h', 37847), ('n', 37325)]
rarest:      [('X', 23), ('&', 16), ('æ', 7), ('\t', 2), ('à', 1)]
speaker lines: 5647 | different speakers: 184

What this tells us. The whole text uses only 75 different characters, and some are very rare: à appears once. The text has no digits, and it uses the curly apostrophe ’, never the straight one '. That matters at the end. There are 5,647 speeches by 184 different speakers.

Frame the problem

• The unit: one position in the text. The input is the characters before it, up to 64. The target is the next character.
• The labels: none written by a person. The text is its own answer key, as in Training objectives.
• The metric: the validation loss. This is the average cross-entropy per character, −ln p of the real next character, on text that the model did not train on (Loss). Lower is better. One way to read it: a loss L means that the model is as unsure as a choice among e^L equally likely characters.
• What the exhibit needs in addition: samples that a visitor can read. The loss is the evidence for Ben, and the samples are the evidence for Nadia’s visitors. The case study uses both.

Split the text: one cut, no shuffling

The validation text must be text that the model never saw. Ben cuts the text once: the first 90% is for training, the last 10% for validation.

cut = int(len(text) * 0.9)
train_text, valid_text = text[:cut], text[cut:]
print(f"train {len(train_text):,} characters | valid {len(valid_text):,} characters")
print("valid starts:", repr(valid_text[:60]))

The output:

train 765,787 characters | valid 85,088 characters
valid starts: ' to your father’s? We’ll to dinner thither.\n\nROMEO.\nI will f'

Why not shuffle the lines and take 10% at random, as you would with tickets? Because neighbouring lines are not independent. A verse line continues the line before it, and a speech answers the one before. If line 2 of a speech is in training and line 3 is in validation, the model has seen most of the context of the answer. The validation loss would then measure memory of these scenes, not the skill of writing new text. One cut keeps validation a block of text that the model has never read. This is the same leakage as in Train, validation and test sets in the machine learning course.

What this tells us. Validation is the second half of Romeo and Juliet, from the middle of a sentence in Act 2. It shares the style and some names with training, so it measures “more of these plays”, not any English text. There is no separate test set. Ben makes only one choice that the validation loss could bias, the model size, so he states this limit instead.

Characters as tokens

The familiar situation. A typewriter has about 80 keys, and every English text typed on it uses only those keys. You never meet an unknown key.

The problem. The ticket router used word tokens, and every word outside its vocabulary became <unk> (Tokens and embeddings). A model that writes cannot work with <unk>. It must be able to spell names such as “Guildenstern” and old forms such as “lov’d”.

The idea. Use one token per character. The vocabulary is then tiny, and no word is unknown, because every word is made of characters. Large language models use subword tokens instead: common words are one token, and rare words are cut into common pieces. From text to tokens in the large language models course shows how. A subword tokenizer is learned from a very large text collection, and it would add a second, unfamiliar step. Characters are the simplest choice that can write anything, and they make the model’s spelling visible. The cost is length: a sentence becomes several times more tokens, so the model’s 64-character window holds only about 12 words.

chars = sorted(set(train_text))                 # the vocabulary: every character of the training text
stoi = {c: i for i, c in enumerate(chars)}      # character -> ID


def encode(s):
    return [stoi[c] for c in s]


def decode(ids):
    return "".join(chars[i] for i in ids)


print(len(chars), "characters in the vocabulary")
print("in valid but not in train:", sorted(set(valid_text) - set(chars)))
print(repr("".join(chars)))
print(encode("ROMEO."), "->", repr(decode(encode("ROMEO."))))
train_ids = torch.tensor(encode(train_text))
valid_ids = torch.tensor(encode(valid_text))
show_shape("train_ids", train_ids)
show_shape("valid_ids", valid_ids)

The output:

75 characters in the vocabulary
in valid but not in train: []
'\t\n !&(),-.:;?ABCDEFGHIJKLMNOPQRSTUVWXYZ[]_abcdefghijklmnopqrstuvwxyzàæ—‘’“”'
[30, 27, 25, 17, 27, 9] -> 'ROMEO.'
train_ids    shape (765787,)
valid_ids    shape (85088,)

What this tells us. The vocabulary comes from the training text only, the habit from Debug the experiment. Here every validation character is also in training, so nothing is lost. “ROMEO.” becomes six IDs, and decode turns them back into the same text. The whole training text is now one long row of 765,787 numbers. show_shape() is the notebook’s helper that prints a tensor’s shape.

Baselines: what counting achieves

Before any network, Ben asks how far counting goes. Three baselines, each with its validation loss:

1. Every character equally likely: the loss is ln 75 = 4.317. A model that learned nothing scores this.
2. Character frequencies: predict each character with its share in the training text, whatever came before.
3. Pair counts (a bigram model): count how often each character follows each other character in training. Predict the next character from the current one only. “Bigram” means “two letters”.

import math

V = len(chars)
print("1. every character equally likely:", round(math.log(V), 3))

counts = torch.bincount(train_ids, minlength=V).float() + 1          # +1: no character has chance 0
p_char = counts / counts.sum()
print("2. character frequencies:         ", round(float(-p_char[valid_ids].log().mean()), 3))

pairs = torch.ones(V, V)                                              # +1 for every pair, as above
pairs.index_put_((train_ids[:-1], train_ids[1:]), torch.ones(len(train_ids) - 1), accumulate=True)
p_next = pairs / pairs.sum(1, keepdim=True)                           # row = this character, column = next
bigram_loss = float(-p_next[valid_ids[:-1], valid_ids[1:]].log().mean())
print("3. pair counts (bigram):          ", round(bigram_loss, 3))

The output:

1. every character equally likely: 4.317
2. character frequencies:          3.313
3. pair counts (bigram):           2.499

The + 1 gives every character and every pair at least one count. Without it, a pair that never occurs in training would get the chance 0, and one such pair in validation would make the loss infinite. index_put_(..., accumulate=True) adds 1 to the cell of each pair in the text: the row is the character, the column the character after it.

What does the pair-count model write? It chooses each next character at random, with the counted chances:

def bigram_generate(prompt, length=200, seed=0):
    g = torch.Generator().manual_seed(seed)
    ids = encode(prompt)
    for _ in range(length):
        ids.append(int(torch.multinomial(p_next[ids[-1]], 1, generator=g)))
    return decode(ids)


print(bigram_generate("ROMEO.\n"))

The output:

ROMEO.
Than thiflsiapan sdl I avees [_]fo l f be—Wes me worcacothtt cas d f t t d, thakn ndulifolld t ’d heu in far be

FI’l, Shar uckncathinorath whal ce s; ke t;
I isk s athowonoulld t.

Prklolt, freue the

What this tells us. Counting pairs cuts the uncertainty from a choice among 75 characters to about 12 (e^2.499 ≈ 12.2). The sample has short “words”, spaces and line breaks, and a few real words (“Than”, “me”, “far”, “be”). It cannot spell longer words, because it sees only one character back. 2.499 is the number to beat. If a transformer cannot beat a table of pair counts, it is not worth its parameters.

The model

Ben builds a decoder from the parts of this course. A character table gives each character a row of 128 numbers, as the word embeddings did (Tokens and embeddings). A position table adds a row for each of the 64 positions. Three transformer blocks with the causal mask follow, so each position reads only itself and the characters before it. A final normalization and an output layer give a score for each of the 75 characters at every position.

The tiny language model, with the tensor shapes for a batch of 32 windows of 64 characters. The embeddings and the blocks with the causal mask are the parts from Modules 5 and 6.

Text description of the diagram

The diagram is a chain of six boxes, from left to right. 1: a text window of 64 characters, such as “ROMEO.”, a new line and “I will”. 2: the character IDs, shape (32, 64): 32 windows of 64 IDs. 3: the character table plus the position table, shape (32, 64, 128). 4: three transformer blocks with the causal mask, shape (32, 64, 128). 5: a norm and the output layer, which give a score for every character, shape (32, 64, 75). 6: the loss, which compares the scores with the same window moved one character to the right. Boxes 3 and 4 are highlighted: they are the parts from Modules 5 and 6.

from torch import nn


class TinyLM(nn.Module):
    """Character + position vectors -> blocks with a causal mask -> a score for every character."""

    def __init__(self, vocab_size, dim=128, heads=4, blocks=3, context=64, dropout=0.1):
        super().__init__()
        self.context = context
        self.embed = nn.Embedding(vocab_size, dim)            # one row of numbers per character
        self.position = nn.Embedding(context, dim)            # one row of numbers per position
        block = nn.TransformerEncoderLayer(dim, heads, 4 * dim, dropout=dropout,
                                           batch_first=True, norm_first=True)
        self.blocks = nn.TransformerEncoder(block, blocks, enable_nested_tensor=False)
        self.norm = nn.LayerNorm(dim)
        self.out = nn.Linear(dim, vocab_size)                 # a score for every character

    def forward(self, ids):                                   # ids: (batch, n)
        n = ids.shape[1]
        h = self.embed(ids) + self.position(torch.arange(n))  # (batch, n, dim)
        mask = nn.Transformer.generate_square_subsequent_mask(n)
        h = self.blocks(h, mask=mask, is_causal=True)         # each position reads only the past
        return self.out(self.norm(h))                         # (batch, n, vocab_size)


torch.manual_seed(0)
model = TinyLM(V)
for name, part in [("character table", model.embed), ("position table", model.position),
                   ("3 blocks", model.blocks), ("final norm", model.norm), ("output layer", model.out)]:
    print(f"{name:16s} {sum(p.numel() for p in part.parameters()):>8,}")
print(f"{'total':16s} {sum(p.numel() for p in model.parameters()):>8,}")

The output:

character table     9,600
position table      8,192
3 blocks          594,816
final norm            256
output layer        9,675
total             622,539

What this tells us. This is the NextWordModel of Training objectives with characters in place of words, one more block and a final norm. As there, nn.TransformerEncoderLayer with a causal mask works as a decoder block. 96% of the 622,539 parameters are in the three blocks. The character table is small (75 × 128 = 9,600), because the vocabulary is small.

One training batch

A training example is a window: 64 characters from a random place in the training text. The answers are the same 64 characters moved one place to the right, so every position has its own next character to predict.

def get_batch(ids, generator, batch_size=32, context=64):
    """batch_size random windows of the text, and the same windows moved one character to the right."""
    starts = torch.randint(len(ids) - context - 1, (batch_size,), generator=generator)
    x = torch.stack([ids[s:s + context] for s in starts])
    y = torch.stack([ids[s + 1:s + context + 1] for s in starts])
    return x, y


x, y = get_batch(train_ids, torch.Generator().manual_seed(0))
show_shape("x", x)
show_shape("y", y)
print("x[0]:", repr(decode(x[0].tolist())[:40]))
print("y[0]:", repr(decode(y[0].tolist())[:40]))
loss_fn = nn.CrossEntropyLoss()
with torch.no_grad():
    scores = model(x)
show_shape("scores", scores)
print("untrained loss:", round(loss_fn(scores.reshape(-1, V), y.reshape(-1)).item(), 3), "| ln 75 =", round(math.log(V), 3))

The output:

x            shape (32, 64)
y            shape (32, 64)
x[0]: 'eak.\n\nSECOND CITIZEN.\nI will hear Cassiu'
y[0]: 'ak.\n\nSECOND CITIZEN.\nI will hear Cassius'
scores       shape (32, 64, 75)
untrained loss: 4.442 | ln 75 = 4.317

What this tells us. One batch holds 32 × 64 = 2,048 predictions. y[0] is x[0] moved by one character, so the answer after “Cassiu” is “s”. The scores have one row of 75 for every position. The untrained model’s loss, 4.442, is close to ln 75: it guesses almost evenly, as the untrained ticket network did with ln 5.

The generation loop

To write, the model repeats one step, as in Training versus inference: score the next character, choose one, add it, run again. The model reads at most its last 64 characters.

@torch.no_grad()
def generate(model, prompt, length=200, temperature=0.8, seed=0):
    """The generation loop: score the next character, choose one, add it, repeat."""
    model.eval()
    g = torch.Generator().manual_seed(seed)
    ids = encode(prompt)
    for _ in range(length):
        scores = model(torch.tensor([ids[-model.context:]]))[0, -1]    # the last position's scores
        if temperature == 0:
            next_id = int(scores.argmax())                              # greedy
        else:
            probs = torch.softmax(scores / temperature, dim=-1)
            next_id = int(torch.multinomial(probs, 1, generator=g))
        ids.append(next_id)
    return decode(ids)

@torch.no_grad() turns off gradients for the whole function, and model.eval() turns off dropout. torch.multinomial draws one character with the given chances, and the seeded generator g makes the draw repeatable. The untrained model writes noise such as qlFT-t[KOLsn)àNR;sWfC….

Train, and watch the samples

The training loop has the five steps of The training loop. Every 100 steps, Ben measures the loss on 20 fixed batches of training text and 20 of validation text. The batches stay the same, so the curves are comparable. At steps 0, 300, 1,000 and 3,000, he saves a sample of 1,000 characters. Each starts from “ROMEO.” and a new line, with seed 0 and temperature 0.8.

import time

STEPS, EVAL_EVERY, SAMPLE_AT = 3000, 100, (0, 300, 1000, 3000)


@torch.no_grad()
def estimate(model, ids, batches=20):
    """Mean loss on 20 fixed random batches (the same batches every time)."""
    model.eval()
    g = torch.Generator().manual_seed(1)
    losses = [loss_fn(model(x).reshape(-1, V), y.reshape(-1)).item()
              for x, y in (get_batch(ids, g) for _ in range(batches))]
    return sum(losses) / len(losses)


torch.manual_seed(0)
model = TinyLM(V)
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-3)
batches = torch.Generator().manual_seed(0)
history, samples = [], {}
started = time.perf_counter()
for step in range(STEPS + 1):
    if step % EVAL_EVERY == 0:
        history.append({"step": step, "train": round(estimate(model, train_ids), 4),
                        "valid": round(estimate(model, valid_ids), 4)})
        if step % 500 == 0:
            print(history[-1], f"{time.perf_counter() - started:.0f} s")
    if step in SAMPLE_AT:
        samples[step] = generate(model, "ROMEO.\n", length=1000, temperature=0.8, seed=0)
    if step == STEPS:
        break
    model.train()                                  # dropout on for training
    x, y = get_batch(train_ids, batches)
    optimizer.zero_grad()
    loss = loss_fn(model(x).reshape(-1, V), y.reshape(-1))
    loss.backward()
    optimizer.step()
print(f"trained {STEPS} steps in {time.perf_counter() - started:.0f} s")

The output:

{'step': 0, 'train': 4.4342, 'valid': 4.4412} 0 s
{'step': 500, 'train': 1.8886, 'valid': 1.9545} 33 s
{'step': 1000, 'train': 1.6678, 'valid': 1.7638} 60 s
{'step': 1500, 'train': 1.5865, 'valid': 1.6757} 88 s
{'step': 2000, 'train': 1.5315, 'valid': 1.6259} 115 s
{'step': 2500, 'train': 1.5025, 'valid': 1.6037} 143 s
{'step': 3000, 'train': 1.4781, 'valid': 1.5777} 171 s
trained 3000 steps in 172 s

The losses are the same in every run on this computer; the seconds change with the load (168 to 212 in five runs). AdamW is Adam with a small weight decay (Overfitting controls). After a short trial, Ben chose a learning rate of 0.002, twice the course’s default, so that the run fits in a few minutes. There are no epochs: 3,000 steps of 32 random windows make 6,144,000 predictions, about 8 for each training character.

• Training text
• Validation text
• Character frequencies (dashed)
• Pair counts (dashed)

Training and validation loss of the tiny language model every 100 steps (mean of 20 fixed batches), with the two counting baselines on the validation text. The model passes the pair counts at step 100 and is still improving at step 3,000.

Text description of the chart

A line chart with the training step, 0 to 3,000, on the horizontal axis and the loss per character, 1 to 4.6, on the vertical axis. Two dashed flat lines are the baselines on the validation text: character frequencies at 3.313 and pair counts at 2.499. The training and validation curves both start at about 4.44. They fall very fast. At step 100 the validation loss is 2.474, just below the pair counts, and at step 300 it is 2.101. Then they fall more slowly: validation 1.764 at step 1,000, 1.626 at step 2,000 and 1.578 at step 3,000. The training curve stays a little below the validation curve, 0.10 lower at the end, and both still slope down.

What this tells us. Read the curves as in Batches, epochs and validation. Both fall together, and the validation loss is still falling at step 3,000. So the model is not overfitting yet, and the last step is the best. More steps would help a little, at the cost of more minutes. The gap of about 0.1 between the curves is small.

The samples

Now read what the model wrote at each point, all with the same prompt, seed and temperature:

for step, sample in samples.items():
    print(f"--- step {step} ---")
    print(sample[:230])

The output:

--- step 0 ---
ROMEO.
qlFT-t[KOLsn)àNR;sWfCIVYvnjZs!H_fc)H‘( )&—tKokvwXPoEcaZc”htH?Iaà;Dw:XkTkZdackn.kOvnvu“:QKQUvcDN’;]Hiu ij:[’e YZ
szI’U,TSDDDvuvsncwtFzPorACn‘._pUUu“hs?;kozt;—ahO[k&æ!a
’M-I;Mey[æ[Y.zEànku;lM-E:Hhuk:Yqjr?ahà)àxiTMu:u“JeJ:M”MJ
--- step 300 ---
ROMEO.
That the lon, ands for aver bed for wo bed youre worcack hit cased cakerst,
Thathinguce on don ’d hin in far be
seave cand thas cathe heatn whal comsicke thy hisk siand worne here.

PORONEMAER.
De thereake, bid unt ther if
--- step 1000 ---
ROMEO.
That the long and did wave seen for which thou well.
How hithin the with damen.

DERIGO.
Where ’t he bid mad be seam you have staw the ray winplace sick the Caesar. and wither’s the Parus thou huse, I have abus unto me:
No
--- step 3000 ---
ROMEO.
I will not bare here I am the world walk,—

SECOND SENGUILDENSTERN.
Thou dock thou do my lov’d ’twas bid mad be sea a cannot conce
For may anyprace safel to the knows,
That eyes the caultree, thus great his bid untony him?

What this tells us. Read the four samples in order; this is the exhibit.

• Step 0: every character is about equally likely, so the model writes noise, rare characters included.
• Step 300, after about 20 seconds: the shape of a play is already there. Lines of verse, a blank line, then a name in capitals with a full stop. The name, “PORONEMAER”, is made up, and most long words are not words (“worcack”, “cakerst”). The common short words are right: “the”, “for”, “bed”, “here”.
• Step 1,000: more real words (“long”, “seen”, “which”, “thou”, “Caesar”). “DERIGO” is close to Roderigo, from Othello.
• Step 3,000: a first line that reads like verse, “I will not bare here I am the world walk,—”. Old forms appear, such as “lov’d” and “’twas”. The speaker “SECOND SENGUILDENSTERN” joins two patterns of real names: “SECOND CITIZEN” and “GUILDENSTERN” from Hamlet. It means nothing.

The samples share fragments, such as “bid mad be sea”, because the same seed gives the same random draws.

The fair comparison

The curves use 20 batches. For the final comparison, Ben measures every model on the whole validation text, under the same metric:

@torch.no_grad()
def full_loss(model, ids):
    """Loss on the whole text, in windows of 64 characters, with dropout off."""
    model.eval()
    n = (len(ids) - 1) // model.context * model.context
    x, y = ids[:n].view(-1, model.context), ids[1:n + 1].view(-1, model.context)
    total = sum(loss_fn(model(xb).reshape(-1, V), yb.reshape(-1)).item() * yb.numel()
                for xb, yb in zip(x.split(256), y.split(256)))
    return total / n


results = {"every character equally likely": math.log(V),
           "character frequencies": float(-p_char[valid_ids].log().mean()),
           "pair counts (bigram)": bigram_loss,
           "tiny transformer": full_loss(model, valid_ids)}
for name, loss in results.items():
    print(f"{name:32s} valid loss {loss:.3f}  ~ a choice among {math.exp(loss):4.1f} characters")
print(f"{'tiny transformer, training text':32s} loss {full_loss(model, train_ids):.3f}")

The output:

every character equally likely   valid loss 4.317  ~ a choice among 75.0 characters
character frequencies            valid loss 3.313  ~ a choice among 27.5 characters
pair counts (bigram)             valid loss 2.499  ~ a choice among 12.2 characters
tiny transformer                 valid loss 1.585  ~ a choice among  4.9 characters
tiny transformer, training text  loss 1.485

Validation loss on the whole held-out text (the second half of Romeo and Juliet). The tiny transformer is far below both counting baselines.

Text description of the chart

A column chart with four columns. All equal (every character equally likely): 4.317. Frequency (character frequencies): 3.313. Pairs (pair counts, the bigram model): 2.499. Tiny transformer: 1.585, highlighted. Each column is lower than the one before it, and the transformer’s column is the shortest by far.

What this tells us. The tiny transformer wins clearly: 1.585 against 2.499 for the pair counts. In plain words, its uncertainty about the next character is like a choice among about 5 characters, against 12 for the pair counts. It wins because it reads up to 64 characters back, not one. So it knows that “Guildenst” is followed by “e”, and that a blank line is followed by a name. In this measurement every window starts fresh, so its first characters have little context. That makes the comparison slightly harder for the transformer, not easier. The loss on the training text, 1.485, is only 0.1 lower: it has not learned the training text by heart.

This is the opposite of the ticket router, where word counts were enough and 4,700 tickets were too few for a transformer. Here the task needs order: a count of characters says nothing about what comes next.

Temperature

The temperature divides the scores before softmax, as in Training versus inference. Below 1, likely characters get even more likely. Above 1, the chances become more even. At 0, the model always takes the top character: greedy choice. First, the model’s own chances for one next character:

prompt = "LEAR.\nGive me the "
with torch.no_grad():
    model.eval()
    scores = model(torch.tensor([encode(prompt)]))[0, -1]
probs = torch.softmax(scores, dim=-1)
for p, i in zip(*probs.topk(8)):
    print(f"{chars[i]!r:6s} {p.item():.3f}")

The output:

'm'    0.121
's'    0.080
'c'    0.068
't'    0.065
'b'    0.062
'w'    0.061
'p'    0.055
'h'    0.043

The model's chances for the character after 'LEAR.' and 'Give me the ', at temperature 1. After a space, many first letters are possible; the top one, m, has only 12%.

Text description of the chart

A column chart with the eight most likely next characters after “Give me the “: m 12%, s 8%, c 7%, t 7%, b 6%, w 6%, p 5%, h 4%. All are first letters of a word, and none has much more than the others. The 20 most likely characters together hold 88% of the chance.

After a space, the model is unsure: many words can follow “Give me the”. Greedy choice always takes m. Sampling at temperature 1 takes m only about one time in eight. Now the same prompt and seed at four temperatures:

for t in (0, 0.5, 1.0, 1.5):
    print(f"--- temperature {t} ---")
    print(generate(model, "LEAR.\n", length=160, temperature=t, seed=0))

The output:

--- temperature 0 ---
LEAR.
I have be the more that that the more of the more of the more of the more of the more of the more of the more of the more of the more of the more of the more of
--- temperature 0.5 ---
LEAR.
I will not be and flower,
The world will be the moor, and the as the corrow
That the diffly come; he be bear of sea a can a like which is not place so king;
If
--- temperature 1.0 ---
LEAR.
That the long and dlaw, villain!

 [_Gloucesternerorcestion, castlemand, damn thou dull clov’d ’tward in my worthy.

 [_Exeunt._]

FIRST CITIZEN.
What ketter hi
--- temperature 1.5 ---
LEAR.
Than thy lon, ands,
Comev’d brinforly- but boy; my cament teach’d
be knocd, and him unmigit com’d him in my worthy.’
, Shall villatizan asnaveplace sackott;—I’l

What this tells us. At 0, the model loops on “the more of” forever: the same trap as SmolLM2’s greedy text in Module 7. At 0.5, the words are almost all real but plain. At 1.0, the model writes a stage direction, [_Exeunt._], with the play’s own markup, and a garbled one before it. At 1.5, the spelling falls apart. The next section measures this trade-off.

What it gets right, and what it gets wrong

Samples are easy to over-read, so Ben counts what a reader checks first. For each sample of 1,000 characters, he measures three shares. The share of words that also appear in the training text. The share of different words (a low share means repetition). And the share of speaker names that are real speakers of the plays.

train_words = set(re.findall(r"[a-z’]+", train_text.lower()))
speakers = set(SPEAKER.findall(train_text))


def judge(sample):
    """What a reader checks first: real words, real speaker names, variety."""
    words = re.findall(r"[a-z’]+", sample.lower())
    names = SPEAKER.findall(sample)
    return {"words": len(words),
            "real words": round(sum(w in train_words for w in words) / max(len(words), 1), 2),
            "different words": round(len(set(words)) / max(len(words), 1), 2),
            "speaker lines": len(names),
            "real names": round(sum(n in speakers for n in names) / max(len(names), 1), 2)}


rows = {"valid text (real Shakespeare)": judge(valid_text[:1000]),
        "pair counts (bigram)": judge(bigram_generate("ROMEO.\n", length=1000))}
rows.update({f"transformer, step {s}": judge(samples[s]) for s in SAMPLE_AT[1:]})
rows.update({f"temperature {t}": judge(generate(model, "ROMEO.\n", length=1000, temperature=t, seed=0))
             for t in (0, 0.5, 1.5)})
for name, row in rows.items():
    print(f"{name:30s}", row)

The output:

valid text (real Shakespeare)  {'words': 193, 'real words': 0.95, 'different words': 0.6, 'speaker lines': 7, 'real names': 1.0}
pair counts (bigram)           {'words': 184, 'real words': 0.18, 'different words': 0.79, 'speaker lines': 3, 'real names': 0.33}
transformer, step 300          {'words': 201, 'real words': 0.49, 'different words': 0.78, 'speaker lines': 3, 'real names': 0.33}
transformer, step 1000         {'words': 204, 'real words': 0.78, 'different words': 0.65, 'speaker lines': 8, 'real names': 0.75}
transformer, step 3000         {'words': 202, 'real words': 0.88, 'different words': 0.59, 'speaker lines': 6, 'real names': 0.83}
temperature 0                  {'words': 251, 'real words': 1.0, 'different words': 0.04, 'speaker lines': 1, 'real names': 1.0}
temperature 0.5                {'words': 206, 'real words': 0.95, 'different words': 0.5, 'speaker lines': 7, 'real names': 1.0}
temperature 1.5                {'words': 155, 'real words': 0.54, 'different words': 0.92, 'speaker lines': 7, 'real names': 0.43}

• Share of real words
• Share of different words

One sample of 1,000 characters at each temperature, from the trained model (0.8 is the end-of-training sample). Low temperature: real words, repeated. High temperature: varied words, many of them misspelled. Real Shakespeare has 95% and 60%.

Text description of the chart

A line chart with the temperature, 0 to 1.5, on the horizontal axis and a share from 0% to 100% on the vertical axis. The line for real words starts at 100% at temperature 0, then 95% at 0.5, 88% at 0.8 and 54% at 1.5. The line for different words starts at 4% at temperature 0, then 50% at 0.5, 59% at 0.8 and 92% at 1.5. The two lines cross between 0.8 and 1.5. For comparison, the real validation text has 95% real words and 60% different words in its first 1,000 characters.

What this tells us. These are single samples of about 200 words each, so treat the numbers as rough. Still, the pattern is clear.

What the model gets right:

• Line shapes and layout. Speeches, blank lines, names in capitals with a full stop, and stage directions in [_ … _], already at step 300.
• Spelling of common words. At step 3,000 and temperature 0.8, 88% of its words are real words of the plays, against 18% for the pair counts. At temperature 0.5, 95%, the same as real Shakespeare.
• Punctuation and old forms. Commas at line ends, “’twas”, “lov’d”, “thou”.
• Variety at temperature 0.8. 59% different words, close to real Shakespeare’s 60%.

What it gets wrong:

• Meaning. No sentence says anything. “Thou dock thou do my lov’d ’twas bid mad be sea a cannot conce” sounds like a play and means nothing.
• Who speaks. It mixes characters of different plays, and about one speaker name in six is invented.
• Facts. It has no knowledge to look up, only patterns of characters. Ben asks about a fact in the plays: Macbeth killed King Duncan.

question = "HORATIO.\nWho killed King Duncan, and why?\n\nHAMLET.\n"
print(generate(model, question, length=150, temperature=0.5, seed=0))

The output:

HORATIO.
Who killed King Duncan, and why?

HAMLET.
The most long and did make the place of the thou worst the traitor man the content to me.

BRUTUS.
What what he see a can a like which is not place so

What this tells us. The model continues the layout perfectly: Hamlet replies, then Brutus from another play speaks. “Traitor” fits the mood. There is no answer, because the model only continues text. A large language model can answer because it saw trillions of tokens and was then tuned to follow questions. Both are missing here.

The final model and an end-to-end prediction

The model at step 3,000 is the final model: the validation loss was still falling, so the last step is the best one. For the exhibit, it must work outside the notebook. Ben saves the weights and everything needed to use them, as in Save and inspect results:

import hashlib, json

torch.save(model.state_dict(), "tiny_lm.pt")
info = {"chars": chars, "config": {"dim": 128, "heads": 4, "blocks": 3, "context": 64, "dropout": 0.1},
        "text": "shakespeare_tragedies.txt",
        "text_sha256": hashlib.sha256(Path("shakespeare_tragedies.txt").read_bytes()).hexdigest(),
        "split": "first 90% train, last 10% valid", "steps": STEPS, "seed": 0,
        "valid_loss": round(results["tiny transformer"], 4), "torch": torch.__version__}
Path("tiny_lm.json").write_text(json.dumps(info, ensure_ascii=False, indent=1), encoding="utf-8")
for f in ("tiny_lm.pt", "tiny_lm.json"):
    print(f"{f:13s} {Path(f).stat().st_size / 1000:,.1f} kB")

The output:

tiny_lm.pt    2,504.4 kB
tiny_lm.json  0.9 kB

The weights are 2.5 MB: 622,539 parameters × 4 bytes. The vocabulary must travel with them: a list in another order would turn every sample into nonsense, with no error. Then Ben writes generate.py for the exhibit computer. It holds its own copy of the TinyLM class, because a state_dict holds only numbers, not code:

generate.py

"""Write text with the saved tiny language model, in a new process.

    python generate.py "PROMPT" [--length 200] [--temperature 0.8] [--seed 0]
"""
import argparse
import json
import sys

import torch
from torch import nn

class TinyLM(nn.Module):
    ...                                    # the same class as above

parser = argparse.ArgumentParser()
parser.add_argument("prompt")
parser.add_argument("--length", type=int, default=200)
parser.add_argument("--temperature", type=float, default=0.8)
parser.add_argument("--seed", type=int, default=0)
args = parser.parse_args()

torch.set_num_threads(2)
info = json.load(open("tiny_lm.json", encoding="utf-8"))
chars = info["chars"]
stoi = {c: i for i, c in enumerate(chars)}
unknown = sorted(set(args.prompt) - set(chars))
if unknown:
    sys.exit(f"These characters are not in the vocabulary: {unknown}")

model = TinyLM(len(chars), **info["config"])
model.load_state_dict(torch.load("tiny_lm.pt"))
model.eval()

g = torch.Generator().manual_seed(args.seed)
ids = [stoi[c] for c in args.prompt]
with torch.no_grad():
    for _ in range(args.length):
        scores = model(torch.tensor([ids[-model.context:]]))[0, -1]
        if args.temperature == 0:
            next_id = int(scores.argmax())
        else:
            probs = torch.softmax(scores / args.temperature, dim=-1)
            next_id = int(torch.multinomial(probs, 1, generator=g))
        ids.append(next_id)
print("".join(chars[i] for i in ids))

The notebook writes the full file. Now a new prompt, in a fresh Python process, compared with the notebook’s own generate():

import subprocess, sys

prompt = "MACBETH.\nWhat news"
run = subprocess.run([sys.executable, "generate.py", prompt, "--length", "150", "--seed", "1"],
                     capture_output=True, text=True, check=True)        # a fresh Python process
print(run.stdout)
print("same as in this notebook:", run.stdout.rstrip("\n") == generate(model, prompt, length=150, seed=1))

The output:

MACBETH.
What news when, sir, they blood the bosiness,
That when daughters in place to the head be watch the done fight,
Look my sight father, she is positites of thee,

same as in this notebook: True

What this tells us. Two files and one program are enough. A new process rebuilds the model and writes exactly the same text as the notebook, for the same prompt and seed. In the notebook, the last task asks you to do the same with your own prompt.

Common mistake

A prompt with a character the model never saw

Symptom: a visitor types ROMEO., a new line and What's the news? on a normal keyboard. The program stops:

exit code 1 | These characters are not in the vocabulary: ["'"]

Cause: the plays use the curly apostrophe ’. The straight one ' on the keyboard never appears in the text, so it has no ID and no row in the character table. The same is true for digits and the straight double quote. A character model has no <unk>: it can read only the 75 characters it was trained on.

Fix: name the missing characters, as generate.py does, instead of crashing with a KeyError. For the exhibit, Ben replaces ' with ’ before generating. Do not silently drop unknown characters: the prompt would change without anyone knowing.

What it means for a real AI system: every model reads only what its tokenizer can encode. Large models use pieces down to single bytes, so that no input is impossible, as From text to tokens shows. Even then, text that is rare in the training data is read worse.

How far is this from a real LLM?

Nadia’s last condition is the sign in the exhibit. Ben compares his model with SmolLM2-135M, the smallest model of a real family, which you met in Training versus inference. Its pretraining numbers come from its model card, as quoted in Training objectives.

	This tiny model	SmolLM2-135M	How many times more
Parameters	622,539	134,515,008	about 216
Transformer blocks	3	30	10
Numbers per token	128	576	4.5
Vocabulary	75 characters	49,152 subword tokens	about 655
Context (tokens it reads at once)	64 characters	8,192 tokens	128
Training predictions	6,144,000 (765,787 characters, each read about 8 times)	2 trillion tokens	about 325,000
Training hardware	2 CPU threads, about 3 minutes	64 H100 data-centre GPUs	—
After pretraining	nothing	nothing: it is a base model. Chat models add instruction tuning.	—

The training row compares predictions: 3,000 steps × 32 windows × 64 characters for the tiny model, 2 trillion tokens for SmolLM2. Each SmolLM2 token is usually several characters, so in text the gap is even larger.

What this tells us. SmolLM2-135M is small for a modern model, and it is still about 216 times larger and trained on about 325,000 times more examples. Even so, Module 7 showed that it rambles and cannot route a ticket from a prompt. The tiny model shows the mechanism of a large language model honestly: the same objective, the same causal mask and the same generation loop. It does not show the ability, which comes from scale and from the training after pretraining. Ben’s sign says: “This model learned the look of a Shakespeare play in three minutes. A chatbot uses the same idea, with hundreds of thousands of times more text and far more parameters.”

Insights and limits

• Shape first, spelling next, meaning never. Layout by step 300 (about 20 seconds), common words by step 1,000, old forms by step 3,000. A few more minutes would not bring meaning or facts.
• Context beats counting: 1.585 against 2.499 for pair counts. Temperature trades correctness against variety: 95% real words at 0.5, 54% at 1.5.
• Limits. There is one sample per setting, so the word and name shares are rough. Validation is the second half of one play that shares names with training, so it measures “more of the same”, not English in general. There is no test set. The time budget stopped training, not the curve. The numbers come from one computer and one seed. Another computer can give slightly different losses, and its samples can look different from the first changed draw on.
• Next steps. Train longer, and stop when the validation loss stops falling. Try a context of 128 characters, and compare the loss and the time. Train on a different author and compare the samples.

Responsible use

Whose text is it? The plays are Shakespeare’s, and they are public domain: anyone may copy them, train on them and show the results. The museum’s sign still names the source: “Trained on six tragedies by William Shakespeare, from Project Gutenberg.” Public domain is a legal status, not a reason to hide where the text came from.

What such a model can be used for: teaching and exhibits like Nadia’s; experiments with model size, temperature and context; and art, clearly labelled as machine-made.

What it must not be used for:

• Not as Shakespeare. Its text is not “lost Shakespeare” or a quotation. Label every sample as generated.
• Not for answers. It has no facts, as the Duncan question showed. It must never sit behind a question box that looks like a chatbot.
• Not on a living author’s text without permission. The same 3-minute recipe on a living author’s books or posts would imitate a person’s style. That raises copyright, consent and impersonation questions that public-domain text avoids. Check the licence of any text first, as in Licensing and data terms in the large language models course.
• Not as an unfiltered public generator. The plays contain violence, insults and the prejudices of their time, and the model recombines them. For a public screen, a person reviews the fixed samples first. The live generator starts only from prompts that the staff chose.

Make it a portfolio piece

This case study shows that you understand a language model from the inside. Make it yours with one change, such as another public-domain author, a context of 128 or a longer run, and publish it.

What to put in the repository:

• README.md (template below), with the loss chart and the four samples.
• The notebook, run from top to bottom, with its outputs.
• generate.py and tiny_lm.json. Add tiny_lm.pt (2.5 MB) only with the PyTorch version that saved it; a reader can also rebuild it with the notebook in a few minutes.
• The text’s source and status, with its SHA-256: “Six tragedies by William Shakespeare, from Project Gutenberg eBook #100. Public domain in the USA. Project Gutenberg header and licence removed.” Do not put the Project Gutenberg name inside a changed copy of the text itself.

A README template:

# A tiny character-level language model, trained on Shakespeare on a CPU

## Question
What can a 0.6-million-parameter decoder learn from six plays in 3 minutes on a CPU?

## Data and status
Six Shakespeare tragedies (850,875 characters), Project Gutenberg eBook #100,
public domain in the USA; header and licence removed. SHA-256 59bd62e6...94b4.
Split: first 90% train, last 10% validation (one cut, no shuffling).

## Model
Characters as tokens (75). Character + position embeddings, 3 transformer
blocks (128 dims, 4 heads) with a causal mask, 622,539 parameters.
AdamW, lr 0.002, 3,000 steps of 32 x 64 characters, seed 0.

## Results (validation loss per character)
Equally likely 4.317 | character frequencies 3.313 | pair counts 2.499 |
tiny transformer 1.585. Samples at steps 0, 300, 1,000, 3,000 in the notebook.

## What it learned, and what it did not
Layout by step 300; 88% real words at step 3,000 (temperature 0.8).
No meaning, no facts. Compared with SmolLM2-135M: 216x fewer parameters,
about 325,000x fewer training predictions.

## Limits
One seed, one computer; validation is the end of one play; no test set.

How to describe it in an interview, in two sentences: “I trained a 600,000-parameter character-level transformer from scratch on six Shakespeare plays, in three minutes on a CPU. It beat pair counts on held-out text by a wide margin, and I measured what it does not learn, meaning and facts, against a real small language model.”

Go further:

• Train 6,000 steps and plot the curves. Predict first: will the validation loss still fall, and will the training loss pull away from it?
• Change the context to 16 characters, then to 128. How do the loss, the time and the samples change?
• Load SmolLM2-135M and compute its loss on the validation text per character (its total loss divided by the number of characters, not tokens). Compare it with 1.585.

Knowledge check

Check your understanding

Ben splits the text with one cut, not by shuffling the lines. Why? What would go wrong with a random 10% of the lines as validation?

Show the answer

Neighbouring lines depend on each other: a verse line continues the line before it, and a speech answers the one before. With randomly chosen lines, most validation lines would have their neighbours in training. The model would have seen almost the same context. So the validation loss would measure memory of these scenes, not the skill of writing new text, and it would look too good. One cut keeps validation a block that the model has never read. It is the same leakage as a test set that shares information with training, as in the machine learning course.

Check your understanding

The tiny transformer has a validation loss of 1.585, the pair counts 2.499. Explain what the two numbers mean in plain words, and why the transformer can be so much better.

Show the answer

The loss is the average surprise at the real next character, −ln p. e to the power of the loss is a number of equally likely choices. It is about 12 characters for the pair counts and about 5 for the transformer. The pair counts see only the current character, so after “Guildenst” they know only “t”. The transformer reads up to 64 characters back through attention with the causal mask. So it can use the whole word, the line and the speaker’s name. That context lowers its uncertainty.

Check your understanding

A visitor says: this model writes like Shakespeare, so a chatbot is just this, but bigger. What is right and what is missing in that sentence?

Show the answer

Right: a chatbot uses the same mechanism, a decoder trained with next-token prediction, with a causal mask and the same generation loop. Missing: scale and further training. SmolLM2-135M, a small real model, has about 216 times more parameters. It trained on 2 trillion tokens, about 325,000 times this model’s predictions, on 64 data-centre GPUs. A chatbot is far larger still, and it is tuned after pretraining to follow instructions. The tiny model learned layout and spelling, but no meaning and no facts. It answered a question about King Duncan with a speech that says nothing.

Recap

• You built a character-level decoder with 622,539 parameters from the course’s parts: a character table, a position table, 3 blocks with the causal mask and an output layer.
• You split the text with one cut, because neighbouring lines leak, and built the vocabulary from training only.
• Counting baselines set the bar: 3.313 for character frequencies, 2.499 for pair counts. The transformer reached 1.585 on the held-out text after about 3 minutes on a CPU.
• The samples showed the order of learning: layout by step 300, words by step 1,000, old forms by step 3,000. Meaning and facts never came. Temperature trades real words against variety.
• You saved the weights with the vocabulary and settings, and reloaded them in a fresh process with the same result.
• Compared with SmolLM2-135M: about 216 times fewer parameters and about 325,000 times fewer training predictions. That gap, not a different idea, separates this model from a real LLM.

Next, the final assignment brings the course together in a project of your own.

Next: Final assignment