Hugging Face Datasets Library: Loading and Processing Data

1. Overview

When working with NLP models or fine-tuning LLMs, you always need a prior step: preparing data in the right format, making it clean, and loading it into your training pipeline. Hugging Face’s datasets library solves exactly this problem — it gives you a unified API to load datasets from the Hub, batch-process them with map/filter, and save them in a format optimized for repeated reads.

After reading this guide, you’ll be able to load datasets from Hugging Face Hub, understand the Dataset/DatasetDict structure, preprocess text with tokenizers, handle large datasets with streaming, and write tests to validate your data processing pipeline.

2. Installation

pip install datasets transformers

The datasets library doesn’t strictly depend on transformers, but the tokenization section later will need it.

import datasets
print(datasets.__version__)

3.0.1

3. Loading datasets from Hugging Face Hub

3.1. Basic load_dataset

from datasets import load_dataset

reviews = load_dataset("imdb")
print(reviews)

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
    })
})

load_dataset fetches the script/metadata for the imdb dataset from the Hub, caches the compressed files locally (~/.cache/huggingface/datasets), and returns a DatasetDict — essentially a dictionary of Dataset objects organized by split (train, test, unsupervised).

3.2. Loading a specific split

train_reviews = load_dataset("imdb", split="train")
print(type(train_reviews), len(train_reviews))

<class 'datasets.arrow_dataset.Dataset'> 25000

Passing split="train" returns a single Dataset directly instead of a DatasetDict, which is convenient when you only need one split.

Note: If a dataset on the Hub requires a config (e.g., glue), you must pass the config name as well: load_dataset("glue", "mrpc"). Without this argument, you’ll get an error asking you to select a config.

4. Exploring Dataset structure

sample = train_reviews[0]
print(sample.keys())
print(sample["label"], sample["text"][:60])

dict_keys(['text', 'label'])
0 I rented I AM CURIOUS-YELLOW from my video store because of...

Accessing by index returns a regular Python dict. Dataset also supports slicing, which returns a dict of lists:

batch = train_reviews[0:3]
print(batch["label"])

[0, 0, 0]

View the schema using features:

print(train_reviews.features)

{'text': Value(dtype='string', id=None),
 'label': ClassLabel(names=['neg', 'pos'], id=None)}

ClassLabel tells you that 0 corresponds to the neg label and 1 to pos — useful when mapping back to label names.

5. Processing data with map, filter, select

5.1. Using filter to filter rows

positive_only = train_reviews.filter(lambda row: row["label"] == 1)
print(len(positive_only))

12500

filter takes a function that returns True/False for each row, iterates through the entire dataset, and keeps only rows that satisfy the condition.

5.2. Using map to transform data

def add_length(row):
    row["text_length"] = len(row["text"])
    return row

train_with_length = train_reviews.map(add_length)
print(train_with_length[0]["text_length"])

1640

map applies a function to each row, returning a dict with new fields (or updated existing fields) — the results are merged back into the original dataset.

5.3. Using map with batches for speed

def add_length_batched(batch):
    batch["text_length"] = [len(t) for t in batch["text"]]
    return batch

train_fast = train_reviews.map(add_length_batched, batched=True, batch_size=1000)

With batched=True, the function receives an entire batch dict (each field is a list) instead of individual rows — this avoids calling the Python function 25,000 times and is significantly faster with large datasets.

5.4. Using select to get a subset

small_train = train_reviews.select(range(100))
print(len(small_train))

100

Use this when you want to quickly test your pipeline on a subset before running it on the full dataset.

6. Text preprocessing with tokenizer

The most common real-world use case: tokenizing text to feed into a model.

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

def tokenize_batch(batch):
    return tokenizer(batch["text"], truncation=True, max_length=256)

tokenized = train_reviews.map(tokenize_batch, batched=True)
print(tokenized.column_names)

['text', 'label', 'input_ids', 'attention_mask']

tokenize_batch takes a batch and returns a dict with input_ids and attention_mask — map automatically merges these new fields into the dataset while keeping existing fields. To drop unnecessary fields (e.g., raw text to save RAM), use remove_columns:

tokenized = train_reviews.map(
    tokenize_batch, batched=True, remove_columns=["text"]
)
print(tokenized.column_names)

['label', 'input_ids', 'attention_mask']

7. Splitting, shuffling, and sharding

