Data Pipeline

This page follows the text from the CSV file to the batches the training loop reads. None of these steps has a backward pass: gradients stop at the embedding table, and the ids have no gradient.

The order in train.main

train.main runs these lines in this order, before the model exists. The comment in the code says the same: split BEFORE windowing.

train.pymainGitHub

40    text = load_corpus(config.corpus_path, config.corpus_column, config.max_rows)

train.pymainGitHub

41    tok = build_tokenizer(text, config.num_merges)

train.pymainGitHub

42    ids = encode_corpus(text, tok)

train.pymainGitHub

43    train_ids, val_ids = split_ids(ids, config.val_fraction)

train.pymainGitHub

44    tr_x, tr_y = make_windows(train_ids, config.seq_len, config.stride)45    va_x, va_y = make_windows(val_ids, config.seq_len, config.stride)

The tokenizer is built on the whole text, including the part that later becomes validation. The split separates the windows, not the vocabulary or the merges.

Next-token prediction

Training has one objective: at each position of a window, predict the next token. For input x[t], the correct answer is y[t], the id right after it in the stream. A causal mask in attention stops a position from reading later ones, so the prediction at t can only use x[0] to x[t]. Softmax & Cross-Entropy covers how the model is scored on each answer.

Note

In general, the probability of a whole sequence can be written as a product, over positions, of each token's probability given the tokens before it. The PDF states this as Eq. 61. This code never computes that product. It computes a mean cross-entropy over windows of at most 16 tokens.

Text to ids

load_corpus reads one column of a CSV file. In config.py the file is "data/corpus.csv", the column is "text", and at most 2000 rows are read.

A column that does not exist raises ValueError.
max_rows limits the rows read from the file, before missing values are dropped.
Each remaining row is stripped of surrounding whitespace.
All rows are joined with newlines, and one final newline is added.

The result is a single string. Rows have no other boundary in it, so a window can run from the end of one row into the start of the next, and the model sees the newline as a token.

From rows to one textIllustrative ExampleHand-picked rows, not from the corpus.

Rows	Text returned
`[" hello ",null,"world"]`	`"hello\nworld\n"`

Code listing: data.load_corpus

data.pyload_corpusGitHub

5def load_corpus(path, column, max_rows) -> str:6    df = pd.read_csv(path, nrows=max_rows)7 8    if column not in df.columns:9        raise ValueError(f"no column {column!r}, the file has {list(df.columns)}")10 11    texts = df[column].dropna().astype(str).str.strip()12    return "\n".join(texts) + "\n"

encode_corpus passes that string to the tokenizer and stores the ids in a one-dimensional int64 array of N ids. Tokenization covers encode.

Code listing: data.encode_corpus

data.pyencode_corpusGitHub

15def encode_corpus(text: str, tok: Tokenizer) -> np.ndarray:16    return np.array(tok.encode(text), dtype=np.int64)

The split

What is it?

The splitImplemented

Data

split_ids cuts the id array at one position. Everything before the cut is for training and everything after it is for validation. Nothing is shuffled, so both parts stay in order.

Why does the model need it?

Validation loss is only useful on text the model did not train on. Because the cut happens before windows exist, no window can hold ids from both sides.

Windows overlap, so the order matters. If train.main made windows first and split them afterwards, the windows around the cut would hold ids from both sides, and validation would partly test on training text. Splitting first removes those windows, at the cost of fewer windows overall.

Mathematics

Where the stream is cut #

\text{cut} = \lfloor (1 - \rho)\, N \rfloor

What does each symbol mean?

Symbol	Meaning	In code	Shape
$\rho$	Fraction kept for validation	`val_fraction`	scalar
$N$	Number of ids in the stream	`len(ids)`	scalar
$\text{cut}$	Number of training ids	`cut`	scalar

Tiny numerical example

The last column counts the windows that would hold ids from both sides if the order were reversed. At stride 1 it equals T in every row.

Splitting before windowingIllustrative ExampleHand-picked lengths, with stride 1. These are not counts from the corpus. The last row uses the N of tests/test_data.py with the settings in config.py.

N	val_fraction	T	cut	Train windows	Validation windows	Windows crossing the cut, if windowing came first
20	0.25	3	15	12	2	3
40	0.2	4	32	28	4	4
3,000	0.1	16	2,700	2,684	284	16

Tensor shapes

ids has shape (N,), train_ids has cut elements and val_ids has N - cut. In config.py, val_fraction is 0.1.

Actual implementation

data.pysplit_idsGitHub

