datasets
Hugging Face Datasets loads and preprocesses ML training data. One call, load_dataset(), pulls any of the public datasets on the Hugging Face Hub or reads your local CSV, JSON, Parquet, or image folders, and hands back a Dataset backed by Apache Arrow. Arrow memory-maps from disk, so a 100 GB dataset does not need 100 GB of RAM. You transform with map and filter, results are cached and reused, and one method converts to PyTorch, TensorFlow, JAX, NumPy, Pandas, or Polars for training.
The default way to move training data in the Hugging Face ecosystem, and deservedly so: Arrow-backed memory mapping and cached map() solve real problems. Budget disk for the cache, pin the major version, and reach for webdataset or mosaicml-streaming when raw distributed throughput is the whole job.
Use it if
- You train or evaluate models on public benchmarks: load_dataset('rajpurkar/squad') replaces a download script, a parser, and a caching layer
- Your dataset is larger than RAM: Arrow memory-mapping and streaming=True let you iterate corpora that never fully touch memory or even disk
- You preprocess with tokenizers: map(fn, batched=True, num_proc=N) with automatic caching is the standard companion to transformers
- You publish data: push_to_hub gives your dataset versioning, a browser preview, and one-line loading for everyone else
- Your work is tabular analytics, not model training: polars and duckdb query CSV and Parquet directly with far less machinery and no cache directory to manage
- Disk space is tight: every map() writes a new Arrow cache file, so a few tokenization passes over a large dataset quietly multiplies its footprint under ~/.cache/huggingface
- You need maximum-throughput sequential streaming for distributed training: webdataset and mosaicml-streaming are built around sharded sequential reads and resumption; IterableDataset covers less of that ground
- You depend on old script-based Hub datasets: loading-script support was removed in the 4.x line, so legacy trust_remote_code datasets no longer load and older tutorials break
Setup reality
pip install datasets brings pyarrow, pandas, fsspec, dill, multiprocess, and huggingface-hub, a heavy but pure-wheel install that works everywhere. Audio, image, and video decoding are extras (datasets[audio] and datasets[vision] pull torchcodec and Pillow), and forgetting them surfaces as errors only when you access a column. The cache in ~/.cache/huggingface grows without bounds until you learn cleanup_cache_files and HF_HOME. Gated datasets need huggingface-cli login first. Upgrading across majors is the real hazard: 4.x removed script-based loading, so pin before you upgrade a training repo.
Patterns
Load a dataset from the Hubload-hub-dataset
from datasets import load_dataset
ds = load_dataset("rajpurkar/squad")
print(ds) # DatasetDict with train/validation
print(ds["train"][0]) # first example as a dictUse the namespaced id (owner/name); bare names from old tutorials often point at deprecated mirrors or fail outright.
Load local CSV, JSON, or Parquetload-local-files
from datasets import load_dataset
ds = load_dataset("csv", data_files="my_data.csv")
ds = load_dataset("parquet", data_files={"train": "data/train/*.parquet",
"test": "data/test/*.parquet"})Everything lands in a 'train' split unless you pass a dict of data_files; glob patterns work.
Stream a dataset without downloading itstream-dataset
ds = load_dataset("HuggingFaceFW/fineweb", split="train", streaming=True)
for example in ds.take(5):
print(example["text"][:80])Streaming returns an IterableDataset: no random indexing, no len(), and map() runs lazily on iteration instead of eagerly with a cache.
Tokenize with batched mapmap-tokenize
from transformers import AutoTokenizer
tok = AutoTokenizer.from_pretrained("bert-base-cased")
def tokenize(batch):
return tok(batch["text"], truncation=True, max_length=512)
ds = ds.map(tokenize, batched=True, num_proc=4)batched=True passes dicts of lists and is dramatically faster; each map writes a new Arrow cache file, so repeated experiments eat disk.
Filter examples by a predicatefilter-rows
short = ds.filter(lambda x: len(x["text"]) < 1000)
short_batched = ds.filter(lambda batch: [len(t) < 1000 for t in batch["text"]], batched=True)The batched form must return a list of booleans, one per row; non-batched filter on millions of rows is slow enough to feel broken.
Make a train/test splitsplit-train-test
splits = ds["train"].train_test_split(test_size=0.1, seed=42)
train_ds, test_ds = splits["train"], splits["test"]Returns a DatasetDict with exactly 'train' and 'test' keys; pass seed or the split changes between runs and poisons comparisons.
Feed a Dataset to a PyTorch DataLoaderformat-for-pytorch
from torch.utils.data import DataLoader
ds = ds.with_format("torch", columns=["input_ids", "attention_mask", "label"])
loader = DataLoader(ds, batch_size=32, shuffle=True)with_format returns a new view without copying data; leaving string columns in scope breaks default collation, so list the columns you train on.
Build a Dataset from dicts, pandas, or a generatorfrom-python-objects
from datasets import Dataset
import pandas as pd
ds = Dataset.from_dict({"text": ["Hello", "World"]})
ds = Dataset.from_pandas(pd.DataFrame({"a": [1, 2]}))
def gen():
for i in range(1000):
yield {"value": i}
ds = Dataset.from_generator(gen)from_generator writes through Arrow as it goes, so it handles data far larger than RAM; from_pandas can smuggle in an unwanted __index_level_0__ column.
Save to disk and reload latersave-and-reload
ds.save_to_disk("./processed_squad")
from datasets import load_from_disk
ds = load_from_disk("./processed_squad")save_to_disk writes Arrow files, and the pair must match: a directory written by save_to_disk only reopens with load_from_disk, not load_dataset.
Publish a dataset to the Hubpush-to-hub
ds.push_to_hub("your-username/my-dataset", private=True)
reloaded = load_dataset("your-username/my-dataset")Requires huggingface-cli login or an HF_TOKEN env var; data uploads as Parquet, and the Hub renders a browsable preview automatically.
Work with audio or image columnsdecode-audio-images
from datasets import load_dataset, Audio
ds = load_dataset("PolyAI/minds14", "en-US", split="train")
ds = ds.cast_column("audio", Audio(sampling_rate=16_000))
sample = ds[0]["audio"] # decoded on accessDecoding happens lazily on access and needs the audio/vision extras installed; cast_column resamples audio on the fly rather than rewriting files.
Shuffle and take a subsetshuffle-and-select
small = ds.shuffle(seed=42).select(range(1000))shuffle builds an index mapping that makes later reads slower; call flatten_indices() afterward if you will iterate the subset many times.
Alternatives
| Package | Registry | Pick it when |
|---|---|---|
| webdataset | PyPI | You want tar-shard streaming tuned for high-throughput distributed training loops. |
| mosaicml-streaming | PyPI | You train from cloud object storage and need deterministic resumption mid-epoch. |
| polars | PyPI | Your job is fast tabular transformation and analysis, not feeding a model trainer. |
| lancedb | PyPI | You need an embedded columnar store with vector search over your samples, not just sequential loading. |