Build a Custom Dataset Class

You've spent hours cleaning your data, formatting it into neat JSON lines, and verifying every label. But the moment you try to feed it into your finetuning script, you hit a wall: the model's tokenizer expects a specific format, your samples don't fit into the batch, and your GPU idles while you debug shape mismatches. The pain is real — most finetuning failures trace back to a poorly designed dataset class, not the model architecture. In this lesson, you'll learn to build a custom dataset class for finetuning that turns your raw data into a smooth, efficient training pipeline, so you can stop wrestling with dataloaders and start training.

The problem this lesson solves

Off-the-shelf datasets like datasets.load_dataset are great for standard benchmarks, but your finetuning task is unique. Your data might be:

• Stored in a proprietary format (Parquet, custom CSV, database dumps)
• Requiring preprocessing like truncation, padding, or special tokens
• Too large to load into memory all at once
• Asymmetric — inputs and targets have different lengths or roles

A generic torch.utils.data.Dataset doesn't know about tokenizers, attention masks, or sequence lengths. Without a custom dataset class, you end up writing messy preprocessing logic inside your training loop, duplicating tokenization calls, and battling with collation errors. The solution is to encapsulate all data handling in one reusable class, so your training script stays clean and your data pipeline is deterministic and efficient.

Core concept / mental model

Think of a custom dataset class as a data factory that produces one fully-prepared training sample at a time. The factory takes a raw recipe (your data file) and a tokenizer, and its job is to output a dict with keys like input_ids, attention_mask, and labels — exactly what your model expects.

Here's the mental model:

• __init__ sets up the raw materials: loads the data, stores the tokenizer, caches any expensive preprocessing.
• __len__ tells PyTorch how many samples exist, so the DataLoader can shuffle and batch them.
• __getitem__ takes an index, retrieves the raw sample, and transforms it into a model-ready format. This is where the magic happens.

This design keeps your code modular. You can swap the tokenizer, change the data source, or adjust the formatting logic without touching the training loop.

How it works step by step

Step 1: Define the class skeleton

Start with the __init__ method that receives the data path and tokenizer. Load your data into a list or a light structure (like a memory-mapped array). If your data is huge, avoid loading everything into RAM — use lazy loading with pyarrow or datasets.

Step 2: Implement __len__

Return the number of samples. This is trivial but essential for batching and epoch lengths.

Step 3: Implement __getitem__

The core logic: for a given index, fetch the raw sample, format it into a prompt-response pair, tokenize with truncation and padding, and return a dict with input_ids, attention_mask, and labels.

Step 4: Choose encoding strategy

Decide whether to tokenize on-the-fly (slow but memory-efficient) or pre-tokenize in __init__ (fast but memory-hungry). The right choice depends on your dataset size and RAM budget.

Step 5: Integrate with DataLoader

Pass your dataset to DataLoader and specify a custom collate_fn for dynamic padding per batch — critical for variable-length sequences.

Hands-on walkthrough

Let's build a FinetuneDataset class for instruction-following data. We'll use the Hugging Face transformers tokenizer and assume your data is a JSON file with input and output fields.

import json
import torch
from torch.utils.data import Dataset
from transformers import AutoTokenizer

class FinetuneDataset(Dataset):
    def __init__(self, data_path, tokenizer, max_length=512):
        with open(data_path) as f:
            self.data = json.load(f)  # list of dicts
        self.tokenizer = tokenizer
        self.max_length = max_length

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

    def __getitem__(self, idx):
        sample = self.data[idx]
        # Format as prompt + response
        prompt = f"Instruction: {sample['instruction']}\nInput: {sample['input']}\nResponse:"
        response = sample["output"]
        full_text = prompt + response

        # Tokenize with truncation; we'll pad later in collate_fn
        encodings = self.tokenizer(
            full_text,
            truncation=True,
            max_length=self.max_length,
            return_tensors="pt"
        )
        input_ids = encodings["input_ids"].squeeze(0)
        attention_mask = encodings["attention_mask"].squeeze(0)

        # For language modeling, labels = input_ids
        return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": input_ids.clone()}

Expected output

When you iterate, you get tensors:

tokenizer = AutoTokenizer.from_pretrained("gpt2")
ds = FinetuneDataset("data.json", tokenizer)
print(len(ds))  # e.g., 1000
sample = ds[0]
print(type(sample))  # dict
print(sample["input_ids"].shape)  # torch.Size([512])
print(sample["labels"])

