Build a Dataset with the Datasets Library

In this LLM Finetuning lesson, you'll learn to create a dataset using the Hugging Face datasets library—covering loading, preprocessing, and saving your data for fine-tuning. Hands-on steps and troubleshooting included.

Focus: build a dataset with the datasets library

Sponsored

You've spent hours collecting the perfect raw data for your LLM fine-tuning project — JSON files, CSVs, maybe a messy folder of text documents. But the moment you try to feed that pile into your training loop, everything falls apart: inconsistent formats, missing columns, memory errors, or a slow bottleneck that makes every epoch feel like a lifetime. Building a clean, efficient dataset is the foundation of any successful fine-tuning run, and the Hugging Face datasets library is the tool that turns chaos into a streamlined, battle-tested pipeline. In this lesson, you'll move from raw, unstructured data to a high-quality dataset ready for tokenization and training — step by step, with hands-on code you can run today.

The problem this lesson solves

Fine-tuning a pretrained LLM is not about clever model tweaks; it's about your data. Garbage in, garbage out has never been more literal: a model trained on inconsistent or noisy data will produce reliably wrong outputs. But before you even think about model quality, you face practical blockers:

  • Format chaos — Your data lives in JSON, CSV, text files, or a DataFrame, and you need one unified format.
  • Memory pressure — Loading a multi-gigabyte dataset into RAM can crash your machine or force you to buy a bigger one.
  • Preprocessing drag — Cleaning, filtering, and mapping transformations often become a maze of custom scripts.
  • Slow iteration — Without a standardized workflow, every experiment requires re-doing the same plumbing.

By the end of this lesson, you'll be able to build a dataset with the datasets library that solves all of these. You'll load raw data from multiple sources, clean it, transform it, and save it in an optimized format that your training loop can consume instantly. No more duct-taped pipelines.

Core concept / mental model

Think of the datasets library as a data warehouse for machine learning. Instead of writing ad hoc code to load files, shuffle rows, and apply transformations, you get a single, unified API that handles the heavy lifting.

At its heart is the Dataset object — think of it as a smart table backed by memory-mapped Apache Arrow storage. This gives you three superpowers:

  • Memory efficiency — The data lives on disk, not in RAM. You access slices or elements on demand, so you can work with datasets much larger than your memory.
  • Speed — Arrow's columnar format is incredibly fast for slicing, sampling, and iterating.
  • Consistency — Every Dataset has a schema (column names and types), so you catch mistakes early.

But that's just the container. The real magic is the pipeline: you load raw data, map transformations across it, filter out bad rows, split it into train/validation sets, and then save it in a format like Parquet that's perfect for training.

Here's the mental model to keep in your head:

You are building a factory — raw materials go in one end (JSON, CSV, text), and a polished, high-quality dataset comes out the other, ready for the tokenizer.

The datasets library is your conveyor belt. Each step — load_dataset, map, filter, train_test_split, save_to_disk — is a station on that belt. You control the sequence, but the heavy machinery is already built.

How it works step by step

Building a dataset with the datasets library follows a predictable pattern. Once you internalize these steps, you can apply them to any fine-tuning project.

  1. Install and import — Get the library and bring the core functions into your namespace.
  2. Load raw data — Use datasets.load_dataset() with a format string ('json', 'csv', 'text', etc.) and a file path.
  3. Inspect the dataset — Check the schema, number of rows, and a few examples to catch surprises early.
  4. Clean and transform — Apply Dataset.map() to process each example (e.g., tokenizing, lowercasing, reformatting). Use Dataset.filter() to drop rows that don't meet your criteria.
  5. Split the data — Use Dataset.train_test_split() to create train and validation sets, so you can monitor overfitting.
  6. Save for training — Use Dataset.save_to_disk() or convert to Parquet/CSV for fast reloading.
  7. Load in your training script — Pull the saved dataset back with load_from_disk() and hand it to your tokenizer.

That's the whole flow. The datasets library's design philosophy is that all of this should feel like a natural extension of Python — not a separate framework.

Hands-on walkthrough

Let's get our hands dirty. We'll build a dataset from scratch — starting with raw JSON, cleaning it, enriching it, and saving it as a Parquet file ready for fine-tuning.

Step 1: Setup

