Datasets Library Của Hugging Face: Tải Và Xử Lý Dữ Liệu

1. Tổng quan

Khi làm việc với model NLP hay fine-tune LLM, bạn luôn cần một bước trước đó: chuẩn bị dữ liệu đúng format, đủ sạch, và nạp được vào pipeline huấn luyện. Thư viện datasets của Hugging Face giải quyết đúng vấn đề này — nó cho bạn một API thống nhất để tải dataset từ Hub, xử lý hàng loạt bằng map/filter, và lưu lại ở dạng tối ưu cho việc đọc lặp lại.

Đọc xong bài này, bạn sẽ tải được dataset từ Hugging Face Hub, hiểu cấu trúc Dataset/DatasetDict, biết cách tiền xử lý văn bản bằng tokenizer, xử lý dataset lớn bằng streaming, và viết được test kiểm tra pipeline xử lý dữ liệu.

2. Cài đặt

pip install datasets transformers

Thư viện datasets không phụ thuộc cứng vào transformers, nhưng phần tokenize ở mục 6 cần đến nó.

import datasets
print(datasets.__version__)

3.0.1

3. Tải dataset từ Hugging Face Hub

3.1. load_dataset cơ bản

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 tải script/metadata của dataset imdb từ Hub, cache file nén về local (~/.cache/huggingface/datasets), rồi trả về một DatasetDict — tức dictionary các Dataset theo split (train, test, unsupervised).

3.2. Chỉ lấy một split cụ thể

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

<class 'datasets.arrow_dataset.Dataset'> 25000

Truyền split="train" trả về trực tiếp một Dataset thay vì DatasetDict, tiện khi bạn chỉ cần một tập.

Lưu ý: Nếu dataset trên Hub yêu cầu config (ví dụ glue), bạn phải truyền thêm tên config: load_dataset("glue", "mrpc"). Thiếu tham số này sẽ ném lỗi yêu cầu chọn config.

4. Khám phá cấu trúc Dataset

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...

Truy cập bằng index trả về một dict Python bình thường. Dataset hỗ trợ cả slicing trả về dict các list:

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

[0, 0, 0]

Xem schema bằng features:

print(train_reviews.features)

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

ClassLabel cho bạn biết 0 tương ứng nhãn neg, 1 là pos — hữu ích khi mapping ngược sang tên nhãn.

5. Xử lý dữ liệu với map, filter, select

5.1. filter để lọc hàng

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

12500

filter nhận một hàm trả True/False cho mỗi row, chạy qua toàn bộ dataset và chỉ giữ lại các row thỏa điều kiện.

5.2. map để biến đổi dữ liệu

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 áp dụng hàm lên từng row, trả về dict chứa field mới (hoặc cập nhật field cũ) — kết quả được gộp lại vào dataset gốc.

5.3. map theo batch để tăng tốc

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)

Với batched=True, hàm nhận cả một batch dict (mỗi field là list) thay vì từng row — tránh gọi hàm Python 25.000 lần, nhanh hơn đáng kể với dataset lớn.

5.4. select để lấy subset

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

100

Dùng khi bạn muốn thử nghiệm pipeline nhanh trên một tập con trước khi chạy full dataset.

6. Tiền xử lý văn bản với tokenizer

Trường hợp thực tế phổ biến nhất: tokenize text để đưa vào 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 nhận một batch và trả dict chứa input_ids, attention_mask — map tự động gộp các field mới này vào dataset, giữ nguyên field cũ. Muốn bỏ field không cần (ví dụ text thô để tiết kiệm RAM), dùng remove_columns:

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

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

7. Chia split, shuffle, shard

7.1. Tự chia train/test

Nếu dataset chỉ có một split, bạn tự chia bằng 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 đảm bảo kết quả chia tái lập được giữa các lần chạy.

7.2. Shuffle và shard

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

6250

shard chia dataset thành num_shards phần đều nhau — hữu ích khi phân phối dữ liệu cho nhiều worker huấn luyện phân tán.

8. Streaming cho dataset lớn

Khi dataset quá lớn để tải hết vào disk/RAM, dùng streaming=True:

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

3

Với streaming, datasets không tải toàn bộ file về mà đọc tuần tự từng record qua iterator. Bạn không dùng được len() hay index trực tiếp, nhưng vẫn map/filter được theo kiểu lazy (chỉ áp dụng khi iterate tới).

Lưu ý: Streaming phù hợp với dataset hàng trăm GB như các tập pretraining. Với dataset nhỏ (vài trăm MB), tải thường (streaming=False) nhanh hơn vì tận dụng được cache và truy cập ngẫu nhiên.

9. Lưu và tải lại dataset

Sau khi xử lý xong, lưu kết quả để không phải chạy lại map mỗi lần:

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 ghi dataset ở format Arrow — đọc lại rất nhanh vì không cần parse lại từ nguồn gốc. Nếu cần tương tác với pandas, chuyển đổi trực tiếp:

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

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

10. Viết test cho pipeline xử lý dữ liệu

Pipeline xử lý dữ liệu nên có test để đảm bảo logic map/filter không âm thầm hỏng khi bạn sửa code. Dưới đây là test bằng 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 tạo dataset nhỏ trong memory, không cần tải gì từ Hub — rất phù hợp cho unit test. Hai test trên kiểm tra đúng hành vi của map (sinh field mới) và filter (lọc đúng điều kiện) mà không phụ thuộc vào dataset thật ngoài 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. Lỗi hay gặp và khi nào không nên dùng

• Quên batched=True khi xử lý dataset lớn: map không batch chạy chậm vì gọi hàm Python cho từng row một. Với hàng trăm nghìn row, luôn ưu tiên batched=True nếu logic cho phép vectorize.
• Dùng index ngẫu nhiên trên dataset streaming: streamed[5] sẽ ném lỗi vì streaming dataset là iterator, không hỗ trợ random access. Chỉ dùng .take(), .skip(), hoặc iterate tuần tự.
• Không set seed khi shuffle/train_test_split: kết quả không tái lập được giữa các lần chạy, gây khó debug khi so sánh metric giữa các thí nghiệm.
• Cache phình to không kiểm soát: mỗi lần map với hàm khác nhau, datasets ghi một file cache mới. Nếu máy bạn ít disk, dùn load_dataset(..., keep_in_memory=True) hoặc định kỳ xóa ~/.cache/huggingface/datasets.
• Không nên dùng datasets cho dữ liệu quan hệ phức tạp (nhiều bảng join với nhau, cần transaction). Đây là thư viện tối ưu cho dữ liệu dạng bảng/text phẳng dùng trong ML, không phải thay thế cho database.

12. Kết luận

Thư viện datasets cho bạn một API gọn để đi từ “tải dataset thô từ Hub” đến “dataset đã tokenize, sẵn sàng huấn luyện” chỉ qua vài lệnh load_dataset, map, filter. Với dataset nhỏ, bạn tải bình thường và tận dụng cache Arrow; với dataset khổng lồ, chuyển sang streaming=True để không phải tải hết về disk. Khi xây pipeline xử lý dữ liệu cho dự án thật, đừng quên viết test cho các hàm map/filter bằng Dataset.from_dict — việc này giúp bạn phát hiện sớm khi logic tiền xử lý bị thay đổi ngoài ý muốn.