Output:

1000
<class 'dict'>
torch.Size([512])
tensor([ 50256, 1234, ...])

Adding dynamic padding with collate_fn

Because each sample has a different length, we pad to the longest in the batch.

from torch.nn.utils.rnn import pad_sequence

def collate_fn(batch):
    input_ids = [item["input_ids"] for item in batch]
    attention_masks = [item["attention_mask"] for item in batch]
    labels = [item["labels"] for item in batch]

    # Pad sequences
    input_ids_padded = pad_sequence(input_ids, batch_first=True, padding_value=tokenizer.pad_token_id)
    attention_masks_padded = pad_sequence(attention_masks, batch_first=True, padding_value=0)
    labels_padded = pad_sequence(labels, batch_first=True, padding_value=-100)  # ignore index for loss

    return {
        "input_ids": input_ids_padded,
        "attention_mask": attention_masks_padded,
        "labels": labels_padded
    }

from torch.utils.data import DataLoader
loader = DataLoader(ds, batch_size=8, shuffle=True, collate_fn=collate_fn)
for batch in loader:
    print(batch["input_ids"].shape)  # (8, max_len)
    break

Output:

torch.Size([8, 512])

Pre-tokenized for speed

To speed up training, pre-tokenize all samples in __init__ (if memory allows).

class PreTokenizedDataset(Dataset):
    def __init__(self, data_path, tokenizer, max_length=512):
        with open(data_path) as f:
            raw = json.load(f)
        self.tokenizer = tokenizer
        self.max_length = max_length
        # Pre-tokenize everything
        self.samples = []
        for sample in raw:
            text = sample["input"] + sample["output"]  # simple concat
            tokens = self.tokenizer(text, truncation=True, max_length=self.max_length)
            self.samples.append(tokens)

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

    def __getitem__(self, idx):
        tokens = self.samples[idx]
        return {
            "input_ids": torch.tensor(tokens["input_ids"]),
            "attention_mask": torch.tensor(tokens["attention_mask"]),
            "labels": torch.tensor(tokens["input_ids"])
        }

Pro tip: Pre-tokenization reduces per-epoch CPU overhead but triples memory usage. Use it only when your data fits easily in RAM.

Compare options / when to choose what

When to choose what:

• On-the-fly for datasets > 5 GB or when you have GPU memory constraints but plenty of CPU time.
• Pre-tokenized for small, curated datasets where speed is critical (e.g., QLoRA experiments).
• datasets with streaming for massive corpora that don't fit on disk or RAM.

Troubleshooting & edge cases

• Shape mismatch errors: If your model expects labels of the same length as input_ids, always pad labels too, and set padding value to -100 (PyTorch's ignore index).
• Missing pad token: Some tokenizers (like GPT-2) don't have a pad token by default. Set tokenizer.pad_token = tokenizer.eos_token to avoid RuntimeError: input_ids must be padded with a pad_token_id.
• Padding on the left vs. right: For causal LLMs, pad on the left if you're generating later, but for pure finetuning, right‑padding is fine. Use padding_side in the tokenizer.
• Truncation causing lost responses: If you truncate the full prompt+response, you might cut off the response entirely. Consider truncating the prompt first to reserve space for the response.
• Memory explosion in collate_fn: If you have extreme length variance, dynamic padding can still create giant batches. Use DataLoader with pin_memory=True and consider gradient accumulation to manage batch sizes.
• Slow epoch with large data: If tokenization dominates, enable num_workers in DataLoader, but ensure your dataset class is picklable (avoid lambdas, use top-level functions).

What you learned & what's next

You now know how to build a custom dataset class for finetuning, including: - Implementing __init__, __len__, and __getitem__. - Choosing between on-the-fly and pre-tokenized strategies. - Creating a custom collate_fn for dynamic padding. - Avoiding common pitfalls like pad token issues and label truncation.

Key takeaways:

• Your dataset class should produce model-ready tensors, not raw text.
• Dynamic padding with -100 labels prevents loss on padding tokens.
• Pre-tokenization trades memory for speed.
• Always handle tokenizer pad tokens explicitly.
• Design for modularity: swap tokenizer or data source without changing training code.

Next step: In the next lesson, you'll learn to integrate this dataset into a full finetuning loop with Hugging Face Trainer, combining your data factory with training arguments and evaluation metrics. You'll move from building blocks to a working training pipeline.