7.1. Creating your own train/test split

If a dataset has only one split, you can split it yourself with train_test_split:

split_result = train_reviews.train_test_split(test_size=0.2, seed=42)
print(split_result)

DatasetDict({
    train: Dataset({
        features: ['text', 'label'],
        num_rows: 20000
    })
    test: Dataset({
        features: ['text', 'label'],
        num_rows: 5000
    })
})

seed=42 ensures reproducible splitting across runs.

7.2. Shuffle and shard

shuffled = train_reviews.shuffle(seed=42)
part = shuffled.shard(num_shards=4, index=0)
print(len(part))

6250

shard splits the dataset into num_shards equal parts — useful when distributing data across multiple workers for distributed training.

8. Streaming for large datasets

When a dataset is too large to fit on disk or in RAM, use streaming=True:

streamed = load_dataset("imdb", split="train", streaming=True)
first_three = list(streamed.take(3))
print(len(first_three))

3

With streaming, datasets doesn’t download the entire file but reads records sequentially through an iterator. You can’t use len() or direct indexing, but you can still map/filter lazily (applied only when iterating).

Note: Streaming is best for datasets with hundreds of GB, like pretraining sets. For small datasets (a few hundred MB), regular loading (streaming=False) is faster because it leverages caching and random access.

9. Saving and reloading datasets

After processing, save the result so you don’t have to re-run map every time:

tokenized.save_to_disk("data/imdb_tokenized")

from datasets import load_from_disk

reloaded = load_from_disk("data/imdb_tokenized")
print(len(reloaded))

25000

save_to_disk writes the dataset in Arrow format — reloading is very fast because no reparsing from the original source is needed. If you need to interact with pandas, convert directly:

df = tokenized.select(range(5)).to_pandas()
print(df.columns.tolist())

['label', 'input_ids', 'attention_mask']

10. Writing tests for your data processing pipeline

Your data processing pipeline should have tests to ensure the map/filter logic doesn’t silently break when you modify code. Here’s a test suite using pytest.

import pytest
from datasets import Dataset


def add_length(row):
    row["text_length"] = len(row["text"])
    return row


def test_add_length_computes_correct_value():
    data = Dataset.from_dict({"text": ["hello", "hi there"]})

    result = data.map(add_length)

    assert result[0]["text_length"] == 5
    assert result[1]["text_length"] == 8


def test_filter_keeps_only_positive_labels():
    data = Dataset.from_dict({"label": [0, 1, 1, 0]})

    positive = data.filter(lambda row: row["label"] == 1)

    assert len(positive) == 2
    assert all(row["label"] == 1 for row in positive)

Dataset.from_dict creates a small in-memory dataset without needing to fetch anything from the Hub — perfect for unit tests. The two tests above verify that map works correctly (creating new fields) and filter keeps only matching rows, without depending on real datasets from the internet.

pytest test_dataset_pipeline.py -v

test_dataset_pipeline.py::test_add_length_computes_correct_value PASSED
test_dataset_pipeline.py::test_filter_keeps_only_positive_labels PASSED

11. Common pitfalls and when not to use it

• Forgetting batched=True for large datasets: Non-batched map is slow because it calls Python for each row individually. With hundreds of thousands of rows, always prefer batched=True if your logic allows vectorization.
• Using random indexing on streamed datasets: streamed[5] will raise an error because streamed datasets are iterators that don’t support random access. Use only .take(), .skip(), or sequential iteration.
• Not setting seed with shuffle/train_test_split: results won’t be reproducible across runs, making it hard to debug when comparing metrics between experiments.
• Uncontrolled cache bloat: each time you map with a different function, datasets writes a new cache file. If your machine has limited disk space, use load_dataset(..., keep_in_memory=True) or periodically clear ~/.cache/huggingface/datasets.
• Don’t use datasets for complex relational data (multiple tables with joins, transactions). This library is optimized for flat tabular/text data used in ML, not as a database replacement.

12. Conclusion

The datasets library gives you a clean API to go from “raw dataset from the Hub” to “tokenized dataset ready for training” in just a few load_dataset, map, and filter calls. For small datasets, load normally and leverage Arrow caching; for huge datasets, switch to streaming=True to avoid downloading everything to disk. When building real data pipelines, always write tests for your map/filter functions using Dataset.from_dict — this helps you catch unintended changes to your preprocessing logic early.