Implementation:Eric mitchell Direct preference optimization Get Dataset

Knowledge Sources	Direct Preference Optimization
Domains	Data_Engineering, Extensibility
Last Updated	2026-02-08 02:00 GMT

Overview

Concrete tool for dispatching dataset loading by name with format validation provided by the direct-preference-optimization repository.

Description

The get_dataset function is a simple dispatcher that maps dataset name strings to their loader functions. It supports three built-in datasets (shp, hh, se) and validates the output format with an assertion.

Usage

Called by get_batch_iterator during data pipeline construction. Users extend this by adding elif branches for custom datasets.

Code Reference

Source Location

• Repository: direct-preference-optimization
• File: preference_datasets.py
• Lines: 163-177

Signature

def get_dataset(
    name: str,
    split: str,
    silent: bool = False,
    cache_dir: str = None,
):
    """Load the given dataset by name. Supported by default are 'shp', 'hh', and 'se'."""

Import

from preference_datasets import get_dataset

I/O Contract

Inputs

Name	Type	Required	Description
name	str	Yes	Dataset identifier: "shp", "hh", "se", or custom name
split	str	Yes	"train" or "test"
silent	bool	No	Suppress progress bars (default False)
cache_dir	str	No	HuggingFace dataset cache directory

Outputs

Name	Type	Description
data	Dict[str, Dict]	Canonical format dict; assertion validates keys are exactly {responses, pairs, sft_target}

Usage Examples

Loading Built-in Dataset

from preference_datasets import get_dataset

data = get_dataset('hh', split='train', cache_dir='.cache')
print(f"Loaded {len(data)} prompts")

Adding Custom Dataset

# In preference_datasets.py, add to get_dataset:
def get_dataset(name, split, silent=False, cache_dir=None):
    if name == 'shp':
        data = get_shp(split, silent=silent, cache_dir=cache_dir)
    elif name == 'hh':
        data = get_hh(split, silent=silent, cache_dir=cache_dir)
    elif name == 'se':
        data = get_se(split, silent=silent, cache_dir=cache_dir)
    elif name == 'custom':  # NEW
        data = get_custom(split, silent=silent, cache_dir=cache_dir)
    else:
        raise ValueError(f"Unknown dataset '{name}'")

    # This assertion validates the format
    assert set(list(data.values())[0].keys()) == {'responses', 'pairs', 'sft_target'}
    return data

Related Pages

Implements Principle

• Principle:Eric_mitchell_Direct_preference_optimization_Dataset_Registration