Tiny Shakespeare Data Pipeline

Before a language model can learn anything, its training text has to become numbers in a shape a neural network can consume. This first phase of the walkthrough builds that pipeline for the tiny-shakespeare text: download it, find its alphabet, turn every character into an integer, hold out a validation set, and cut the result into batches of short chunks. Every number below was computed from the text itself with tools/experiments/gpt/dataset_stats.py in this site's repository, which reproduces what the step files print.

The code for this phase is in steps/1_getting_started of the repository. Each step file is a complete program; each code block below is the difference between one step file and the next.

Project layout: one script for the whole model

Production machine-learning projects split configuration, data loading, models and training across many files. This walkthrough deliberately keeps everything in a single script, train_gpt.py, inside a gpt_text_model folder, with the training text stored under data/local/text/. Everything from the tokenizer to the final Transformer lives in that one file. The monolithic layout is unusual, but it keeps every component visible and makes the interactions between them explicit: as the model grows from a bigram table to a Transformer, each change appears in the same place as everything it touches.

Downloading the training text

The script keeps a small catalogue of public-domain texts, with tiny-shakespeare as the default, and downloads the chosen one only if it is not already on disk. The dataset name is taken from the command line, so python train_gpt.py moby_dick trains on Moby-Dick instead. Step file: train_gpt_1_3.py.

# Allow dataset selection via command-line argument
if len(sys.argv) > 1:
    dataset_name = sys.argv[1].lower()  # Get dataset name from command line
else:
    dataset_name = "shakespeare"  # Default dataset

# Check if the chosen dataset exists in the dictionary
if dataset_name not in datasets:
    raise ValueError(f"Dataset '{dataset_name}' is not available. Choose from {list(datasets.keys())}")

url = datasets[dataset_name]

# Use os.path.join to ensure cross-platform compatibility
data_dir = os.path.join("data", "local", "text")
data_file = os.path.join(data_dir, f"{dataset_name}.txt")

# Make sure data directory exists
os.makedirs(data_dir, exist_ok=True)

# Download the dataset if it's not already present
if not os.path.isfile(data_file):
    print(f"Downloading {dataset_name} dataset from {url}...")
    try:
        response = requests.get(url)
        response.raise_for_status()  # Check for successful request
        with open(data_file, 'w') as f:
            f.write(response.text)
        print(f"Dataset saved to {data_file}.")
    except requests.exceptions.RequestException as e:
        print(f"Error downloading dataset: {e}")
else:
    print(f"Dataset already exists at {data_file}.")

Three details matter in practice. The path is built with os.path.join, so it works on Windows and Unix alike. The path is relative to the directory the script is run from, so running it from inside gpt_text_model creates gpt_text_model/data/local/text/shakespeare.txt. And the file is written and later read with Python's default text encoding, which on Windows is not UTF-8; this is harmless for tiny-shakespeare, which is pure ASCII, but the later tokenized version passes encoding='utf-8' explicitly so that other texts survive the round trip. (The next step file trims the catalogue to two entries while developing; the finished train_gpt.py restores all nine.)

A quick check that the setup works is to run the script and look at the first lines of the downloaded file:

First Citizen:
Before we proceed any further, hear me speak.

All:
Speak, speak.

The file is a sequence of speeches, each a speaker's name and a colon on one line, followed by the lines spoken. It has 1,115,394 characters.

Building the character vocabulary

Modern language models split text into words or word pieces, a choice discussed in Tokenization. This model starts with the simplest possible scheme: every character is a token. The vocabulary is the sorted list of distinct characters in the text. Step file: train_gpt_1_4.py.

--- train_gpt_1_3.py
+++ train_gpt_1_4.py
@@ -53 +46,14 @@

+####################################################################################################
+# Load the Dataset and Create Vocabulary
+
+# Read the dataset into a string
+with open(data_file, 'r') as f:
+    text = f.read()
+    # Create a set of all unique characters in the text (vocabulary)
+    chars = sorted(list(set(text)))
+    vocab_size = len(chars)
+
+    # Print vocabulary size and characters
+    print(f"Vocabulary size: {vocab_size}")
+    print(f"Vocabulary: {chars}")

For tiny-shakespeare the vocabulary has 65 characters: the newline, the space, the punctuation marks ! $ & ' , - . : ; ?, the digit 3, and the 26 capital and 26 lower-case letters. Sorting puts the newline first (index 0) and the space second (index 1), because their character codes are the smallest.

Character-level tokens have two advantages for learning. They are transparent: the model's input and output can be read directly. And the vocabulary is closed: every string the model will ever see in this text is a sequence of these 65 symbols, so there is no "unknown word" problem. The price is length. A word-level vocabulary might have 50,000 entries but write a line in a dozen tokens, while characters need one token per letter, so the model must learn spelling before it can learn anything about words.