19def split_ids(ids, val_fraction):20    cut = int((1 - val_fraction) * len(ids))21    return ids[:cut], ids[cut:]

Code and math connection

Equation	Role	Lines	Code
Where the stream is cut	Computes it	data.py:20	`cut = int((1 - val_fraction) * len(ids))`
Where the stream is cut	Computes it	data.py:21	`return ids[:cut], ids[cut:]`
Where the stream is cut	Calls it	train.py:43	`train_ids, val_ids = split_ids(ids, config.val_fraction)`
Where the stream is cut	Tests it	tests/test_data.py:37	`assert len(train) == 90 and len(val) == 10`

Where is it used?

Source

data.py, lines 19-21, split_ids

PDF reference

Section 4.2, p. 6

Used by

Called by: train.main

Tested by: tests.test_data.test_split_is_contiguous_and_disjoint, tests.test_data.test_window_counts_for_n_3000

Calls it: train.py, line 43 (main), for Where the stream is cut
Tests it: tests/test_data.py, line 37 (test_split_is_contiguous_and_disjoint), for Where the stream is cut

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Where the stream is cut	Computes it	`data.py`	`split_ids`	20	GitHub: data.py, line 20
Where the stream is cut	Computes it	`data.py`	`split_ids`	21	GitHub: data.py, line 21

Sliding windows

What is it?

Sliding windowsImplemented

Data

make_windows slides a window of T ids along the stream, moving k ids each time. The inputs are the window. The targets are the window that starts one id later, so each input row moves one place left and gains one new id at the end.

Why does the model need it?

The model reads a fixed number of tokens at once, and training needs a correct answer for every position. Shifting by one makes each target the id that comes right after the input at that position.

Mathematics

Where windows start Derived from implementation#

s \in \{0,\, k,\, 2k,\, \dots\}, \qquad s < N - T

The last target of a window is ids[s + T], which must exist. That is why s stops before N - T.

Inputs #

x^{(s)}_{t} = \text{ids}_{s+t}, \qquad t = 0, \dots, T-1

Targets #

y^{(s)}_{t} = \text{ids}_{s+t+1}, \qquad t = 0, \dots, T-1

Targets are inputs shifted by one #

y^{(s)}_{t} = x^{(s)}_{t+1}, \qquad t = 0, \dots, T-2

What does each symbol mean?

Symbol	Meaning	In code	Shape
$s$	Start position of a window	`starts`	(W,) for all of them
$t$	Position inside the window	`np.arange(seq_len)`	scalar
$k$	How far the window moves each time	`stride`	scalar
$T$	Window length	`seq_len`	scalar
$N$	Number of ids given to make_windows	`len(ids)`	scalar
$x$	Inputs, one row per window	`inputs`	(W, T)
$y$	Targets, one row per window	`targets`	(W, T)

Tiny numerical example

Edit the ids, the window length or the stride.

Sliding windowsIllustrative ExampleHand-picked ids that you can edit. They are not ids from the corpus. Tests compare the engine behind this table with the real make_windows.

Token ids

Window length T

Stride k

N = 10, W = ceil((10 - 4) / 1) = 6

Start s	Inputs	Targets
0	`3, 1, 4, 1`	`1, 4, 1, 5`
1	`1, 4, 1, 5`	`4, 1, 5, 2`
2	`4, 1, 5, 2`	`1, 5, 2, 6`
3	`1, 5, 2, 6`	`5, 2, 6, 5`
4	`5, 2, 6, 5`	`2, 6, 5, 3`
5	`2, 6, 5, 3`	`6, 5, 3, 5`

In every row, the targets are the inputs moved one place left, with one new id at the end.

Tensor shapes

train.main uses 16 as T and 1 as k. At stride 1, neighbouring windows share all but one id.

Inputs of the training windowsdata.py:31

train_ids(N,)
tr_x(W, T)(W, 16)

N is the number of ids. It depends on the corpus. W is the number of windows. It depends on the corpus. T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt.

Targets of the training windowsdata.py:32

train_ids(N,)
tr_y(W, T)(W, 16)

make_windows returns two pandas DataFrames, with one row per window and one column per position. They are not NumPy arrays until iterate_batches or validation_loss calls to_numpy. The two DataFrames share the same row index, which keeps each input row with its target row.

Actual implementation

data.pymake_windowsGitHub

24def make_windows(ids, seq_len, stride = 1):25    if len(ids) <= seq_len:26        raise ValueError(f"need more than seq_len={seq_len} tokens, got {len(ids)}")27 28    starts = np.arange(0, len(ids) - seq_len, stride)29    rows = starts[:, None] + np.arange(seq_len)30 31    inputs = pd.DataFrame(ids[rows])32    targets = pd.DataFrame(ids[rows + 1])33    return inputs, targets