First, install the library (if you haven't) and import the essentials.

pip install datasets
from datasets import Dataset, load_dataset, load_from_disk

Step 2: Load raw JSON

Imagine you have a records.json file with conversational data from a customer support chatbot.

[
  {"input": "I need a refund", "output": "We understand. Could you provide your order ID?"},
  {"input": "My order hasn't arrived", "output": "Let's track it. What's your tracking number?"},
  {"input": "the app crashes on startup", "output": "Sorry for the issue. Have you tried reinstalling?"}
]

Load it and take a look:

ds = load_dataset("json", data_files="records.json", split="train")
print(ds)
print(ds[:1])

Expected output:

Dataset({
    features: ['input', 'output'],
    num_rows: 3
})
{'input': ['I need a refund'], 'output': ['We understand. Could you provide your order ID?']}

Notice how load_dataset with "json" guesses the schema. You can see the column names and row count instantly.

Step 3: Clean and transform with map

Now let's clean the data. For example, we might want to normalize whitespace and enforce lowercase for the input field (but keep the output capitalized).

def clean_example(example):
    example["input"] = example["input"].strip().lower()
    example["output"] = example["output"].strip()
    return example

cleaned_ds = ds.map(clean_example)
print(cleaned_ds["input"])

Expected output:

['I need a refund', 'My order hasn't arrived', 'the app crashes on startup']

The map function applies your Python function to every row, and it's fast because it's executed in Rust behind the scenes. It also supports batched=True for even more speed when processing many rows.

Step 4: Filter bad rows

Not all data is useful. Maybe you want to drop examples where the input is too short or contains placeholder text like [unknown].

def is_valid(example):
    return len(example["input"]) > 5 and "[unknown]" not in example["input"]

filtered_ds = cleaned_ds.filter(is_valid)
print(filtered_ds)

Expected output:

Dataset({
    features: ['input', 'output'],
    num_rows: 3
})

Here all rows pass, but in a real dataset this is where you'd remove noise.

Step 5: Split into train/validation

You need a held-out set to evaluate your fine-tuned model.

split_ds = filtered_ds.train_test_split(test_size=0.2, seed=42)
print(split_ds)

Expected output:

DatasetDict({
    train: Dataset({
        features: ['input', 'output'],
        num_rows: 2
    })
    test: Dataset({
        features: ['input', 'output'],
        num_rows: 1
    })
})

Now you have a DatasetDict — a dictionary of datasets for train and test. This is the standard structure for training.

Step 6: Save and reload

The final step is to persist our work.

# Save to disk (Arrow format, fast to reload)
split_ds.save_to_disk("my_dataset")

# Or save as Parquet for interop with other tools
split_ds["train"].to_parquet("train.parquet")

# Reload later
reloaded = load_from_disk("my_dataset")
print(reloaded)

Expected output:

DatasetDict({
    train: Dataset({
        features: ['input', 'output'],
        num_rows: 2
    })
    test: Dataset({
        features: ['input', 'output'],
        num_rows: 1
    })
})

You now have a clean, split, and saved dataset that you can load in any training script with just one line.

Compare options / when to choose what

The datasets library is not the only game in town. Here's how it stacks up:

Approach Pros Cons Best when…
datasets library Memory-mapped, fast, built-in splits & transformation, integrates with Hugging Face Hub Extra dependency, learning curve, might be overkill for tiny data Building a serious fine-tuning pipeline, working with large data
Pandas DataFrame Familiar syntax, great for exploration Loads everything into RAM, slower for big data, no built-in ML features Quick data exploration or small datasets (<1M rows)
Pure Python (lists/dicts) Zero dependencies, full control Manual everything, prone to bugs, no memory optimization Prototyping with a handful of examples
CSV/JSON on disk Universal format, easy to share Requires custom loading code, slow for repeated access, no schema enforcement Data exchange with other teams or systems

The rule of thumb: for any fine-tuning project beyond a toy example, use the datasets library. It gives you performance and structure that translate directly to a smoother training loop.

Variations worth knowing

  • load_dataset with data_files — You can pass a dictionary of file patterns to create a DatasetDict with custom splits (e.g., {'train': 'train/.json', 'test': 'test/.json'}). This is great for data that's already organized into folders.
  • Dataset.from_dict — If your data is already in Python memory (e.g., a list of dicts from an API), you can skip load_dataset and create a Dataset directly.
  • Preprocessing with tokenize_function — The map function is perfect for adding a 'labels' column after tokenization, which is the exact input format expected by Hugging Face Transformers.

Troubleshooting & edge cases

Even with a great library, things go wrong. Here are the most common issues and how to fix them.

Problem: load_dataset can't guess the file format

If your file has a non-standard extension (e.g., .txt with CSV content), the library might not parse it correctly.

Solution: Explicitly pass the format string:

ds = load_dataset("csv", data_files="data.txt", sep="\t")

Problem: Memory error when loading a huge file

Loading a 10 GB JSON file exhausts RAM.

Solution: Use Apache Arrow's memory mapping. The datasets library already does this under the hood, but you need to make sure you're not converting to Python lists or pandas unnecessarily. If you must, use Dataset.select or Dataset.shard to work with a subset. Also, consider converting your raw data to Parquet first — it's more compact and faster to load.

Problem: map is too slow

On large datasets, sequential mapping crawls.

Solution: Use batched=True and num_proc to process examples in batches and parallelize with multiprocessing:

ds = ds.map(clean_example, batched=True, num_proc=4)

Problem: Schema mismatch after map

If your transformation function returns a new field with a different type, you might see errors.

Solution: Check the dataset features after mapping:

print(ds.features)

If you added a field that's a list, make sure the library can infer the shape. Sometimes you need to specify the column types with Dataset.cast_column().

Problem: Duplicate rows or order not preserved

The datasets library may reorder or lose duplicates if you're not careful with filter and map.

Solution: Preserve order explicitly by using Dataset.map with keep_in_memory=True or by sorting afterwards. For deduplication, use Dataset.unique() on a column and then Dataset.filter to keep only first occurrences.

What you learned & what's next

You've just built a dataset with the datasets library — from raw JSON to a clean, split, and saved dataset ready for tokenization. You can now:

  • Load raw data from various formats using load_dataset.
  • Inspect dataset schema with print(ds) and ds.features.
  • Transform and clean data with Dataset.map and Dataset.filter.
  • Split your dataset into train and validation sets with train_test_split.
  • Save and reload your dataset with save_to_disk and load_from_disk.

You also know when to reach for the datasets library versus simpler tools like pandas, and you can troubleshoot the most common pitfalls.

The next lesson in this track is Tokenization and preprocessing for fine-tuning. You'll take the dataset you built here and feed it through a tokenizer, preparing input IDs and attention masks — the exact format your model expects. Your clean dataset is now the solid foundation for that step.

Practice recap

Take the dataset you built in this lesson and add a new column length containing the character length of the input field using map. Then filter out rows where length is under 10 characters. Finally, save the updated dataset to disk and reload it to confirm the schema. This mirrors the preprocessing you'll do with tokenizers next.

Common mistakes

  • Forgetting to specify the split when loading, resulting in a DatasetDict instead of a Dataset — use split="train".
  • Applying transformations that change the schema without checking ds.features, leading to silent type coercion or errors.
  • Loading a huge dataset into memory with Dataset.to_pandas() — defeats the memory-mapping advantage. Stay with ds[:n] or ds.select().
  • Saving the dataset as CSV/JSON instead of Parquet or save_to_disk, making reloads slower and less memory-efficient.
  • Ignoring the batched=True option in map, causing painfully slow preprocessing on large datasets.

Variations

  1. Use Dataset.from_dict when your data is already in Python lists or dicts to skip file I/O.
  2. Pass a dictionary of file patterns to load_dataset for pre-built train/test splits from separate folders.
  3. Use streaming=True in load_dataset for enormous datasets that exceed RAM, processing one example at a time.

Real-world use cases

  • Building a conversational dataset from customer support chat logs to fine-tune a domain-specific chatbot.
  • Creating a question-answer dataset from documentation pages to fine-tune a RAG retrieval model.
  • Aggregating scraped product reviews into a classification dataset for sentiment analysis fine-tuning.

Key takeaways

  • The datasets library provides a memory-mapped, Arrow-backed Dataset object that lets you handle large data without exhausting RAM.
  • The standard pipeline is load → map (transform) → filter (clean) → split → save — reusable for any fine-tuning project.
  • Using Dataset.map with batched=True dramatically speeds up preprocessing on large datasets.
  • Always split your data with train_test_split to create a validation set for monitoring overfitting.
  • Save your processed dataset with save_to_disk or Parquet for fast, efficient reloading in training scripts.
  • The datasets library integrates with Hugging Face Hub, making sharing and versioning datasets simple.

Sponsored

Sponsored

Discussion

Questions, corrections, and tips help everyone reading this page.

0 comments

Add a comment

Shown publicly with your comment.

Be constructive · max 4,000 characters

No comments yet — start the thread.

Related tutorials, quizzes, and articles for this topic.