The Datasets library

Setup#

pip install datasets

Tested with datasets 5.0. Nothing here downloads a model, and the data is five rows we write ourselves — the whole lesson runs offline after install.

Building a dataset from a file you own#

import csv
from datasets import load_dataset

rows = [
    {"text": "The biryani arrived cold",              "stars": 1},
    {"text": "Perfect filter coffee, came back twice", "stars": 5},
    {"text": "Service was slow but the dosa was good", "stars": 3},
    {"text": "Overpriced and the AC was broken",       "stars": 2},
    {"text": "Best vada pav in the neighbourhood",     "stars": 5},
]
with open("reviews.csv", "w", newline="", encoding="utf-8") as f:
    w = csv.DictWriter(f, fieldnames=["text", "stars"]); w.writeheader(); w.writerows(rows)

ds = load_dataset("csv", data_files="reviews.csv", split="train")
print(ds)
print("features :", ds.features)
print("one row  :", ds[0])
print("one column:", ds["stars"])
print("a slice  :", ds[:2]["text"])

good = ds.filter(lambda r: r["stars"] >= 4)
print("kept     :", good.num_rows, good["text"])
print(ds.train_test_split(test_size=0.4, seed=0))

Dataset({
    features: ['text', 'stars'],
    num_rows: 5
})
features : {'text': Value('string'), 'stars': Value('int64')}
one row  : {'text': 'The biryani arrived cold', 'stars': 1}
one column: Column([1, 5, 3, 2, 5])
a slice  : ['The biryani arrived cold', 'Perfect filter coffee, came back twice']
kept     : 2 Column(['Perfect filter coffee, came back twice', 'Best vada pav in the neighbourhood'])
DatasetDict({
    train: Dataset({
        features: ['text', 'stars'],
        num_rows: 3
    })
    test: Dataset({
        features: ['text', 'stars'],
        num_rows: 2
    })
})

Progress bars for the generating and filtering steps also appear on your terminal. They are written to the error stream, so they are left out above.

Naming your labels#

from datasets import Dataset, ClassLabel

ds = Dataset.from_dict({
    "text":  ["cold biryani", "great coffee", "slow service", "best vada pav"],
    "label": [0, 1, 0, 1],
})
ds = ds.cast_column("label", ClassLabel(names=["bad", "good"]))
print("features:", ds.features)
print("row 0   :", ds[0], "->", ds.features["label"].int2str(ds[0]["label"]))
print("in-memory (no cache files):", ds.cache_files)
print("arrow bytes:", ds.data.nbytes)

torch_ds = ds.with_format("torch")
print("torch format:", torch_ds[0])

features: {'text': Value('string'), 'label': ClassLabel(names=['bad', 'good'])}
row 0   : {'text': 'cold biryani', 'label': 0} -> bad
in-memory (no cache files): []
arrow bytes: 97
torch format: {'text': 'cold biryani', 'label': tensor(0)}

The walkthrough#

Indexing by row gives a dict; indexing by name gives a column. ds[0] is one example. ds["stars"] is every value in that column. ds[:2]["text"] slices rows first, then picks a column. Reading that pattern correctly saves a lot of confused debugging.

A Column is not a Python list. In datasets 5 the column accessor returns a lazy Column object. Iterate it, index it, or wrap it in list(...) when a plain list is what you need.

ClassLabel turns integers into names that travel with the data. Now int2str and str2int exist, train_test_split(stratify_by_column="label") becomes legal, and the Trainer can print readable metric names. Do this early — retrofitting label names after a confusing evaluation is the harder path.

from_dict builds in memory; load_dataset builds on disk. Notice cache_files is empty above. Data created from a Python dict lives in RAM. Data loaded from a file is written to an Arrow file and memory-mapped, which is where the bigger-than-RAM property comes from.

with_format("torch") changes what indexing returns, not what is stored. The label came back as a tensor while text stayed a string, because strings have no tensor form. with_format("torch", columns=["input_ids", "label"]) is the usual call before feeding a DataLoader.

Splitting is seeded, so it is reproducible. train_test_split(test_size=0.4, seed=0) gives the same two tables on every machine. Leaving the seed out gives a different split each run, and metrics that wobble for no visible reason — see train/test split.

Common mistakes#

Reaching for pandas out of habit. ds.to_pandas() works and is the right move for a plot or a quick group-by. It also loads everything into memory, which defeats the point on a large dataset. Convert a sample, not the table.

Expecting filter and map to modify in place. They return a new dataset and leave the original untouched. ds.filter(...) on its own line changes nothing. Assign the result.

Splitting after shuffling badly, or not at all. train_test_split shuffles by default. If your rows are grouped by class, an unshuffled split hands you a train set of one class. If rows are duplicated across the split, your test score is a memory test — see leakage.

Loading community datasets without reading them. A Hub dataset can be tens of gigabytes and carry a license that forbids your use. Check the size on the dataset page and the license field before load_dataset starts filling your disk.

Assuming loading scripts still run. Datasets 4 removed Python loading scripts from the Hub. Old tutorials that pass trust_remote_code=True now fail; those datasets have mostly been re-published as plain Parquet, which loads with no argument at all.

Try it yourself#

Add a length column with ds.add_column("length", [len(t) for t in ds["text"]]), then sort by it with ds.sort("length"). Print ds.cache_files before and after. Which operations wrote a new Arrow file, and which returned a view?

What to learn next#

• map, batched and the cache — turning raw text into model inputs, once.
• Data collators — the step between a Dataset and a batch.
• Datasets too big for RAM — the same problem, solved from the PyTorch side.