A dataset that behaves like a list and is not one

Loading, and what comes back

from datasets import load_dataset

ds = load_dataset("imdb")
print(ds)

DatasetDict({
    train: Dataset({features: ['text', 'label'], num_rows: 25000}),
    test:  Dataset({features: ['text', 'label'], num_rows: 25000}),
    unsupervised: Dataset({features: ['text', 'label'], num_rows: 50000}),
})

A DatasetDict is a dictionary of splits. Each Dataset has features (named, typed columns) and a row count. Ask for one split directly and you get the Dataset alone:

train = load_dataset("imdb", split="train")
print(train[0])                    # one row, as a dict
print(train[:3]["text"][0][:80])   # a slice returns a dict of lists
print(train.features)

That last line is worth internalising. Indexing a row gives a dict; slicing gives a dict of lists. train[0] is {'text': ..., 'label': 1}; train[:3] is {'text': [...], 'label': [...]}. Half the confusing errors in a first hour with this library come from expecting a list of dicts.

Features are typed, and the types do work

print(train.features["label"])
# ClassLabel(names=['neg', 'pos'], id=None)

print(train.features["label"].int2str(1))   # 'pos'

ClassLabel carries the names, so you never have to guess whether 1 means positive: the same problem id2label solves on the model side. Other types you will meet are Value("string"), Value("int32"), Sequence, Image, and Audio, and the last two decode lazily, which the media lesson returns to.

Split slicing without writing Python

small = load_dataset("imdb", split="train[:1%]")
middle = load_dataset("imdb", split="train[100:200]")
both  = load_dataset("imdb", split="train+test")

train[:1%] is the habit worth forming. Develop on one per cent, then remove the suffix. A pipeline debugged on 250 rows runs unchanged on 25,000, and you will have saved yourself ten minutes per mistake.

Splitting when the dataset has no split

parts = ds["train"].train_test_split(test_size=0.1, seed=42, stratify_by_column="label")

Three things here are deliberate. seed makes it reproducible, so the numbers you quote next week come from the same split. stratify_by_column keeps class proportions equal across the two halves, which matters enormously when a class is rare: an unstratified 10% test split of a dataset with 2% positives can easily land with almost none, and your recall figure then measures nothing.

And a validation set is a third split, not a second use of the test set:

tmp = ds["train"].train_test_split(test_size=0.2, seed=42)
val_test = tmp["test"].train_test_split(test_size=0.5, seed=42)
# tmp['train'] = 80% train, val_test['train'] = 10% val, val_test['test'] = 10% test

You tune on validation. You touch test once, at the end. Every time you look at the test score and change something, it becomes a little more like a training set, and the number it reports becomes a little more optimistic than reality.

Loading your own files

No upload required: load_dataset takes formats directly:

ds = load_dataset("csv", data_files="reviews.csv")
ds = load_dataset("json", data_files={"train": "tr.jsonl", "test": "te.jsonl"})
ds = load_dataset("parquet", data_files="data/*.parquet")
ds = load_dataset("imagefolder", data_dir="photos/")   # labels from folder names

So the whole library is available for local work. You do not need to publish anything to get streaming, mapping, caching and typed columns.

Configurations, when one name holds many datasets

Some repositories hold several variants, and the second positional argument chooses between them:

from datasets import get_dataset_config_names
print(get_dataset_config_names("mozilla-foundation/common_voice_17_0")[:5])

hi = load_dataset("mozilla-foundation/common_voice_17_0", "hi", split="train")

For multilingual corpora the configuration is usually the language, and loading without it either fails with a list of options or silently gives you the default, which, more often than is comfortable, is English. Print the configuration names before you assume.

Sorting, shuffling, selecting

ds = ds.shuffle(seed=42)
small = ds.select(range(1000))
pos = ds.filter(lambda r: r["label"] == 1)
ds = ds.remove_columns(["unused"]).rename_column("text", "sentence")

None of these copy the data. select and shuffle build an indices map (a list of row positions) and the underlying table stays where it is. This is why shuffling a 40 GB dataset is instant and takes no memory, and it is a genuine surprise the first time.

The exception is filter, which has to read every row to decide. On a large dataset, pass num_proc to parallelise it.

Do this now

Load any dataset with split="train[:1%]", print features, print row zero, and print a slice of three. Then run train_test_split with a seed and check that running it twice gives the same first row. That reproducibility is what makes every later comparison meaningful.