Build GPT · 4 of 10 · Python

Batch next-token examples

A language model learns one skill: given some text, guess what comes next. Training examples for that cost nothing to make. Take any stretch of tokens as the input. The target is the same stretch moved one place to the right, so each position's target is the token that follows it. One window of 8 tokens gives 8 guesses to learn from at once.

Do the exercise in the app

Learn it

Words to know

context window: The number of tokens the model reads at one time.
next-token prediction: The training task of a language model: given the tokens so far, give a high score to the token that comes next.
batch: A group of examples that the model processes together in one step.

How it works

Pick a start position s and a window length context.
Input: the tokens from s up to, but not including, s + context.
Target: the tokens from s + 1 up to, but not including, s + context + 1. It is the input shifted by one.
Do this for every start position. Collect all the inputs in one list and all the targets in another.

Worked example

Input
tokens = [10, 20, 30, 40, 50]
starts = [0, 2]
context = 2

Output
[[[10, 20], [30, 40]], [[20, 30], [40, 50]]]

Start	Input	Target	What the model learns
0	[10, 20]	[20, 30]	after 10 comes 20; after 10, 20 comes 30
2	[30, 40]	[40, 50]	after 30 comes 40; after 30, 40 comes 50

	pos 0	pos 1	pos 2	pos 3	pos 4
tokens	10	20	30	40	50
input	10	20
target		20	30

The window at start 0. The target row holds the same tokens as the input row, one position later: 20 is the answer for 10, and 30 is the answer for 10, 20.

Each target is its input moved one place to the right. The last target token, 30 or 50, lies one past the end of the input window.

Why it matters

Every training step of the workshop's GPT starts by cutting a batch of windows like these from the training text, at random start positions.

Your turn

Cut input and target windows for next-token prediction out of a token sequence.

tokens: the token IDs of the training text
starts: the start position of each window; all valid, possibly none
context: the length of each window; positive
Returns: The list [inputs, targets]. Each holds one window per start, of length context; a target window begins one position after its input window.

Example 1

Input
tokens = [10, 20, 30, 40, 50]
starts = [0, 2]
context = 2

Output
[[[10, 20], [30, 40]], [[20, 30], [40, 50]]]

The window at start 0 is [10, 20] with target [20, 30]. The window at start 2 is [30, 40] with target [40, 50].

Example 2

Input
tokens = [0, 1, 2, 3]
starts = [0, 1]
context = 2

Output
[[[0, 1], [1, 2]], [[1, 2], [2, 3]]]

Starts 0 and 1 overlap. That is fine: each still pairs an input with the tokens one place later.

Example 3

Input
tokens = [0, 1]
starts = []
context = 1

Output
[[], []]

No starts, so both lists are empty.

Hints

Each hint gives away a little more. Open one only when you need it.

Hint 1: Where to start
A Python slice tokens[a:b] takes positions a up to, but not including, b. Each window is one slice.
Hint 2: The key step
The input slice begins at s. The target slice begins at s + 1. Both are context long.
Hint 3: The outline
Build inputs as tokens[s : s + context] for each s in starts, and targets as tokens[s + 1 : s + context + 1]. Return [inputs, targets].

Solution

From here on the page gives the method away. If you are still solving, go back to the hints.

Step by step

The window the model reads. A slice stops before its end index, so s : s + context holds exactly context tokens.
The same window one place later: what should come next. Position i of the target is the token that follows position i of the input. That pairing is the whole training signal.

The same numbered steps are the comments in the solution.

Python

def solve(tokens, starts, context):
    inputs = []
    targets = []
    for s in starts:
        # 1. The window the model reads
        inputs.append(tokens[s : s + context])
        # 2. The same window one place later: what should come next
        targets.append(tokens[s + 1 : s + context + 1])
    return [inputs, targets]

Watch out

A target equal to the input teaches the model to copy, not to predict. The shift by one is the lesson.
The target needs one token past the end of the input, so the last valid start is len(tokens) - context - 1.

Check yourself

How is the target window related to the input window?

It is the same stretch of tokens moved one position to the right.

Why does one window give several training examples?

Every position has its own target: the next token. A window of 8 tokens gives 8 predictions to learn from.