Splitting, Combining, Concatenating

~20 min · datasets, split

Level 0Scout

0 XP0/50 lessons0/10 achievements

0/120 XP to next level120 XP to go0% complete

train_test_split, stratify, seed

ds.train_test_split(test_size=0.1, seed=42, stratify_by_column="label") returns a DatasetDict with train and test splits. Stratification preserves class balance — required for any classification eval that isn't a uniformly distributed dataset.

concatenate_datasets and interleave_datasets

concatenate_datasets([a, b, c]) stacks rows. interleave_datasets([a, b, c], probabilities=[0.6, 0.3, 0.1]) samples from sources at each step. The latter is the right tool when you're mixing curated + crawled data with a target ratio.

Code

Stratified split + concatenate·python

from datasets import load_dataset, concatenate_datasets, interleave_datasets

a = load_dataset("stanfordnlp/imdb", split="train")
print(a.features["label"])  # ClassLabel

splits = a.train_test_split(test_size=0.1, seed=42, stratify_by_column="label")
print({k: len(v) for k, v in splits.items()})

# Concat (stacked rows)
combined = concatenate_datasets([splits["train"], splits["test"]])
print(len(combined), len(a))

# Interleave (probabilistic mix during training)
b = load_dataset("stanfordnlp/imdb", split="test")
mixed = interleave_datasets([a, b], probabilities=[0.7, 0.3], seed=42, stopping_strategy="first_exhausted")
print(type(mixed), len(mixed))

External links

Exercise

Take any classification dataset. Run a stratified 80/10/10 train/dev/test split with seed=42. Verify class balance is preserved across all three. Then run a non-stratified split with the same seed and compare class balance — usually noticeably worse on small datasets.

Progress

Progress is local-only — sign in to sync across devices.

← PreviousCustom Datasets from Local Files Next →Caching, push_to_hub, and Reproducibility

Spotted a bug or have feedback on this page?Report an Issue

Comments 0

🔔 Reply notifications (sign in)

No comments yet — be the first.