Code and math connection

Equation	Role	Lines	Code
Where windows start	Computes it	data.py:28	`starts = np.arange(0, len(ids) - seq_len, stride)`
Inputs	Computes it	data.py:29	`rows = starts[:, None] + np.arange(seq_len)`
Inputs	Computes it	data.py:31	`inputs = pd.DataFrame(ids[rows])`
Targets	Computes it	data.py:32	`targets = pd.DataFrame(ids[rows + 1])`
Targets are inputs shifted by one	Computes it	data.py:32	`targets = pd.DataFrame(ids[rows + 1])`
Targets are inputs shifted by one	Tests it	tests/test_data.py:21	`assert (y.to_numpy()[:, :-1] == x.to_numpy()[:, 1:]).all()`

Where is it used?

Source

data.py, lines 24-33, make_windows

PDF reference

Section 4.2, Eq. 17, p. 6

Used by

Called by: train.main

Tested by: tests.test_data.test_batches_cover_every_window_once, tests.test_data.test_no_shuffle_keeps_order, tests.test_data.test_stride_skips_windows, tests.test_data.test_targets_are_inputs_shifted_by_one, tests.test_data.test_too_few_tokens_raises, tests.test_data.test_window_counts_for_n_3000, tests.test_data.test_window_table_matches_hand_example

Tests it: tests/test_data.py, line 21 (test_targets_are_inputs_shifted_by_one), for Targets are inputs shifted by one

Status

Implemented

Implemented in

Equation	Role	File	Function	Lines	GitHub
Where windows start	Computes it	`data.py`	`make_windows`	28	GitHub: data.py, line 28
Inputs	Computes it	`data.py`	`make_windows`	29	GitHub: data.py, line 29
Inputs	Computes it	`data.py`	`make_windows`	31	GitHub: data.py, line 31
Targets	Computes it	`data.py`	`make_windows`	32	GitHub: data.py, line 32
Targets are inputs shifted by one	Computes it	`data.py`	`make_windows`	32	GitHub: data.py, line 32

Window count

What is it?

Window countDerived from implementation

Data

The number of windows is the number of start positions. The code never computes this formula. It follows from the array of starts that make_windows creates.

Why does the model need it?

The window count sets how many batches make up an epoch, and so how many steps training runs. The learning-rate schedule needs that total.

Mathematics

np.arange(0, N - T, k) lists the multiples of k from 0 up to, but not including, N - T. Counting them gives the formula.

Number of windows #

W = \left\lceil \frac{N - T}{k} \right\rceil, \qquad N > T

If N is not greater than T, make_windows raises ValueError before it creates any start. With a stride above 1, the ids after the last window may go unused.

What does each symbol mean?

Symbol	Meaning	In code	Shape
$W$	Number of windows	`len(inputs)`	scalar
$N$	Number of ids given to make_windows	`len(ids)`	scalar
$T$	Window length	`seq_len`	scalar
$k$	Stride	`stride`	scalar

Tiny numerical example

The formula and make_windowsIllustrative ExampleHand-picked lengths. The last column is the number of rows the real make_windows returned.

N	T	k	Formula	make_windows
10	4	1	6	6
10	4	2	3	3
10	4	3	2	2
10	4	6	1	1
10	4	7	1	1
5	4	1	1	1
10	1	1	9	9
40	16	1	24	24

Tensor shapes

W is the number of rows of tr_x and tr_y, as shown under Sliding windows. It depends on the corpus, so it stays a symbol.

Actual implementation

The code is the starts line of make_windows, shown under Sliding windows. No function computes W from the formula.

Code and math connection

Equation	Role	Lines	Code
Number of windows	Computes it	data.py:28	`starts = np.arange(0, len(ids) - seq_len, stride)`
Number of windows	Computes it	data.py:25	`if len(ids) <= seq_len:`
Number of windows	Tests it	tests/test_data.py:44	`assert len(make_windows(train, 16)[0]) == 2684`

Where is it used?

Source

data.py, lines 24-33, make_windows

PDF reference

Section 4.2, Eq. 18, p. 6

Used by

Called by: train.main

Tests it: tests/test_data.py, line 44 (test_window_counts_for_n_3000), for Number of windows

Status

Derived from implementation

Implemented in

Equation	Role	File	Function	Lines	GitHub
Number of windows	Computes it	`data.py`	`make_windows`	28	GitHub: data.py, line 28
Number of windows	Computes it	`data.py`	`make_windows`	25	GitHub: data.py, line 25

