HuggingFace Datasets

Overview

The datasets library by HuggingFace provides memory-efficient, Apache Arrow-backed access to tens of thousands of public datasets and your own local/remote data. It supports streaming for datasets that don't fit in RAM and integrates directly with transformers tokenizers and training loops.

When to Use

Use this skill when:

• Loading datasets from the HuggingFace Hub or local files (Parquet, JSONL, CSV, Arrow).
• Streaming datasets too large for memory.
• Applying preprocessing, tokenization, or filtering at scale with .map() and .filter().
• Pushing curated datasets to the Hub for sharing or versioning.
• Interleaving or concatenating multiple datasets.

Do not use for:

• Small in-memory pandas DataFrames with < 100K rows — use pandas directly.
• Model inference — use transformers or vllm skills.
• Embedding computation — use embedding-analysis skill.

Installation

uv pip install datasets pyarrow huggingface_hub

Quick Start

Loading

from datasets import load_dataset, load_from_disk

# From the Hub
dataset = load_dataset("imdb", split="train")

# From local files
dataset = load_dataset("parquet", data_files="data/*.parquet", split="train")

# From disk (pre-saved)
dataset = load_from_disk("./my_saved_dataset")

# Streaming (never loads full dataset into memory)
dataset = load_dataset("c4", "en", split="train", streaming=True)

Inspection

print(dataset)                    # Dataset info
print(dataset.shape)              # (n_rows, n_columns)
print(dataset.features)           # Schema
print(dataset[0])                 # First row
print(dataset[:5])                # First 5 rows as dict
print(dataset.column_names)       # Column list

Core Operations

Map (Apply Function)

# Basic mapping
def add_length(example):
    example["text_length"] = len(example["text"])
    return example

dataset = dataset.map(add_length)

# Batched mapping (faster for large datasets)
dataset = dataset.map(add_length, batched=True, batch_size=1000)

# Tokenization with transformers
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")

def tokenize(batch):
    return tokenizer(batch["text"], truncation=True, padding="max_length", max_length=512)

dataset = dataset.map(tokenize, batched=True)

Filter

# Quality filter: remove short or low-quality examples
dataset = dataset.filter(lambda x: len(x["text"]) > 100)
dataset = dataset.filter(lambda x: x["score"] > 0.5 if x["score"] is not None else False)

# Language filter
from langdetect import detect
dataset = dataset.filter(lambda x: detect(x["text"]) == "en")

Split

# Random split
train_test = dataset.train_test_split(test_size=0.2, seed=42)
train = train_test["train"]
test = train_test["test"]

# Stratified split (requires class labels)
from datasets import ClassLabel, DatasetDict

split_dataset = dataset.class_encode_column("label").train_test_split(
    test_size=0.2, stratify_by_column="label", seed=42
)

Select, Shuffle, Sort

# Shuffle
dataset = dataset.shuffle(seed=42)

# Select specific indices
dataset = dataset.select(range(1000))

# Sort
dataset = dataset.sort("timestamp")

Concatenate and Interleave

from datasets import concatenate_datasets, interleave_datasets

# Concatenate
combined = concatenate_datasets([dataset_a, dataset_b])

# Interleave (round-robin mixing)
mixed = interleave_datasets([dataset_a, dataset_b], probabilities=[0.7, 0.3], seed=42)

Streaming Pipeline Pattern

For datasets that don't fit in RAM, compose a streaming pipeline:

dataset = load_dataset("bigcode/the-stack", "python", split="train", streaming=True)

# Filter first (cheap operation)
dataset = dataset.filter(lambda x: len(x["content"]) > 500)

# Then map (expensive operation only on filtered data)
dataset = dataset.map(enrich_with_heuristics)

# Take only what you need
for i, example in enumerate(dataset):
    if i >= 10000:
        break
    process(example)

Save and Publish

Save Locally

# Save to disk (Arrow format, fastest reload)
dataset.save_to_disk("./my_dataset")

# Export to other formats
dataset.to_parquet("./my_dataset.parquet")
dataset.to_json("./my_dataset.jsonl")
dataset.to_csv("./my_dataset.csv")

Push to HuggingFace Hub

from huggingface_hub import login
login()

dataset.push_to_hub("my-username/my-curated-dataset")

# With a DatasetDict (train/val/test)
dataset_dict = DatasetDict({
    "train": train_dataset,
    "validation": val_dataset,
    "test": test_dataset,
})
dataset_dict.push_to_hub("my-username/my-curated-dataset")

Performance Tips

1. Use batched=True in .map() — 10-100x faster than row-level mapping.
2. Filter before map — avoid expensive operations on data you'll discard.
3. Use streaming + take for exploration before committing to a full load.
4. Prefer Parquet over JSONL/CSV for on-disk storage — columnar access is faster.
5. Set num_proc for CPU-bound map operations:

dataset = dataset.map(heavy_function, batched=True, num_proc=4)

Quality Gate

A dataset load pipeline is correct when:

• .features matches the expected schema.
• Streaming mode works for datasets > available RAM.
• Filter operations are deterministic (fixed seed where relevant).
• The dataset is saved to disk or pushed to Hub with all splits.