Implementation:Hiyouga LLaMA Factory Data Utils
| Knowledge Sources | |
|---|---|
| Domains | Data Processing, Cloud Storage |
| Last Updated | 2026-02-06 19:00 GMT |
Overview
Data Utils provides shared data types, dataset merging/splitting utilities, and cloud storage access functions for the LLaMA Factory data pipeline.
Description
The module defines foundational types: the Role enum (user, assistant, system, function, observation), the DatasetModule TypedDict for train/eval dataset pairs, and the SLOTS type alias for template slot sequences. It provides merge_dataset for combining multiple datasets via concatenation or interleaving (with configurable strategies: interleave_under, interleave_over, interleave_once), split_dataset for train/validation splitting supporting both map and iterable datasets, and get_dataset_module for converting HuggingFace DatasetDict objects into the internal DatasetModule format. Cloud storage utilities include setup_fs (creates fsspec filesystem objects for S3 or GCS) and read_cloud_json (reads JSON/JSONL files from cloud storage with automatic anonymous/credential-based access fallback).
Usage
Use these utilities when working with the data loading and processing pipeline. The merge and split functions are called by the dataset loader, while the cloud storage functions enable loading training data from S3 or GCS buckets.
Code Reference
Source Location
- Repository: Hiyouga_LLaMA_Factory
- File: src/llamafactory/data/data_utils.py
- Lines: 1-203
Signature
SLOTS = list[Union[str, set[str], dict[str, str]]]
class Role(StrEnum):
USER = "user"
ASSISTANT = "assistant"
SYSTEM = "system"
FUNCTION = "function"
OBSERVATION = "observation"
class DatasetModule(TypedDict):
train_dataset: Optional[Union["Dataset", "IterableDataset"]]
eval_dataset: Optional[Union["Dataset", "IterableDataset", dict[str, "Dataset"]]]
def merge_dataset(
all_datasets: list[Union["Dataset", "IterableDataset"]],
data_args: "DataArguments",
seed: int,
) -> Union["Dataset", "IterableDataset"]: ...
def split_dataset(
dataset: Optional[Union["Dataset", "IterableDataset"]],
eval_dataset: Optional[Union["Dataset", "IterableDataset", dict[str, "Dataset"]]],
data_args: "DataArguments",
seed: int,
) -> tuple[dict, dict]: ...
def get_dataset_module(dataset: Union["Dataset", "DatasetDict"]) -> "DatasetModule": ...
def setup_fs(path: str, anon: bool = False) -> "fsspec.AbstractFileSystem": ...
def read_cloud_json(cloud_path: str) -> list[Any]: ...
Import
from llamafactory.data.data_utils import (
Role,
SLOTS,
DatasetModule,
merge_dataset,
split_dataset,
get_dataset_module,
read_cloud_json,
)
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| all_datasets | list[Dataset or IterableDataset] | Yes (for merge) | List of datasets to merge |
| data_args | DataArguments | Yes | Data configuration with mix_strategy, val_size, streaming, buffer_size, interleave_probs |
| seed | int | Yes | Random seed for shuffling and splitting |
| dataset | Dataset or IterableDataset | No | Main training dataset to split |
| eval_dataset | Dataset, IterableDataset, or dict | No | Explicit evaluation dataset(s) |
| cloud_path | str | Yes (for cloud) | S3 (s3://) or GCS (gs://, gcs://) path to JSON/JSONL files |
Outputs
| Name | Type | Description |
|---|---|---|
| merged dataset | Dataset or IterableDataset | Single merged dataset from merge_dataset |
| (train_dict, eval_dict) | tuple[dict, dict] | Train and evaluation splits from split_dataset |
| DatasetModule | DatasetModule | Converted dataset module from get_dataset_module |
| list[Any] | list | Parsed JSON records from read_cloud_json |
Usage Examples
from llamafactory.data.data_utils import merge_dataset, split_dataset, read_cloud_json
# Merge multiple datasets
merged = merge_dataset([dataset_a, dataset_b], data_args, seed=42)
# Split into train/eval
train_dict, eval_dict = split_dataset(merged, None, data_args, seed=42)
# Read training data from S3
records = read_cloud_json("s3://my-bucket/training-data/alpaca.jsonl")
Related Pages
- Hiyouga_LLaMA_Factory_Data_Formatter - Template formatters that use the SLOTS type
- Hiyouga_LLaMA_Factory_API_Chat - API layer that uses Role enum for role mapping