Lesson 18 / 33

Loading Data with the datasets Library

Load Hub, CSV and JSON data into a DatasetDict and stream data too large for memory.

By · 15 years in web development

Arrow-backed tables with splits

load_dataset() loads data from a Hub dataset id or from local files such as CSV, JSON Lines or Parquet. The result is usually a DatasetDict, a dictionary of named splits like train, validation and test, each a Dataset. A Dataset behaves like a table: ds.column_names, ds.features, len(ds), ds[0] for one row and ds["text"] for a column. Under the hood the data lives in Apache Arrow files that are memory-mapped from disk, so a dataset bigger than your RAM can still be indexed quickly without loading it all. Processed results are cached on disk too, so rerunning the same transformation is instant. For truly huge corpora, streaming=True returns an IterableDataset that reads examples lazily over the network; you lose random access and len, but you can start training without downloading everything. If a dataset has no validation split, create one with train_test_split and a fixed seed so your evaluation is repeatable.

Hub data, local files and a validation split

The seed makes the split identical on every run.

from datasets import load_dataset

imdb = load_dataset("stanfordnlp/imdb")          # DatasetDict: train, test, unsupervised
print(imdb)
print(imdb["train"][0]["text"][:80], imdb["train"][0]["label"])

local = load_dataset("csv", data_files={"train": "reviews.csv"})
split = local["train"].train_test_split(test_size=0.1, seed=42)
train_ds, val_ds = split["train"], split["test"]

stream = load_dataset("HuggingFaceFW/fineweb", split="train", streaming=True)
print(next(iter(stream))["text"][:80])

Check label balance before training

Run ds.to_pandas()["label"].value_counts() on a sample before you train. A 95:5 split between classes means plain accuracy will look great even for a useless model; use F1 or per-class metrics instead.

Quick check: Why can the datasets library work with a dataset larger than your laptop's RAM?

  • Arrow files are memory-mapped from disk, so rows are read on demand
  • It compresses the data into GPU memory
  • It automatically drops rows that do not fit
  • It always streams data from the Hub over the network
Answer

Arrow files are memory-mapped from disk, so rows are read on demand — Datasets are stored as Arrow files and memory-mapped. Streaming is a separate, optional mode.