Pre-training: preparing the inputs

Once we have converted our text corpus to tokens using our tokeniser, we then extract a collection of inputs and targets which we’ll use to train our LLM. Each input is a sequence of tokens, and each target is the same sequence shifted by one. For example, if our starting text is:

“It was a bright cold day in April, and the clocks were striking thirteen. Winston Smith, his chin nuzzled into his breast in an effort to escape the vile wind, slipped quickly through the glass doors”

Then our inputs and targets might be:

Our full LLM implementation uses the following code to achieve this:

class NextTokenDataset(Dataset):
    """Split text into fixed-length input and target sequences."""

    def __init__(self, text, sequence_length):
        self.token_ids = torch.tensor(tokenise(text), dtype=torch.long)
        self.sequence_length = sequence_length

        # Each target is the corresponding input shifted one token to the left.
        # Non-overlapping starts keep this example small and easy to inspect.
        self.start_positions = list(
            range(
                0,
                len(self.token_ids) - sequence_length,
                sequence_length,
            )
        )

        if not self.start_positions:
            raise ValueError(
                "the input text must contain more tokens than sequence_length"
            )

    def __len__(self):
        return len(self.start_positions)

    def __getitem__(self, index):
        start = self.start_positions[index]
        stop = start + self.sequence_length
        inputs = self.token_ids[start:stop]
        targets = self.token_ids[start + 1 : stop + 1]
        return inputs, targets

For a given input and target pair, each token position in the input has the corresponding “next token” in the same position of the target:

At position 0 the model is asked to predict ␠was from It. At position 1 it predicts ␠a after seeing It, ␠was, and so on. The causal attention mask prevents each position from looking ahead at the target it is supposed to predict.

Group sequences into batches

The dataset returns one input/target pair at a time. A PyTorch DataLoader groups several pairs so the model can process them together:

def create_data_loader(text_path, sequence_length, batch_size):
    text = Path(text_path).read_text(encoding="utf-8")
    dataset = NextTokenDataset(text, sequence_length)
    return DataLoader(dataset, batch_size=batch_size, shuffle=True)

shuffle=True here just randomises the order in which each input/target pair is added to a batch if we go through the text corpus multiple times. A production LLM may have a more sophisticated approach to sampling from the text corpus, but the basic idea is exactly the same.