Lesson 1 - Batching & Data Loading

Welcome to Batching & Data Loading#

You’ve built the mini-GPT and proven every gradient correct. Before you can train it, you need to turn the corpus — one long ribbon of character ids — into the batches the model actually consumes. That is what this lesson builds: a small, fast data pipeline that samples random chunks of text, pairs each input with the next-token targets it should predict, and holds back a slice for validation.

It’s the least glamorous part of training and the easiest to get subtly wrong, so we’ll verify it carefully on the real Lantern Bay corpus.

By the end of this lesson, you will be able to:

• Explain why a training example is a fixed-length chunk of the token stream
• Build the input/target pair for a chunk by shifting the tokens by one
• Implement get_batch to sample a batch of chunks at random starting positions
• Split the corpus into training and validation sets and say what each is for

From One Long Stream to Fixed-Length Chunks#

Our corpus is a single 1-D array of token ids — 10,020 of them, over a vocabulary of 24 characters. A transformer, though, works on fixed-length sequences of length block_size (its context window). So a training example is just a contiguous slice of the stream: block_size tokens as the input, and the same slice shifted one position to the right as the targets.

That shift is the whole idea of language modeling from Module 6: the token at each position should predict the token that follows it. So for a chunk starting at index i:

• input x = data[i : i + block_size]
• target y = data[i+1 : i + block_size + 1]

By construction y[t] is the character that comes right after x[t]. Let’s confirm that on the real corpus.

import numpy as np

CORPUS = (
    "the lamp keeper walks the shore at dusk. "
    "the lamp keeper lights the lantern when the fog rolls in. "
    "the boats come home when the lantern glows. "
    "the boats wait outside the bay when the fog is thick. "
    "a small boat drifts near the rocks at dawn. "
    "the keeper rings the bell when a boat drifts near the rocks. "
    "the tide rises at dusk and falls at dawn. "
    "the gulls call over the bay when the tide is low. "
    "the keeper counts the boats and writes the count in a book. "
    "the lantern burns all night and rests at dawn. "
) * 20

chars = sorted(set(CORPUS))
vocab_size = len(chars)
stoi = {c: i for i, c in enumerate(chars)}
itos = {i: c for i, c in enumerate(chars)}
encode = lambda s: [stoi[c] for c in s]
decode = lambda ids: "".join(itos[i] for i in ids)
data = np.array(encode(CORPUS), dtype=np.int64)

print("vocab_size:", vocab_size, "| corpus length:", len(data))

vocab_size: 24 | corpus length: 10020

The vocabulary is the 24 distinct characters (lowercase letters, space, and a period), and the encoded corpus is 10,020 tokens long. This is the same tokenization every lesson in the course uses.

Sampling a Batch#

Training one chunk at a time would be slow and noisy. Instead we stack batch_size chunks — each from a different random starting position — into one array and process them together. Random starts matter: they decorrelate the examples in a batch and, over many steps, cover the whole corpus without us having to march through it in order.

The only bookkeeping is the last legal start index. To have a full target y, a chunk needs block_size + 1 tokens available, so the largest valid start is len(data) - block_size - 1.

def get_batch(data, block_size, batch_size, rng):
    ix = rng.integers(0, len(data) - block_size - 1, size=batch_size)
    x = np.stack([data[i : i + block_size] for i in ix])
    y = np.stack([data[i + 1 : i + block_size + 1] for i in ix])
    return x, y

rng = np.random.default_rng(42)
x, y = get_batch(data, block_size=8, batch_size=4, rng=rng)

print("x shape:", x.shape, "| y shape:", y.shape)
print("shift holds (x[:,1:] == y[:,:-1]):", bool((x[:, 1:] == y[:, :-1]).all()))
print("x[0]:", repr(decode(x[0])))
print("y[0]:", repr(decode(y[0])))

x shape: (4, 8) | y shape: (4, 8)
shift holds (x[:,1:] == y[:,:-1]): True
x[0]: '. the ke'
y[0]: ' the kee'

Both x and y are (batch_size, block_size) = (4, 8). The check x[:, 1:] == y[:, :-1] confirms that y is x slid one step left: reading the first row, the input '. the ke' is paired with the target ' the kee', so at each position the model is asked to predict the very next character ('.' → ' ', ' ' → 't', 't' → 'h', and so on). Because we passed a seeded rng, this batch is exactly reproducible.

Holding Out a Validation Set#

If we measure quality only on the data we train on, we can’t tell learning from memorizing. So we carve the corpus into a training set and a smaller validation set up front, and only ever sample training batches from the first. In Lesson 4 we’ll watch both losses to detect overfitting.

n = int(len(data) * 0.9)
train_data = data[:n]
val_data = data[n:]
print("train tokens:", len(train_data), "| val tokens:", len(val_data))

train tokens: 9018 | val tokens: 1002

A 90/10 split gives 9,018 training tokens and 1,002 held out. From here on, get_batch(train_data, ...) produces the batches we learn from, and get_batch(val_data, ...) produces the batches we only ever evaluate on.

Practice Exercises#

Exercise 1: Confirm the shift on a different slice#

Draw a batch with block_size=12, batch_size=2 from a fresh default_rng(0) and verify x[:, 1:] == y[:, :-1] is still True. Decode one row’s x and y to see the one-character shift.

Exercise 2: Count the tokens seen per step#

For batch_size=32, block_size=32, how many next-token prediction targets does one batch contain? Compute batch_size * block_size.

Exercise 3: The largest legal start index#

With len(data) = 10020 and block_size = 8, what is the largest index get_batch may sample, and why does rng.integers use len(data) - block_size - 1 as its exclusive upper bound?

Summary#

A training example is a fixed-length chunk of the token stream, and its targets are the same chunk shifted one position right — the next-token objective, made concrete. get_batch stacks batch_size such chunks from random starting positions into (batch_size, block_size) input and target arrays, and a seeded generator makes each batch reproducible. Splitting the corpus into training and validation sets up front is what will let us tell genuine learning from memorization.

Key Concepts#

• Chunk: a contiguous block_size-length slice of the corpus; the unit of a training example.
• Input/target shift: y = x shifted one token left, so position t predicts token t+1.
• Batch: batch_size chunks from random starts, stacked to shape (batch_size, block_size).
• Train/validation split: disjoint slices of the corpus; train batches teach, validation batches only measure.

Why This Matters#

The data pipeline is where silent training bugs hide: an off-by-one in the shift means the model is quietly trained to predict the current token instead of the next one, and the loss will still go down while the model learns nothing useful. Building get_batch yourself, and checking the shift explicitly, means the numbers you see in the next lessons are measuring the right thing. Random sampling with a held-out split is exactly how real language models are trained — only the corpus size changes.

Continue Building Your Skills#

You can now turn the corpus into an endless stream of reproducible training batches, each a stack of chunks paired with their next-token targets, with a validation set set aside. The next lesson builds the other half of a training step: the optimizer. You’ll implement Adam from scratch — its running averages of the gradient and its square, and the bias correction that makes its early steps sane — and see why it, not plain gradient descent, is what makes transformer training converge quickly.