Hugging FaceのDatasetsライブラリ: データ読み込みと処理

1. 概要

NLPモデルやLLMのファインチューニングを行う際、常に必要な準備がある。データを正しい形式に整え、十分にクリーニングして、学習パイプラインに読み込める状態にすることだ。Hugging FaceのDatasetsライブラリはこの問題をまさに解決する。Hubからデータセットを読み込み、map/filterで一括処理し、繰り返し読み込みに最適な形式で保存するための統一されたAPIを提供している。

この記事を読み終わったら、Hugging Face Hubからデータセットを読み込み、Dataset/DatasetDictの構造を理解し、tokenizerでテキストを前処理し、streamingで大規模データセットを扱い、データ処理パイプラインをテストするコードが書けるようになるだろう。

2. インストール

pip install datasets transformers

Datasetsライブラリはtransformersに強く依存していないが、セクション6のtokenize部分では必要になる。

import datasets
print(datasets.__version__)

3.0.1

3. Hugging Face Hubからのデータセット読み込み

3.1. 基本的な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はHubからdataset imdbのスクリプト/メタデータを読み込み、圧縮ファイルをローカル（~/.cache/huggingface/datasets）にキャッシュして、DatasetDictを返す。つまり、split（train、test、unsupervised）ごとのDatasetの辞書である。

3.2. 特定のsplitだけを取得する

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

<class 'datasets.arrow_dataset.Dataset'> 25000

split="train"を渡すと、DatasetDictの代わりに直接Datasetを返す。1つのセットだけが必要な場合に便利だ。

注意: Hub上のデータセットが設定（例えばglue）を要求する場合、設定名を追加で渡す必要がある: load_dataset("glue", "mrpc")。このパラメータがないとエラーが発生し、設定の選択を要求される。

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

インデックスでアクセスすると、通常のPython辞書が返される。Datasetはslicingもサポートし、listの辞書が返される：

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

[0, 0, 0]

featuresでスキーマを確認する：

print(train_reviews.features)

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

ClassLabelは0がラベルnegに対応し、1がposに対応することを示してくれる。ラベルを名前にマッピングバックする際に便利だ。

5. map、filter、selectでのデータ処理

5.1. 行をフィルタリングするfilter

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

12500

filterは各行に対してTrue/Falseを返す関数を受け取り、データセット全体を走査して、条件を満たす行のみを保持する。

5.2. データを変換するmap

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は各行に関数を適用し、新しいフィールドを含む辞書を返す（または既存フィールドを更新する）。結果は元のデータセットにマージされる。

5.3. バッチ処理で高速化するmap

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)

batched=Trueでは、関数は各行ではなくbatch全体の辞書（各フィールドはlist）を受け取る。Python関数を25,000回呼び出すのを避け、大規模データセットではかなり高速化される。

5.4. 部分集合を取得するselect

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

100

フルデータセットを実行する前に、部分集合でパイプラインをすばやくテストしたいときに使用する。

6. tokenizerによるテキストの前処理

最も一般的な実際のユースケースはモデルに入力するためにテキストをtokenizeすることだ。

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はbatchを受け取り、input_idsとattention_maskを含む辞書を返す。mapはこれらの新しいフィールドを自動的にデータセットにマージし、既存フィールドは保持される。不要なフィールド（例えば、RAMを節約するため生のtext）を削除したい場合はremove_columnsを使う：

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

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

7. split、shuffle、shardの分割

7.1. train/testの手動分割

データセットがsplitを1つだけ持つ場合、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は実行ごとに分割結果を再現可能にする。

7.2. ShuffleとShard

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

6250

shardはデータセットをnum_shardsの等しい部分に分割する。分散学習の複数のworkerにデータを配布するときに便利だ。

8. 大規模データセットのストリーミング

データセットがdisk/RAMに完全に読み込むには大きすぎる場合、streaming=Trueを使う：

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

3

streamingでは、datasetsはファイル全体をダウンロードせず、iteratorを介して順番に各レコードを読む。len()や直接indexは使えないが、lazy方式のmap/filterは使える（iterateするときだけ適用される）。

注意: Streamingはpretrainingセットのような数百GBのデータセットに適している。小規模なデータセット（数百MB）では、キャッシュとランダムアクセスが利用できるため、通常の読み込み（streaming=False）の方が高速である。

9. データセットの保存と再読み込み

処理が終わったら、結果を保存して、毎回mapを実行し直す必要がないようにする：

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はデータセットをArrow形式で書き込む。元のソースから再度解析する必要がないため、読み込みは非常に高速だ。pandasと相互作用する必要がある場合は、直接変換できる：

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

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

10. データ処理パイプラインのテスト作成

データ処理パイプラインはテストを持つべきだ。コード修正時にmap/filterのロジックが静かに壊れないことを確認するためである。以下は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はメモリ内に小さなデータセットを作成し、Hubから何も読み込む必要がない。unit testに非常に適している。上記の2つのテストは、インターネット上の実際のデータセットに依存することなく、map（新しいフィールドを生成）とfilter（正しい条件でフィルタリング）の正しい動作を確認する。

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. よくあるエラーと使用してはいけない場合

• 大規模データセット処理時にbatched=Trueを忘れる: mapがbatch処理しないと遅くなる。各行に対してPython関数を呼び出すためだ。数十万行の場合、ロジックがベクトル化を許す場合は常にbatched=Trueを優先する。
• streaming datasetでランダムindexを使う: streamed[5]はエラーを発生させる。streaming datasetはiteratorであり、random accessをサポートしないからだ。.take()、.skip()、または順序通りのiterateのみを使用する。
• shuffle/train_test_split時にseedを設定しない: 実行間で結果が再現不可能になり、実験間でメトリクスを比較するときのデバッグが困難になる。
• 管理されていないキャッシュが膨れ上がる: 異なる関数でmapするたびに、datasetsは新しいキャッシュファイルを書き込む。ディスク容量が限られている場合はload_dataset(..., keep_in_memory=True)を使うか、定期的に~/.cache/huggingface/datasetsを削除する。
• 複雑な関連データにdatasetsを使うべきではない（多くのテーブルが結合され、トランザクションが必要な場合）。これはML用の平坦な表形式/テキストデータに最適化されたライブラリであり、データベースの代わりではない。

12. まとめ

Datasetsライブラリは、「HubからRawデータセットを読み込む」から「tokenize済みで学習準備完了のデータセット」へと進むための整潔なAPIを提供している。わずかload_dataset、map、filterコマンドで可能だ。小規模なデータセットの場合は通常通り読み込みArrowキャッシュを活用する。巨大なデータセットの場合はstreaming=Trueに切り替えてdiskに全部ダウンロードしなくてよくする。実際のプロジェクトのデータ処理パイプラインを構築する際、Dataset.from_dictを使用してmap/filter関数のテストを書くことを忘れずに。これはあなたが前処理ロジックが意図しない変更に気づくのに役立つだろう。