Encoding text as integers

A neural network computes with numbers, so each character is mapped to its index in the vocabulary, and back. Two dictionaries do the work: stoi (string to integer) and itos (integer to string). encode turns text into a list of integers, and decode reverses it. Step file: train_gpt_1_7.py.

--- train_gpt_1_4.py
+++ train_gpt_1_7.py
@@ -3,2 +3,5 @@
 import sys
+
+import torch
+

@@ -59 +62,31 @@
     print(f"Vocabulary: {chars}")
+
+####################################################################################################
+# Encode and Decode Text
+
+# Create encoder and decoder mappings
+stoi = {ch: i for i, ch in enumerate(chars)}  # String to integer
+itos = {i: ch for i, ch in enumerate(chars)}  # Integer to string
+
+# Encoder: Convert string to list of integers
+def encode(s):
+    return [stoi[c] for c in s]
+
+# Decoder: Convert list of integers back to string
+def decode(l):
+    return ''.join([itos[i] for i in l])
+
+# Test the encoding and decoding
+test_string = "hi there"
+encoded = encode(test_string)
+decoded = decode(encoded)
+
+print(f"Original: {test_string}")
+print(f"Encoded: {encoded}")
+print(f"Decoded: {decoded}")
+
+# Encode the entire dataset into a sequence of integers
+data = torch.tensor(encode(text), dtype=torch.long)
+
+# Print the first 1,000 elements of the tensor
+print(data[:1000])

The test string "hi there" encodes to [46, 47, 1, 58, 46, 43, 56, 43]: h is 46, i is 47, the space is 1, and so on, and decoding the list gives back the string. During training the model is fed encoded text; during generation its output is decoded.

The same step wraps the whole encoded text in a PyTorch tensor, torch.tensor(encode(text), dtype=torch.long): a one-dimensional array of 1,115,394 64-bit integers. Tensors are the arrays on which PyTorch performs its matrix operations, on the CPU or a GPU, and torch.long is the integer type that PyTorch's embedding layers expect as indices. With this, the pipeline from raw text to model input is: text, characters, integer indices, tensor.

Holding out a validation set

A model that only ever sees its training text could simply memorize it. To tell memorization from learning, a part of the data is set aside and never trained on. Step file: train_gpt_1_8.py.

--- train_gpt_1_7.py
+++ train_gpt_1_8.py
@@ -92 +92,13 @@
 print(data[:1000])
+
+####################################################################################################
+# Preparing the Validation Set
+
+# Split the dataset into training and validation sets
+n = int(0.9 * len(data))  # 90% of the data for training
+train_data = data[:n]
+val_data = data[n:]
+
+# Print the sizes of the training and validation sets
+print(f"Training data size: {len(train_data)} characters")
+print(f"Validation data size: {len(val_data)} characters")

The split is 90 percent for training and 10 percent for validation: 1,003,854 characters and 111,540 characters. The split is by position, not random: the validation set is the last tenth of the file, so it contains whole scenes the model never sees. The validation set acts as a final exam. When the loss on it starts rising while the training loss keeps falling, the model is fitting peculiarities of the training text rather than general patterns, the same gap between training and held-out performance described for classifiers in Evaluating Classifiers. For a text generator, the goal is new passages in Shakespeare's manner, not recited quotations.

Chunks of text and the context length

A Transformer is not trained on the whole text at once. It is trained on short chunks, and the chunk length, called block_size or the context length, is the most text the model can look at when predicting the next character. The walkthrough starts with a block size of 8 to keep the examples readable; the final model uses 256. Step file: train_gpt_1_9.py.

--- train_gpt_1_8.py
+++ train_gpt_1_9.py
@@ -104 +104,14 @@
 print(f"Validation data size: {len(val_data)} characters")
+
+####################################################################################################
+# Chunking the Data (Block Size)
+
+# Set block size (context length)
+block_size = 8
+
+# Get a sample chunk of the training set
+x = train_data[:block_size]  # Input chunk of size block_size
+y = train_data[1:block_size + 1]  # Target chunk, offset by one character
+
+print(f"Input chunk (x): {x}")
+print(f"Target chunk (y): {y}")

The input chunk x is the first 8 characters of the training data, and the target chunk y is the same window shifted one character to the right:

Every character of x is followed, in the text, by the character in the same position of y. Except at the chunk's edges, each character serves both as an input and as a target. The context length is a trade-off: too short, and the model cannot see enough to predict well; too long, and each training step costs more, since attention compares every position with every other.

Eight training examples in one chunk

A chunk of length 8 is not one training example but eight, one for each amount of context from one character to eight. Step file: train_gpt_1_10.py.

