📦 Dataset Pipeline

Data Loading and Preprocessing for All Training Stages

Files

Dataset Classes

~218

Lines

🐍

Python

🔵 Pretrain 🟢 SFT · Chat Template 🟡 DPO · Preference 🔴 RLAIF · RL

🍳 Kitchen Analogy

Think of the dataset module as ingredient preparation stations in a professional kitchen. Each station (Dataset class) takes raw ingredients (JSONL data) and prepares them differently depending on the dish being made — the Pretrain station washes and chops text into uniform pieces, the SFT station marinates conversations into chat templates, the DPO station carefully selects the best and worst preparations for comparison, and the RLAIF station stages prompts for the chef (RL agent) to improvise with. Every station outputs precisely formatted batches, ready for the training "oven."

Architecture Overview

classDiagram
    class Dataset {
        <<PyTorch>>
        +__len__()
        +__getitem__(index)
    }
    class PretrainDataset {
        -data: List
        -max_length: int
        +__init__(data_path, tokenizer, max_length)
        +__getitem__(index) → (X, Y, loss_mask)
    }
    class SFTDataset {
        -data: List
        +create_chat_prompt(conversations)
        +generate_labels(input_ids)
        +__getitem__(index) → (X, Y, loss_mask)
    }
    class DPODataset {
        -data: List
        +__getitem__(index) → (x_cho, y_cho, m_cho, x_rej, y_rej, m_rej)
    }
    class RLAIFDataset {
        -data: List
        +create_chat_prompt(conversations)
        +__getitem__(index) → (prompt, answer)
    }
    Dataset <|-- PretrainDataset
    Dataset <|-- SFTDataset
    Dataset <|-- DPODataset
    Dataset <|-- RLAIFDataset

Key Classes & Functions

Name	Type	File	Description
`PretrainDataset`	Class	lm_dataset.py	Loads JSONL, tokenizes with BOS/EOS, pads to max_length; labels=-100 for padding tokens
`SFTDataset`	Class	lm_dataset.py	Applies chat template via `apply_chat_template()`, computes loss only on assistant turns
`DPODataset`	Class	lm_dataset.py	Pairs chosen/rejected conversations with corresponding loss masks
`RLAIFDataset`	Class	lm_dataset.py	Builds prompts with `add_generation_prompt=True` for PPO/GRPO/SPO training
`pre_processing_chat`	Function	lm_dataset.py	Randomly injects system prompts (10 bilingual options, 20% chance) via `add_system_ratio`
`post_processing_chat`	Function	lm_dataset.py	Strips empty `<think></think>` tags with 95% probability
`create_chat_prompt`	Method	lm_dataset.py	Handles tool calling by injecting function schemas into the system message
`generate_labels`	Method	lm_dataset.py	Masks non-assistant tokens to -100; loss only between `bos_id` and `eos_id` markers

Data Flow

sequenceDiagram
    participant JSONL as JSONL Files
    participant HF as HuggingFace Loader
    participant DS as Dataset Class
    participant TK as Tokenizer
    participant DL as DataLoader
    participant TR as Trainer

    JSONL->>HF: load_dataset("json")
    HF->>DS: raw samples
    DS->>TK: tokenize text/chat
    TK-->>DS: input_ids, attention_mask
    DS->>DS: generate_labels / loss_mask
    DS->>DL: (X, Y, loss_mask) batches
    DL->>TR: collated mini-batches
    TR->>TR: compute loss (masked)

Interactions & Dependencies

Component	Depends On	Interaction
`PretrainDataset`	HuggingFace datasets, tokenizer	Loads JSONL, tokenizes raw text to fixed-length sequences
`SFTDataset`	tokenizer.apply_chat_template	Formats multi-turn conversations using the model's chat template
`SFTDataset`	pre_processing_chat	Randomly adds system prompt before tokenization
`SFTDataset`	post_processing_chat	Strips empty think tags to control reasoning format
`DPODataset`	SFTDataset logic	Reuses chat template + labeling for chosen/rejected pairs
`RLAIFDataset`	tokenizer, reward model	Produces prompts; answers used downstream for reward computation
All Datasets	PyTorch DataLoader	Consumed by DataLoader with custom collate or default batching
`generate_labels`	bos_id / eos_id tokens	Scans token IDs to create loss masks targeting assistant segments

Deep-Dive Topics

💬 Chat Template Processing

How apply_chat_template() converts multi-turn conversations into token sequences, handling system messages, tool calls, and role markers for proper formatting.

🎯 Loss Mask Generation

The logic behind generate_labels(): scanning for bos/eos markers to compute loss only on assistant responses, avoiding training on user prompts or padding.

⚖️ DPO Data Format

How chosen/rejected conversation pairs are structured, tokenized in parallel, and returned with separate loss masks for Direct Preference Optimization training.

Directory Map

dataset/
├── __init__.py           # Package init
└── lm_dataset.py        # ~218 lines — All dataset classes & helpers
    ├── pre_processing_chat() # System prompt injection
    ├── post_processing_chat() # Think-tag cleanup
    ├── PretrainDataset       # Raw text tokenization
    ├── SFTDataset            # Chat template + loss masking
    ├── DPODataset            # Preference pair formatting
    └── RLAIFDataset          # RL prompt generation

Configuration Reference

Parameter	Location	Default	Description
`max_length`	PretrainDataset / SFTDataset	512	Maximum sequence length; shorter sequences are padded, longer truncated
`add_system_ratio`	pre_processing_chat	0.2	Probability of injecting a random system prompt (20%)
`empty_think_ratio`	post_processing_chat	0.05	Probability of keeping empty think tags (5% keep, 95% strip)
`data_path`	All Dataset classes	—	Path to JSONL training data file
`tokenizer`	All Dataset classes	—	HuggingFace tokenizer instance with chat template
`bos_id`	SFTDataset.generate_labels	tokenizer.bos_id	Begin-of-sequence marker for assistant response boundaries
`eos_id`	SFTDataset.generate_labels	tokenizer.eos_id	End-of-sequence marker for assistant response boundaries
`add_generation_prompt`	RLAIFDataset	True	Appends generation prompt token for RL inference

Code Examples

Loading an SFT Dataset

Python — SFT dataset initialization

from dataset.lm_dataset import SFTDataset
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("MiniMind2")
dataset = SFTDataset(
    data_path="data/sft_train.jsonl",
    tokenizer=tokenizer,
    max_length=512
)

# Each item: (input_ids, labels, loss_mask)
X, Y, loss_mask = dataset[0]
print(f"Sequence length: {len(X)}")
print(f"Active loss tokens: {loss_mask.sum().item()}")

DPO Preference Pair Access

Python — DPO dataset usage

from dataset.lm_dataset import DPODataset
from torch.utils.data import DataLoader

dpo_dataset = DPODataset(
    data_path="data/dpo_train.jsonl",
    tokenizer=tokenizer,
    max_length=512
)

loader = DataLoader(dpo_dataset, batch_size=4, shuffle=True)

for batch in loader:
    x_chosen, y_chosen, mask_chosen, \
    x_rejected, y_rejected, mask_rejected = batch
    # Feed to DPO trainer for preference optimization
    break