HuggingFace Datasets: load_dataset, map, filter and Streaming Reference Guide

# Core
pip install datasets                # ≥ 3.0

# Optional extras — install what you need
pip install "datasets[audio]"       # soundfile + librosa + torchcodec
pip install "datasets[vision]"      # Pillow + image features
pip install "datasets[s3]"          # s3fs for cloud-backed paths

# To upload to the Hub
pip install huggingface_hub
hf auth login                       # paste a token from huggingface.co/settings/tokens

# Optional speed-ups
pip install pyarrow                 # bundled, but pin if you need a version

from datasets import load_dataset
from transformers import AutoTokenizer

ds = load_dataset("imdb")
tok = AutoTokenizer.from_pretrained("bert-base-uncased")

# Batched map = orders of magnitude faster than per-example
def tokenize(batch: dict) -> dict:
    return tok(batch["text"], truncation=True, max_length=128)

encoded = ds.map(
    tokenize,
    batched=True,
    batch_size=1000,
    num_proc=4,                   # multi-process
    remove_columns=["text"],      # drop raw text once tokenized
    desc="Tokenizing",
)

# Cast labels to PyTorch tensors lazily
encoded = encoded.with_format("torch", columns=["input_ids", "attention_mask", "label"])
print(encoded["train"][0])

from datasets import load_dataset

# Streaming = no full download, no disk caching
ds = load_dataset(
    "HuggingFaceFW/fineweb",
    split="train",
    streaming=True,
)

ds = ds.shuffle(seed=42, buffer_size=10_000)   # approximate shuffle
ds = ds.filter(lambda x: len(x["text"]) > 200)
ds = ds.map(lambda x: {"chars": len(x["text"])})

# IterableDataset — iterate in order; no random access
for i, row in enumerate(ds.take(5)):
    print(i, row["chars"], row["text"][:80])

# Wrap for PyTorch DataLoader
from torch.utils.data import DataLoader
loader = DataLoader(ds.with_format("torch"), batch_size=32)

from datasets import load_dataset, DatasetDict, Features, ClassLabel, Value

# 1 · Load a CSV from disk and a JSONL from a URL — same call shape
local = load_dataset("csv", data_files="reviews.csv", split="train")
remote = load_dataset(
    "json",
    data_files="https://example.com/reviews.jsonl",
    split="train",
)

# 2 · Concatenate and define an explicit schema
from datasets import concatenate_datasets
ds = concatenate_datasets([local, remote])

ds = ds.cast(Features({
    "text":  Value("string"),
    "label": ClassLabel(names=["neg", "pos"]),
}))

# 3 · Clean, dedupe, split
ds = ds.filter(lambda x: x["text"] and x["text"].strip())
ds = ds.shuffle(seed=42)

# train / val / test = 80 / 10 / 10
train_test = ds.train_test_split(test_size=0.2, seed=42, stratify_by_column="label")
val_test = train_test["test"].train_test_split(test_size=0.5, seed=42)

final = DatasetDict({
    "train": train_test["train"],
    "validation": val_test["train"],
    "test": val_test["test"],
})

# 4 · Save locally + push to the Hub
final.save_to_disk("./reviews_dataset")
final.push_to_hub("my-org/reviews", private=True)