Tokenized Dataset with Sliding Window

A causal language model trains on batches of token ids, and the data pipeline's job is to produce those batches on demand. Each batch holds (B, T) input ids and (B, T) target ids, where the target is the input shifted left by one. The pipeline has to deliver that contract deterministically and reproducibly, from a corpus that may be several gigabytes of raw text.

The pipeline has four pieces. The tokenizer from the previous lesson turns text into one long flat list of ids. A sliding window slices that list into training examples. A custom Dataset exposes the examples as tensors. A DataLoader batches them and shuffles them with a known seed.

A causal LM consumes ids of shape (B, T), where B is the batch size and T is the context length. The target at position t is the input at position t+1, so every training example covers T+1 raw ids. The window stride sets how much consecutive examples overlap.

The slicer stays inside the corpus. If the last window lacks the ids to fill T+1 positions, the slicer drops it. Padding the tail with <|pad|> is also valid, but it complicates the loss mask. Here we drop.

A sliding window is a fixed-size slice that moves along the id stream by a set stride. A pretraining corpus is one long stream of ids, and with only non-overlapping windows every example would put its boundaries at the same T-aligned positions. Adjusting the stride moves those boundaries around, so the model sees more varied predict-next-token tasks.

A stride of T gives non-overlapping windows. A stride of T // 2 gives fifty-percent overlap and doubles the effective dataset. A stride of 1 gives maximum overlap and multiplies the dataset by a factor of T. The cost is more compute per epoch, and the benefit is more boundary diversity. Most pretraining runs use a stride equal to the context length. The corpus is already much larger than the model can finish in one epoch, so the boundary-diversity argument is weaker.

A PyTorch Dataset is a class with two required methods. __len__ returns the number of examples, and __getitem__ returns one example as a pair of tensors. Ours stores the encoded id stream and the stride. Indexing computes the start of the window on the fly, so the memory cost is one copy of the id stream no matter how many examples the stride produces.

The shift-by-one happens inside __getitem__. It returns (input, target) where input = window[:-1] and target = window[1:], both as PyTorch long tensors. The training loop treats them as ground truth.

A deterministic shuffle is one that produces the same order every time you restart the run. A DataLoader with shuffle=True draws from a PyTorch random generator, so passing an explicit torch.Generator seeded per epoch fixes the order. That matters when you want to compare two runs that differ in a single hyperparameter. Without a seed the two runs see the data in different orders, and the loss curves diverge for reasons unrelated to the change.

The seed contract is epoch_seed = base_seed + epoch_index. make_dataloader takes base_seed and epoch as arguments and seeds the generator with their sum. The caller passes the epoch index and increments it at the top of each epoch. A re-run with the same base seed sees the same order in every epoch.

The default PyTorch sampler picks indices uniformly at random with replacement disabled, which is what pretraining wants. Finetuning on a small dataset uses the same contract. The DataLoader assembles a batch by calling __getitem__ B times and stacking the results. Every example has the same length by construction, so no padding logic is needed.

We keep num_workers=0 for simplicity. In a production run, workers parallelize the __getitem__ calls. Here that is mostly a no-op, since the work is a slice of an in-memory tensor, but the same Dataset API supports workers cleanly.

For an id stream of length N, a context length T, and a stride S, the number of examples is max(0, 1 + (N - (T + 1)) // S). The Dataset exposes that calculation as a static method, so the trainer can compute total steps per epoch without iterating.

There is no disk streaming. The corpus is encoded fully in memory and held as a single tensor. For a few million ids that is well under a hundred megabytes, which suits this build. Disk streaming is a separate concern: it plugs in by replacing the storage and keeps the Dataset contract.

There is also no per-document handling. The corpus is treated as one continuous id stream. When the corpus is built from multiple documents, the next-document boundary is encoded by inserting <|endoftext|> ids, and the model learns to predict around the boundary.

main.py defines two classes (MiniBPE and SlidingWindowDataset) plus helpers. SlidingWindowDataset is the PyTorch Dataset. make_dataloader returns a configured DataLoader with a seeded generator. _encode_corpus_to_ids trains the inline BPE tokenizer on the corpus, then encodes it into one flat id list. The demo at the bottom builds a small tokenizer in-process, encodes a built-in corpus, constructs the dataset and dataloader, prints one batch, and asserts the shape contract. The tests in code/tests/test_dataset.py pin the window count formula, the shift-by-one property, the deterministic shuffle, and the stride trade-off.

Run the demo, then change the context length from 16 to 32 and watch the number of examples per epoch fall. That number is your steps-per-epoch budget.