Implementation:Huggingface Alignment handbook Get Dataset
| 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
- Repository: alignment-handbook
- File: src/alignment/data.py (lines 26-79)
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),
)