--- train_gpt_1_9.py
+++ train_gpt_1_10.py
@@ -117 +117,10 @@
 print(f"Target chunk (y): {y}")
+
+####################################################################################################
+# Creating Multiple Examples from one Chunk
+
+for t in range(block_size):
+    context = x[:t + 1]  # Context is all characters up to and including position t
+    target = y[t]  # Target is the t-th character in the target chunk
+    print(f"Context: {context} -> Target: {target}")
+

For the first chunk of the text, the eight examples are:

This does two things. It extracts as much training signal as possible from each chunk. And it trains the model to predict from any amount of context between one character and block_size, which is exactly what generation requires: generation starts from a single character and builds up. Short contexts teach which characters follow which; longer ones teach the shapes of words and short phrases.

Batching chunks for parallel training

Processing one chunk at a time would leave most of a GPU idle. A batch stacks several chunks into one tensor so that they are processed together, and the gradient used for each update is averaged over all of them, which also makes training steadier. Step file: train_gpt_1_11.py.

--- train_gpt_1_10.py
+++ train_gpt_1_11.py
@@ -126 +126,10 @@

+####################################################################################################
+# Batching the Data
+
+batch_size = 4  # Number of chunks processed in parallel
+x_batch = torch.stack([train_data[i:i + block_size] for i in range(0, batch_size * block_size, block_size)])
+y_batch = torch.stack([train_data[i + 1:i + block_size + 1] for i in range(0, batch_size * block_size, block_size)])
+
+print(f"Input batch (x_batch): {x_batch}")
+print(f"Target batch (y_batch): {y_batch}")

With batch_size = 4 and block_size = 8, the input batch is a tensor of shape (4, 8): four sequences of eight characters. This step takes four consecutive chunks from the start of the text, First Ci, tizen: plus a newline and B, efore we, and proceed. Each row has its own targets, shifted by one character as before, so one batch holds 4 × 8 = 32 prediction tasks. When the model processes this tensor, it computes all 32 predictions at once.

Sampling batches at random positions

Taking chunks in order would show the model the text in the same sequence every time. The last step of the pipeline samples each chunk's starting position at random from the whole training set, so each batch mixes passages from across the plays. Step file: train_gpt_1_12.py.

--- train_gpt_1_11.py
+++ train_gpt_1_12.py
@@ -135 +117,39 @@
 print(f"Target batch (y_batch): {y_batch}")
+
+####################################################################################################
+# Get a Batch of Data
+
+# Set random seed for reproducibility
+random.seed(42)
+
+# Function to get a batch of data
+def get_batch(split, block_size=8, batch_size=4):
+    """
+    Generate a batch of input-output pairs for training or validation.
+
+    Args:
+        split (str): 'train' or 'val', determines which dataset to use.
+        block_size (int): Length of each sequence in the batch.
+        batch_size (int): Number of sequences processed in parallel.
+
+    Returns:
+        x_batch (torch.Tensor): Input batch.
+        y_batch (torch.Tensor): Target batch.
+    """
+    data = train_data if split == 'train' else val_data
+
+    # Generate random starting positions for the batch
+    ix = [random.randint(0, len(data) - block_size - 1) for _ in range(batch_size)]
+
+    # Stack the input and target sequences
+    x_batch = torch.stack([data[i:i + block_size] for i in ix])
+    y_batch = torch.stack([data[i + 1:i + block_size + 1] for i in ix])
+
+    return x_batch, y_batch
+
+# Get a batch of training data
+xb, yb = get_batch('train')
+
+# Print the input and target batches
+print(f"Input batch (xb): {xb}")
+print(f"Target batch (yb): {yb}")

The function get_batch takes a split argument, 'train' or 'val', so the same code samples from either set. The random seed makes the sampling reproducible: with random.seed(42), the first batch starts at training positions 670,487, 116,739, 26,225 and 777,572, whose eight-character inputs are ings: le, gods pr, for our and y, give,. Run again, the program draws exactly the same batch. (The step file uses Python's random module for the positions; Karpathy's lecture code uses torch.randint. Either works; they produce different batches.)

The pipeline is now complete. Shakespeare's text has become batches of shape (batch, time) of character indices, each paired with a target batch of the same shape holding the next character at every position. The next phase, The Bigram Language Model, builds the first model that consumes them.

References

• Karpathy, A. (2015). char-rnn [Software repository]. https://github.com/karpathy/char-rnn
• Karpathy, A. (2023). Let's build GPT: from scratch, in code, spelled out [Video lecture]. https://www.youtube.com/watch?v=kCc8FmEb1nY. Code: https://github.com/karpathy/ng-video-lecture
• Code for this article: gpt_text_model/steps/1_getting_started in github.com/mdb42/ml-demos.