Jump to content

Connect SuperML | Leeroopedia MCP: Equip your AI agents with best practices, code verification, and debugging knowledge. Powered by Leeroo — building Organizational Superintelligence. Contact us at founders@leeroo.com.

Implementation:Huggingface Alignment handbook Get Dataset

From Leeroopedia


Knowledge Sources
Domains NLP, Data_Engineering
Last Updated 2026-02-07 00:00 GMT

Overview

Concrete tool for loading single datasets from HuggingFace Hub for alignment training, provided by the alignment-handbook library.

Description

The get_dataset function is the alignment-handbook's unified dataset loading interface. In single dataset mode, it delegates to datasets.load_dataset with the dataset name and config from ScriptArguments. The function validates that either dataset_name or dataset_mixture is provided, ensuring clear data source specification.

Usage

Import this function when loading standard benchmark datasets for SFT or preference training without mixture blending. This is the default mode for most alignment-handbook recipes.

Code Reference

Source Location

Signature

def get_dataset(args: ScriptArguments) -> DatasetDict:
    """Load a dataset or a mixture of datasets based on the configuration.

    Args:
        args (ScriptArguments): Script arguments containing dataset configuration.

    Returns:
        DatasetDict: The loaded datasets.
    """

Import

from alignment import get_dataset, ScriptArguments

I/O Contract

Inputs

Name Type Required Description
args ScriptArguments Yes Script arguments with dataset_name or dataset_mixture set
args.dataset_name Optional[str] Conditional HuggingFace dataset ID (e.g., "HuggingFaceH4/ultrachat_200k"). Required if dataset_mixture is not set
args.dataset_config Optional[str] No Dataset configuration name (e.g., subset or version)
args.dataset_mixture Optional[DatasetMixtureConfig] Conditional Mixture config. Required if dataset_name is not set

Outputs

Name Type Description
return DatasetDict Dictionary with train split (and optionally test split). Contains columns appropriate for the training objective (messages for SFT, chosen/rejected for DPO/ORPO)

Usage Examples

Single Dataset Loading (SFT)

from alignment import get_dataset, ScriptArguments

# ScriptArguments populated from YAML config
# dataset_name: HuggingFaceH4/ultrachat_200k
dataset = get_dataset(script_args)

# Access splits
train_data = dataset["train"]
print(f"Training samples: {len(train_data)}")
# Training samples: 200000

Single Dataset Loading (DPO/ORPO)

from alignment import get_dataset, ScriptArguments

# For DPO: dataset_name = "HuggingFaceH4/ultrafeedback_binarized"
# For ORPO: dataset_name = "argilla/distilabel-capybara-dpo-7k-binarized"
dataset = get_dataset(script_args)

# DPO/ORPO datasets have chosen/rejected columns
train_data = dataset["train"]
print(train_data.column_names)
# ['prompt', 'chosen', 'rejected']

Usage in Training Script

# From scripts/sft.py
dataset = get_dataset(script_args)

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=dataset[script_args.dataset_train_split],
    eval_dataset=(
        dataset[script_args.dataset_test_split]
        if training_args.eval_strategy != "no"
        else None
    ),
    processing_class=tokenizer,
    peft_config=get_peft_config(model_args),
)

Related Pages

Implements Principle

Requires Environment

Page Connections

Double-click a node to navigate. Hold to expand connections.
Principle
Implementation
Heuristic
Environment