Data loading, batching, and efficient sequence packing

PyTorch Dataset and DataLoader: your LM input pipeline

You’ll feed your language model with batches of token IDs; Dataset and DataLoader are the standard way to do that in PyTorch.

Implement a Dataset for pre-tokenized data

Assume you already have a long list (or array) of token IDs, e.g. from a tokenizer:

pythontokens = [101, 42, 13, 99, 7, 50256, 123, 987, ...] # pre-tokenized corpus

For next-token language modeling, each training example is usually:

  • input: a sequence of tokens
  • target: the same sequence shifted by one position (the “next token” labels)

Mechanism:

  • You choose a fixed seq_len.
  • You slice the long token stream into chunks of length seq_len + 1.
  • For each chunk:
    • input = chunk[:-1] (all tokens except the last)
    • target = chunk[1:] (all tokens except the first)

So if chunk = [10, 20, 30, 40], then:

  • input = [10, 20, 30]
  • target = [20, 30, 40]

Here’s a simple Dataset that does exactly that:

pythonimport torch from torch.utils.data import Dataset class LMDataset(Dataset): def __init__(self, token_ids, seq_len): # token_ids: 1D list / numpy array / tensor of ints self.tokens = torch.tensor(token_ids, dtype=torch.long) self.seq_len = seq_len # how many full (seq_len + 1) segments fit? self.num_segments = (len(self.tokens) - 1) // self.seq_len def __len__(self): # each segment gives one (input, target) pair return self.num_segments def __getitem__(self, idx): # continuous span in token space start = idx * self.seq_len end = start + self.seq_len + 1 # need one extra for the target shift chunk = self.tokens[start:end] # shape: [seq_len + 1] x = chunk[:-1] # [seq_len] y = chunk[1:] # [seq_len] return x, y

Example you can run:

pythontokens = list(range(20)) # [0, 1, 2, ..., 19] ds = LMDataset(tokens, seq_len=5) print("len(ds) =", len(ds)) for i in range(3): x, y = ds[i] print(f"i={i}, x={x.tolist()}, y={y.tolist()}")

You should see:

  • i=0: x=[0,1,2,3,4], y=[1,2,3,4,5]
  • i=1: x=[5,6,7,8,9], y=[6,7,8,9,10]
  • etc.

Mechanism summary:

  • __len__ tells PyTorch how many indexable examples exist.
  • __getitem__ defines how to turn an index into your (input, target) pair.

If your underlying tokens length is smaller than seq_len + 1, the integer division will make num_segments zero and your dataset will be empty; always check len(ds) > 0 for your chosen seq_len.

1 / 5