Implementation:Huggingface Datasets Dataset Train Test Split

Knowledge Sources	Huggingface Datasets HF Datasets Docs
Domains	Data_Engineering, ML_Preprocessing
Last Updated	2026-02-14 18:00 GMT

Overview

Concrete tool for splitting a dataset into train and test subsets provided by the HuggingFace Datasets library.

Description

The train_test_split method divides a dataset into a DatasetDict containing "train" and "test" splits. It supports proportional splitting (float values between 0 and 1) and absolute splitting (integer values specifying exact counts). The method supports optional shuffling with a seed for reproducibility, and stratified splitting that preserves class distributions across both splits. If neither test_size nor train_size is provided, the default test size is 0.25. This method is similar to scikit-learn's train_test_split.

Usage

Use Dataset.train_test_split when your dataset does not have predefined splits and you need to create train/test partitions for model evaluation, especially when you need reproducible or stratified splits.

Code Reference

Source Location

• Repository: datasets
• File: src/datasets/arrow_dataset.py
• Lines: L4635-L4916

Signature

@transmit_format
@fingerprint_transform(
    inplace=False,
    randomized_function=True,
    fingerprint_names=["train_new_fingerprint", "test_new_fingerprint"],
    ignore_kwargs=["load_from_cache_file", "train_indices_cache_file_name", "test_indices_cache_file_name"],
)
def train_test_split(
    self,
    test_size: Union[float, int, None] = None,
    train_size: Union[float, int, None] = None,
    shuffle: bool = True,
    stratify_by_column: Optional[str] = None,
    seed: Optional[int] = None,
    generator: Optional[np.random.Generator] = None,
    keep_in_memory: bool = False,
    load_from_cache_file: Optional[bool] = None,
    train_indices_cache_file_name: Optional[str] = None,
    test_indices_cache_file_name: Optional[str] = None,
    writer_batch_size: Optional[int] = 1000,
    train_new_fingerprint: Optional[str] = None,
    test_new_fingerprint: Optional[str] = None,
) -> "DatasetDict":

Import

from datasets import load_dataset

ds = load_dataset("cornell-movie-review-data/rotten_tomatoes", split="validation")
splits = ds.train_test_split(test_size=0.2, seed=42)

I/O Contract

Inputs

Name	Type	Required	Description
test_size	Union[float, int, None]	No	Proportion (float in (0,1)) or absolute number (int) of test samples. Defaults to 0.25 if both sizes are None.
train_size	Union[float, int, None]	No	Proportion (float in (0,1)) or absolute number (int) of train samples. Defaults to complement of test_size.
shuffle	bool	No	Whether to shuffle data before splitting. Defaults to True.
stratify_by_column	Optional[str]	No	Column name of labels for stratified splitting. Must be a ClassLabel column.
seed	Optional[int]	No	Random seed for reproducibility.
generator	Optional[np.random.Generator]	No	NumPy random Generator for computing the permutation.
keep_in_memory	bool	No	Keep split indices in memory. Defaults to False.
load_from_cache_file	Optional[bool]	No	Use cached indices if available.
train_indices_cache_file_name	Optional[str]	No	Cache file path for train split indices.
test_indices_cache_file_name	Optional[str]	No	Cache file path for test split indices.
writer_batch_size	Optional[int]	No	Rows per write operation. Defaults to 1000.
train_new_fingerprint	Optional[str]	No	Fingerprint for the train set after transform.
test_new_fingerprint	Optional[str]	No	Fingerprint for the test set after transform.

Outputs

Name	Type	Description
return	DatasetDict	A dictionary with "train" and "test" keys, each containing a Dataset.

Usage Examples

Basic Usage

from datasets import load_dataset

ds = load_dataset("cornell-movie-review-data/rotten_tomatoes", split="validation")

# Basic 80/20 split with seed
splits = ds.train_test_split(test_size=0.2, seed=42)
print(splits)
# DatasetDict({
#     train: Dataset({ features: ['text', 'label'], num_rows: 852 })
#     test: Dataset({ features: ['text', 'label'], num_rows: 214 })
# })

# Stratified split
ds_imdb = load_dataset("stanfordnlp/imdb", split="train")
splits = ds_imdb.train_test_split(test_size=0.2, stratify_by_column="label")

Related Pages

Implements Principle

• Principle:Huggingface_Datasets_Train_Test_Splitting