Shuffling and batching

What is it?

Shuffling and batchingDerived from implementation

Data

iterate_batches puts the windows in a random order, using the random generator it is given, and hands them out in groups of batch_size. A shorter last batch is kept.

train.main creates one generator from 42 (train.py:37) and uses it for everything random. iterate_batches receives it once per epoch and passes it to pandas, which moves the generator along. So each epoch gets a new order, and a repeated run with the same settings gets the same orders.

Why does the model need it?

One update from all the windows would be slow, and one update from a single window would be noisy. A new order each epoch changes which windows are averaged together.

Mathematics

Batches in one epoch #

\text{batches per epoch} = \left\lceil \frac{W}{B} \right\rceil

What does each symbol mean?

Symbol	Meaning	In code	Shape
$W$	Number of windows	`len(tr_x)`	scalar
$B$	Windows in one batch	`batch_size`	scalar

Tiny numerical example

Two epochs of iterate_batchesIllustrative ExampleHand-picked sizes: 10 windows, batch_size 4, generator seed 42. The order is the stored result of running the real function, because it comes from pandas and NumPy. Each number is the start of one window. This is not the order of a real training run, which also depends on every earlier use of the generator.

Batch sizes: 4, 4, 2

Epoch	Batch 1	Batch 2	Batch 3
1	`2, 9, 1, 6`	`3, 8, 5, 7`	`4, 0`
2	`5, 9, 6, 2`	`0, 1, 7, 3`	`4, 8`

Tensor shapes

B is 16, except in a shorter last batch. The batch is a NumPy array, but the model has no batch axis: batch_loss_and_grads takes one row at a time (train.py:17). Training Loop covers how the results are averaged.

From all windows to one sequencedata.py:43

tr_x(W, T)(W, 16)
xb(B, T)(16, 16)
x(T,)(16,)

W is the number of windows. It depends on the corpus. T = seq_len = 16. T is at most seq_len: generation uses the last seq_len ids, or fewer for a short prompt. B = batch_size = 16. A final short batch has fewer rows.

Actual implementation

data.pyiterate_batchesGitHub

36def iterate_batches(inputs, targets, batch_size, rng, shuffle = True):37    order = inputs.index38    if shuffle:39        order = inputs.sample(frac=1, random_state=rng).index40 41    for i in range(0, len(order), batch_size):42        rows = order[i : i + batch_size]43        yield inputs.loc[rows].to_numpy(), targets.loc[rows].to_numpy()

train.pymainGitHub

59    for epoch in range(config.epochs):60        for xb, yb in iterate_batches(tr_x, tr_y, config.batch_size, rng):

Code and math connection

Equation	Role	Lines	Code
Batches in one epoch	Computes it	data.py:41	`for i in range(0, len(order), batch_size):`
Batches in one epoch	Computes it	data.py:39	`order = inputs.sample(frac=1, random_state=rng).index`
Batches in one epoch	Written inline	train.py:54	`steps_per_epoch = math.ceil(len(tr_x) / config.batch_size)`
Batches in one epoch	Tests it	tests/test_data.py:58	`assert sizes == [16, 16, 14]`

Where is it used?

Source

data.py, lines 36-43, iterate_batches

PDF reference

Not stated in the PDF

Used by

Called by: train.main

Tested by: tests.test_data.test_batches_cover_every_window_once, tests.test_data.test_no_shuffle_keeps_order

Written inline in train.py, line 54. The same count, computed for the learning-rate schedule.

Tests it: tests/test_data.py, line 58 (test_batches_cover_every_window_once), for Batches in one epoch

Status

Derived from implementation

Implemented in

NoteInline equivalent: the pipeline calls data.iterate_batches, and also computes the same formula in place, with no call, at train.py:54.

Equation	Role	File	Function	Lines	GitHub
Batches in one epoch	Computes it	`data.py`	`iterate_batches`	41	GitHub: data.py, line 41
Batches in one epoch	Computes it	`data.py`	`iterate_batches`	39	GitHub: data.py, line 39
Batches in one epoch	Written inline	`train.py`	`main`	54	GitHub: train.py, line 54

Every symbol in data.py

Symbol	Lines	On this page
`data.load_corpus`	data.py:5-12	Text to ids
`data.encode_corpus`	data.py:15-16	Text to ids
`data.split_ids`	data.py:19-21	The split
`data.make_windows`	data.py:24-33	Sliding windows
`data.iterate_batches`	data.py:36-43	Shuffling and batching