Datasets documentation
ストリーム
ストリーム
Dataset streamingを使用すると、ダウンロードせずにデータセットを扱うことができます。 データセットを反復処理する際に、データがストリーミングされます。 これは、次のような場合に特に役立ちます。
- 非常に大きなデータセットのダウンロードを待たされるのは望ましくありません。
- データセットのサイズが、コンピューター上の利用可能なディスク容量を超えています。
- データセットのサンプルをいくつかだけを素早く探索したい場合

例えば、HuggingFaceFW/fineweb データセットの英語分割は45テラバイトですが、ストリーミングを使用すればすぐに利用できます。以下に示すように、load_dataset() で streaming=True を設定することで、データセットをストリーミングできます。
>>> from datasets import load_dataset
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True)
>>> print(next(iter(dataset)))
{'text': 'How AP reported in all formats from tornado-stricken regionsMarch 8, 2012\nWhen the first serious bout of tornadoes of 2012 blew through middle America in the middle of the night, they touched down in places hours from any AP bureau...', ...,
'language_score': 0.9721424579620361, 'token_count': 717}Dataset streamingを使用すると、変換なしでローカルファイルから構成されるデータセットを扱うことができます。 この場合、データセットを反復処理する際に、ローカルファイルからデータがストリーミングされます。 これは、次のような場合に特に役立ちます。
- 非常に大きなローカルデータセットのArrow変換を待たされるのは望ましくありません。
- 変換後のファイルサイズが、お使いのコンピューターの利用可能なディスク容量を超えます。
- データセットのサンプルをいくつかだけを素早く探索したい場合
- Parquetデータセットの特定の列のみを読み込む、または効率的にフィルタリングしたい場合
例えば、oscar-corpus/OSCAR-2201のような、数百の圧縮済みJSONLファイルからなるローカルデータセットをストリーミングして、すぐに使用できます。
>>> from datasets import load_dataset
>>> data_files = {'train': 'path/to/OSCAR-2201/compressed/en_meta/*.jsonl.gz'}
>>> dataset = load_dataset('json', data_files=data_files, split='train', streaming=True)
>>> print(next(iter(dataset)))
{'id': 0, 'text': 'Founded in 2015, Golden Bees is a leading programmatic recruitment platform dedicated to employers, HR agencies and job boards. The company has developed unique HR-custom technologies and predictive algorithms to identify and attract the best candidates for a job opportunity.', ...Parquetはカラム形式のフォーマットであり、カラムのサブセットのみをストリーミングおよび読み込み、不要なカラムを無視できます。Parquetはカラムの統計情報(ファイルレベルおよび行グループレベル)などのメタデータも保存するため、効率的なフィルタリングが可能です。datasets.packaged_modules.parquet.ParquetConfigのcolumnsおよびfilters引数を使用して、Parquetデータセットをストリーミングし、カラムを選択してフィルタを適用します:
>>> from datasets import load_dataset
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True, columns=["url", "date"])
>>> print(next(iter(dataset)))
{'url': 'http://%[email protected]/Content/Press-Release/2012/How-AP-reported-in-all-formats-from-tornado-stricken-regions', 'date': '2013-05-18T05:48:54Z'}
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True, filters=[("language_score", ">=", 0.99)])
>>> print(next(iter(dataset)))
{'text': 'Everyone wishes for something. And lots of people believe they know how to make their wishes come true with magical thinking.\nWhat is it? "Magical thinking is a belief in forms of causation, with no known physical basis," said Professor Emily Pronin of Princeton...', ...,
'language_score': 0.9900368452072144, 'token_count': 716}ストリーミングモードでデータセットを読み込むと、従来のDatasetオブジェクトの代わりに、IterableDatasetと呼ばれる新しいデータセット型のインスタンスが作成されます。 この特別なデータセット型には、以下に示す独自の処理メソッドのセットがあります。
IterableDataset は、モデルの学習のような反復的なタスクに有用です。 ランダムアクセスが必要なタスクには IterableDataset を使用しないでください。for ループで全体を反復処理する必要があり、イテラブルデータセットの最後の例を取得するには、それ以前のすべての例を反復処理しなければなりません。 詳細は、Dataset と IterableDataset のガイドを参照してください。
カラムインデックス
特定の列の値を反復処理したい場合があるかもしれません。幸いなことに、IterableDatasetは列のインデックス指定をサポートしています:
>>> from datasets import load_dataset
>>> dataset = load_dataset("allenai/c4", "en", streaming=True, split="train")
>>> print(next(iter(dataset["text"])))
Beginners BBQ Class Taking Place in Missoula!...Datasetから変換
既存のDatasetオブジェクトがある場合、to_iterable_dataset()関数を使用してIterableDatasetに変換できます。これは、load_dataset()でstreaming=True引数を設定するよりも実際には高速です。データがローカルファイルからストリーミングされるためです。
>>> from datasets import load_dataset
# faster 🐇
>>> dataset = load_dataset("ethz/food101")
>>> iterable_dataset = dataset.to_iterable_dataset()
# slower 🐢
>>> iterable_dataset = load_dataset("ethz/food101", streaming=True)The to_iterable_dataset() function supports sharding when the IterableDataset is instantiated. This is useful when working with big datasets, and you’d like to shuffle the dataset or to enable fast parallel loading with a PyTorch DataLoader.
>>> import torch
>>> from datasets import load_dataset
>>> dataset = load_dataset("ethz/food101")
>>> iterable_dataset = dataset.to_iterable_dataset(num_shards=64) # shard the dataset
>>> iterable_dataset = iterable_dataset.shuffle(buffer_size=10_000) # shuffles the shards order and use a shuffle buffer when you start iterating
dataloader = torch.utils.data.DataLoader(iterable_dataset, num_workers=4) # assigns 64 / 4 = 16 shards from the shuffled list of shards to each worker when you start iteratingシャッフル
通常のDatasetオブジェクトと同様に、IterableDatasetもIterableDataset.shuffle()でシャッフルできます。
buffer_size 引数は、ランダムにサンプルを抽出するためのバッファのサイズを制御します。データセットに100万のサンプルがあり、buffer_size を1万に設定したとします。IterableDataset.shuffle() は、バッファ内の最初の1万のサンプルからランダムにサンプルを選択します。バッファ内の選択されたサンプルは、新しいサンプルと置き換えられます。デフォルトでは、バッファサイズは1,000です。
>>> from datasets import load_dataset
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True)
>>> shuffled_dataset = dataset.shuffle(seed=42, buffer_size=10_000)IterableDataset.shuffle()は、一度に最大
max_buffer_input_shards個のシャード(デフォルトは10)を使用してバッファを埋めるため、シャッフルの品質が大幅に向上します。また、データセットが複数のファイルに分割されている場合、シャードの順序もシャッフルします。
Parquet データセットをストリーミングする場合、ファイルが1つまたは少数の場合は、シャッフルの前にreshard()を使用してください。これにより、ファイル単位ではなく行グループ単位でデータセットをシャーディングでき、シャッフルの品質が向上します。
シャッフル
各エポック後にデータセットをシャッフルし直したい場合があります。そのためには、エポックごとに異なるシードを設定する必要があります。エポックの間にIterableDataset.set_epoch()を呼び出して、データセットに現在のエポックを通知してください。
シードは実質的にinitial seed + current epochになります。
>>> for epoch in range(epochs):
... shuffled_dataset.set_epoch(epoch)
... for example in shuffled_dataset:
... ...Split dataset
データセットを分割するには、以下の2つの方法があります。
- IterableDataset.take() は、データセットの最初の
n個の例を返します:
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True)
>>> dataset_head = dataset.take(2)
>>> list(dataset_head)
[{'text': "How AP reported in all formats from tor...},
{'text': 'Did you know you have two little yellow...}]- IterableDataset.skip() は、データセットの最初の
n個の例を省略し、残りの例を返します。
>>> train_dataset = shuffled_dataset.skip(1000)
takeとskipは、シャードの順序を固定するため、以降のshuffle呼び出しを無効にします。データセットを分割する前にshuffleしてください。
シャード
🤗 Datasets はシャーディングをサポートしており、非常に大きなデータセットを事前に定義された数のチャンクに分割できます。shard() で num_shards パラメータを指定して、データセットを分割するシャードの数を決定します。また、返すシャードを index パラメータで指定する必要があります。
例えば、amazon_polarityデータセットには4つのシャード(この場合は4つのParquetファイル)があります:
>>> from datasets import load_dataset
>>> dataset = load_dataset("fancyzhx/amazon_polarity", split="train", streaming=True)
>>> print(dataset)
IterableDataset({
features: ['label', 'title', 'content'],
num_shards: 4
})データセットを2つのチャンクにシャーディングした後、最初のチャンクには2つのシャードのみが含まれます:
>>> dataset.shard(num_shards=2, index=0)
IterableDataset({
features: ['label', 'title', 'content'],
num_shards: 2
})IterableDataset.reshard() を使用して、データセットのシャード数を増やすことができます。
>>> dataset.reshard()
IterableDataset({
features: ['label', 'title', 'content'],
num_shards: 3600
})リシャーディングの仕組みは、データセットのファイル形式に依存します。 例えば、Parquet の場合は、シャードごとに1つのファイルを持つ代わりに、行グループを使用してリシャーディングを行います。 各形式での動作については、IterableDataset.reshard()のドキュメントを参照してください。
リシャード後もdataset.num_shards==1のままの場合、代わりにIterableDataset.skip()とIterableDataset.take()を使用してチャンク分割を行う必要があります。
連結
同じカラムタイプを共有する場合は、別々のデータセットを連結できます。concatenate_datasets()を使用してデータセットを連結します:
>>> from datasets import concatenate_datasets, load_dataset
>>> stories = load_dataset("ajibawa-2023/General-Stories-Collection", split="train", streaming=True)
>>> stories = stories.select_columns(["text"]) # only keep the 'text' column
IterableDataset({
features: Unknown,
num_shards: 10
})
>>> wiki = load_dataset("wikimedia/wikipedia", "20231101.en", split="train", streaming=True)
>>> wiki = wiki.select_columns(["text"]) # only keep the 'text' column
IterableDataset({
features: ['text'],
num_shards: 41
})
>>> bert_dataset = concatenate_datasets([stories, wiki])
IterableDataset({
features: ['text'],
num_shards: 51
})連結されたデータセットのシャードは、入力データセットのシャードです。
axis=1を設定することで、行数が同じであれば、2つのデータセットを水平方向に連結することもできます:
>>> from datasets import IterableDataset
>>> stories_ids = IterableDataset.from_dict({"ids": list(range(num_stories))})
>>> stories_with_ids = concatenate_datasets([stories, stories_ids], axis=1)この場合、入力データセットから整列していないシャードで終わることを避けるため、連結されたデータセットには1つのシャードのみが含まれます。
インターリーブ
interleave_datasets() は、カラムの型が同じであれば、IterableDataset を他のデータセットと組み合わせることができます。組み合わせられたデータセットは、元の各データセットから交互にサンプルを返します。
>>> from datasets import interleave_datasets
>>> es_dataset = load_dataset('allenai/c4', 'es', split='train', streaming=True)
IterableDataset({
features: Unknown,
num_shards: 2048
})
>>> fr_dataset = load_dataset('allenai/c4', 'fr', split='train', streaming=True)
IterableDataset({
features: Unknown,
num_shards: 2048
})
>>> multilingual_dataset = interleave_datasets([es_dataset, fr_dataset])
IterableDataset({
features: ['text', 'timestamp', 'url'],
num_shards: 2048
})
>>> list(multilingual_dataset.take(2))
[{'text': 'Comprar Zapatillas para niña en chancla con goma por...'},
{'text': 'Le sacre de philippe ier, 23 mai 1059 - Compte Rendu...'}]各元のデータセットからのサンプリング確率を定義し、各データセットのサンプリングと組み合わせ方をより細かく制御できます。probabilities引数に、希望するサンプリング確率を設定してください。
>>> multilingual_dataset_with_oversampling = interleave_datasets([es_dataset, fr_dataset], probabilities=[0.8, 0.2], seed=42)
>>> list(multilingual_dataset_with_oversampling.take(2))
[{'text': 'Comprar Zapatillas para niña en chancla con goma por...'},
{'text': 'Chevrolet Cavalier Usados en Bogota - Carros en Vent...'}]最終的なデータセットの約80%はes_datasetで構成され、20%はfr_datasetで構成されています。
stopping_strategy を指定することもできます。デフォルトの戦略である first_exhausted はサブサンプリング戦略であり、いずれかのデータセットのサンプルが尽きた時点でデータセットの構築が停止します。
stopping_strategy=all_exhausted を指定すると、オーバーサンプリング戦略が実行されます。この場合、すべてのデータセットのすべてのサンプルが少なくとも1回追加された時点で、データセットの構築が停止します。実際には、あるデータセットが尽きた場合、停止条件が満たされるまで、そのデータセットの先頭に戻って処理が続行されます。
サンプリング確率が指定されていない場合、新しいデータセットには max_length_datasets*nb_dataset samples が含まれることに注意してください。
また、すべてのサンプルがちょうど1回だけ参照されることを保証するために、stopping_strategy=all_exhausted_without_replacement も利用できます。
シャーディングを用いた並列処理を適切に行うために、インターリーブされたデータセットのシャードには、すべての入力データセットから少なくとも1つのシャードが含まれます。したがって、インターリーブされたデータセットのシャーディングレベルは、入力データセットのシャーディングレベルの最小値となります。
例えば、入力データセットのシャード数がそれぞれ32、48、128の場合、インターリーブされたデータセットのシャード数は32 = min(32, 48, 128)となり、各新しいシャードには、最初のデータセットから1つのシャード、2番目のデータセットから1〜2つのシャード、3番目のデータセットから4つのシャードが含まれます。
リネーム、削除、キャスト
以下のメソッドを使用すると、データセットの列を変更できます。これらのメソッドは、列の名前を変更したり、列を削除したり、列を新しい特徴量のセットに変更したりする際に便利です。
名前の変更
データセットの列名を変更する必要がある場合は、IterableDataset.rename_column()を使用します。元の列に関連付けられた特徴量は、元の列をその場で置き換えるのではなく、実際には新しい列名の下に移動されます。
IterableDataset.rename_column() に元の列名と新しい列名を指定します:
>>> from datasets import load_dataset
>>> dataset = load_dataset('allenai/c4', 'en', streaming=True, split='train')
>>> dataset = dataset.rename_column("text", "content")削除
1つ以上の列を削除する必要がある場合は、IterableDataset.remove_columns()に削除する列の名前を指定します。複数の列を削除するには、列名のリストを指定します:
>>> from datasets import load_dataset
>>> dataset = load_dataset('allenai/c4', 'en', streaming=True, split='train')
>>> dataset = dataset.remove_columns('timestamp')キャスト
IterableDataset.cast() は、1 つ以上のカラムのフィーチャー型を変更します。このメソッドは、新しい Features を引数として受け取ります。以下のサンプルコードは、ClassLabel と Value のフィーチャー型を変更する方法を示しています。
>>> from datasets import load_dataset
>>> dataset = load_dataset('nyu-mll/glue', 'mrpc', split='train', streaming=True)
>>> dataset.features
{'sentence1': Value('string'),
'sentence2': Value('string'),
'label': ClassLabel(names=['not_equivalent', 'equivalent']),
'idx': Value('int32')}
>>> from datasets import ClassLabel, Value
>>> new_features = dataset.features.copy()
>>> new_features["label"] = ClassLabel(names=['negative', 'positive'])
>>> new_features["idx"] = Value('int64')
>>> dataset = dataset.cast(new_features)
>>> dataset.features
{'sentence1': Value('string'),
'sentence2': Value('string'),
'label': ClassLabel(names=['negative', 'positive']),
'idx': Value('int64')}キャストは、元のフィーチャー型と新しいフィーチャー型が互換性がある場合にのみ機能します。例えば、元の列が1と0のみを含む場合、フィーチャー型
Value('int32')の列をValue('bool')にキャストできます。
IterableDataset.cast_column() を使用すると、1つのカラムのみのフィーチャー型を変更できます。カラム名と新しいフィーチャー型を引数として渡します:
>>> dataset.features
{'audio': Audio(sampling_rate=44100, mono=True)}
>>> dataset = dataset.cast_column("audio", Audio(sampling_rate=16000))
>>> dataset.features
{'audio': Audio(sampling_rate=16000, mono=True)}マップ
通常のDataset向けのDataset.map()関数と同様に、🤗 Datasetsには、IterableDatasetを処理するためのIterableDataset.map()が用意されています。IterableDataset.map()は、サンプルがストリーミングされる際に、処理をオンザフライで適用します。
mapメソッドは、データセットの各サンプルに対して処理関数を適用できます。サンプルごとに独立して適用することも、バッチ単位で適用することも可能です。この関数は、新しい行や列を作成することさえできます。
以下の例は、IterableDatasetをトークナイズする方法を示しています。この関数は、dictを受け取り、dictを出力する必要があります:
>>> def add_prefix(example):
... example['text'] = 'My text: ' + example['text']
... return example次に、IterableDataset.map() を使用して、この関数をデータセットに適用します:
>>> from datasets import load_dataset
>>> dataset = load_dataset('allenai/c4', 'en', streaming=True, split='train')
>>> updated_dataset = dataset.map(add_prefix)
>>> list(updated_dataset.take(3))
[{'text': 'My text: Beginners BBQ Class Taking Place in Missoula!\nDo you want to get better at making...',
'timestamp': '2019-04-25 12:57:54',
'url': 'https://klyq.com/beginners-bbq-class-taking-place-in-missoula/'},
{'text': 'My text: Discussion in \'Mac OS X Lion (10.7)\' started by axboi87, Jan 20, 2012.\nI\'ve go...',
'timestamp': '2019-04-21 10:07:13',
'url': 'https://forums.macrumors.com/threads/restore-from-larger-disk-to-smaller-disk.1311329/'},
{'text': 'My text: Foil plaid lycra and spandex shortall with metallic slinky insets. Attached metall...',
'timestamp': '2019-04-25 10:40:23',
'url': 'https://awishcometrue.com/Catalogs/Clearance/Tweens/V1960-Find-A-Way'}]別の例を見てみましょう。今回は、IterableDataset.map()を使用してカラムを削除します。カラムを削除する場合、そのカラムの内容がマッピング関数に渡された後に削除されます。これにより、マッピング関数はカラムが削除される前にその内容を利用できます。
IterableDataset.map()のremove_columns引数で、削除するカラムを指定します:
>>> updated_dataset = dataset.map(add_prefix, remove_columns=["timestamp", "url"])
>>> list(updated_dataset.take(3))
[{'text': 'My text: Beginners BBQ Class Taking Place in Missoula!\nDo you want to get better at making...'},
{'text': 'My text: Discussion in \'Mac OS X Lion (10.7)\' started by axboi87, Jan 20, 2012.\nI\'ve go...'},
{'text': 'My text: Foil plaid lycra and spandex shortall with metallic slinky insets. Attached metall...'}]バッチ処理
IterableDataset.map() は、バッチ単位での処理にも対応しています。batched=True を設定することで、バッチに対して操作を実行できます。デフォルトのバッチサイズは 1000 ですが、batch_size 引数で調整できます。これにより、トークナイズ、長い文を短いチャンクに分割する処理、データ拡張など、多くの興味深いアプリケーションが可能になります。
トークナイズ
>>> from datasets import load_dataset
>>> from transformers import AutoTokenizer
>>> dataset = load_dataset("allenai/c4", "en", streaming=True, split="train")
>>> tokenizer = AutoTokenizer.from_pretrained('distilbert-base-uncased')
>>> def encode(examples):
... return tokenizer(examples['text'], truncation=True, padding='max_length')
>>> dataset = dataset.map(encode, batched=True, remove_columns=["text", "timestamp", "url"])
>>> next(iter(dataset))
{'input_ids': [101, 4088, 16912, 22861, 4160, 2465, 2635, 2173, 1999, 3335, ..., 0, 0, 0],
'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ..., 0, 0]}バッチ化された map 処理のドキュメントで、バッチ処理の他の例を確認できます。イテラブルなデータセットでも同じように動作します。
フィルター
Dataset.filter() を使用して、述語関数に基づいてデータセットの行をフィルタリングできます。指定された条件に一致する行を返します。
>>> from datasets import load_dataset
>>> dataset = load_dataset('HuggingFaceFW/fineweb', streaming=True, split='train')
>>> start_with_ar = dataset.filter(lambda example: example['text'].startswith('San Francisco'))
>>> next(iter(start_with_ar))
{'text': 'San Francisco 49ers cornerback Shawntae Spencer will miss the rest of the sea...}with_indices=True を設定すると、Dataset.filter() はインデックスによるフィルタリングも実行できます:
>>> even_dataset = dataset.filter(lambda example, idx: idx % 2 == 0, with_indices=True)
>>> list(even_dataset.take(3))
[{'text': 'How AP reported in all formats from tornado-stricken regionsMarch 8, 2012 Whe...},
{'text': 'Car Wash For Clara! Now is your chance to help! 2 year old Clara Woodward has...},
{'text': 'Log In Please enter your ECode to log in. Forgotten your eCode? If you create...}]バッチ
batch メソッドは、IterableDataset をバッチのイテラブルに変換します。これは、トレーニングループ内でバッチを扱う場合や、バッチ化された入力を想定するフレームワークを使用する場合に特に役立ちます。
map関数を使用してデータのバッチに関数を適用する際にも「Batch Processing」オプションがありますが、これは上記のMap セクションで説明されています。ここで説明するbatchメソッドはこれとは異なり、データセットからバッチを作成するためのより直接的な方法を提供します。
batchメソッドは次のように使用できます。
from datasets import load_dataset
# Load a dataset in streaming mode
dataset = load_dataset("some_dataset", split="train", streaming=True)
# Create batches of 32 samples
batched_dataset = dataset.batch(batch_size=32)
# Iterate over the batched dataset
for batch in batched_dataset:
print(batch)
breakこの例では、batched_dataset は依然として IterableDataset ですが、イテレーションで得られる各アイテムは、単一のサンプルではなく、32サンプルのバッチになります。 このバッチングは、データセットをイテレーションする際にオンザフライで実行されるため、IterableDataset のメモリ効率の高い特性が保たれます。
batchメソッドは、drop_last_batchパラメータも提供しています。 これをTrueに設定すると、指定されたbatch_sizeより小さい最後のバッチが破棄されます。 これは、下流の処理ですべてのバッチが同じサイズである必要がある場合に便利です:
batched_dataset = dataset.batch(batch_size=32, drop_last_batch=True)トレーニングループ内でストリーミング
IterableDatasetは、トレーニングループに統合できます。まず、データセットをシャッフルします:
>>> seed, buffer_size = 42, 10_000
>>> dataset = dataset.shuffle(seed, buffer_size=buffer_size)最後に、シンプルなトレーニングループを作成して、トレーニングを開始します:
>>> import torch
>>> from torch.utils.data import DataLoader
>>> from transformers import AutoModelForMaskedLM, DataCollatorForLanguageModeling
>>> from tqdm import tqdm
>>> dataset = dataset.with_format("torch")
>>> dataloader = DataLoader(dataset, collate_fn=DataCollatorForLanguageModeling(tokenizer))
>>> device = 'cuda' if torch.cuda.is_available() else 'cpu'
>>> model = AutoModelForMaskedLM.from_pretrained("distilbert-base-uncased")
>>> model.train().to(device)
>>> optimizer = torch.optim.AdamW(params=model.parameters(), lr=1e-5)
>>> for epoch in range(3):
... dataset.set_epoch(epoch)
... for i, batch in enumerate(tqdm(dataloader, total=5)):
... if i == 5:
... break
... batch = {k: v.to(device) for k, v in batch.items()}
... outputs = model(**batch)
... loss = outputs[0]
... loss.backward()
... optimizer.step()
... optimizer.zero_grad()
... if i % 10 == 0:
... print(f"loss: {loss}")データセットのチェックポイントを保存して反復を再開する
トレーニングループが停止した場合、中断した箇所からトレーニングを再開したいことがあるでしょう。そのためには、モデルとオプティマイザー、およびデータローダーのチェックポイントを保存できます。
Iterable データセットでは、特定の例のインデックスへのランダムアクセスを提供しないため、チェックポイントからの再開には、モデルやオプティマイザーと同様に、IterableDataset.state_dict() と IterableDataset.load_state_dict() を使用できます。
>>> iterable_dataset = Dataset.from_dict({"a": range(6)}).to_iterable_dataset(num_shards=3)
>>> for idx, example in enumerate(iterable_dataset):
... print(example)
... if idx == 2:
... state_dict = iterable_dataset.state_dict()
... print("checkpoint")
... break
>>> iterable_dataset.load_state_dict(state_dict)
>>> print(f"restart from checkpoint")
>>> for example in iterable_dataset:
... print(example)返り値:
{'a': 0}
{'a': 1}
{'a': 2}
checkpoint
restart from checkpoint
{'a': 3}
{'a': 4}
{'a': 5}内部では、iterable datasetは現在読み込まれているshardと、そのshard内のexampleのインデックスを追跡し、この情報をstate_dictに保存します。
チェックポイントから再開する場合、データセットは以前に読み取られたすべてのシャードをスキップして、現在のシャードから再開します。 その後、シャードを読み取り、チェックポイントの正確なサンプルに到達するまでサンプルをスキップします。
したがって、すでに反復処理済みのシャードを再読み込みしないため、データセットの再起動は比較的速くなります。ただし、データセットの再開は通常、即座には完了しません。現在のシャードの先頭から読み込みを再開し、チェックポイントの位置に到達するまでサンプルをスキップする必要があるためです。
これはtorchdataのStatefulDataLoaderと組み合わせて使用できます:
>>> from torchdata.stateful_dataloader import StatefulDataLoader
>>> iterable_dataset = load_dataset("deepmind/code_contests", streaming=True, split="train")
>>> dataloader = StatefulDataLoader(iterable_dataset, batch_size=32, num_workers=4)
>>> # checkpoint
>>> state_dict = dataloader.state_dict() # uses iterable_dataset.state_dict() under the hood
>>> # resume from checkpoint
>>> dataloader.load_state_dict(state_dict) # uses iterable_dataset.load_state_dict() under the hood
.shuffle()が使用されている場合を除き、再開時はチェックポイントが保存された位置から正確に再開されます。ただし、再開時にシャッフルバッファ内のサンプルは失われ、バッファは新しいデータで再充填されます。
保存
イテラブルなデータセットの準備が整ったら、Parquet形式のHugging Face Datasetとして保存し、後でload_dataset()を使って再利用できます。
Hugging Face 上のデータセットリポジトリの名前を引数として push_to_hub() に渡すことで、データセットを保存できます。これにより、データセットを反復処理しながら、データを Hugging Face に段階的にアップロードします:
dataset.push_to_hub("username/my_dataset")デフォルトでは、IterableDataset.push_to_hub() は dataset.num_shards 個のファイルを出力します。出力ファイル数を指定するには、num_shards を設定してください。出力シャード数が dataset.num_shards より多い場合、各入力シャードは、データを1パスで処理しながら、その中の例をラウンドロビン方式で割り当てることで、専用の出力シャードのグループに分散されます。例の数が要求されたシャード数より少ない場合、残りのシャードは空の Parquet ファイルになります。
データセットが複数のシャードで構成されている場合(dataset.num_shards > 1)、複数のプロセスを使用して並列にアップロードできます。map()またはfilter()ステップを適用した場合は、並列実行により高速化されるため、特に有用です。
dataset.push_to_hub("username/my_dataset", num_proc=8)load_dataset()関数を使用して、データセットを再読み込みします:
from datasets import load_dataset
reloaded_dataset = load_dataset("username/my_dataset")エクスポート
🤗 Datasetsはエクスポートにも対応しているため、他のアプリケーションでデータセットを扱うことができます。以下の表は、現在エクスポートに対応しているファイル形式を示しています。
| ファイル形式 | エクスポート方法 |
|---|---|
| CSV | IterableDataset.to_csv() |
| JSON | IterableDataset.to_json() |
| Parquet | IterableDataset.to_parquet() |
| SQL | IterableDataset.to_sql() |
| メモリ内のPythonオブジェクト | IterableDataset.to_pandas(), IterableDataset.to_polars() or IterableDataset.to_dict() |
例えば、次のようにデータセットをCSVファイルにエクスポートします。
>>> dataset.to_csv("path/of/my/dataset.csv")データセットが大きい場合、シャードごとに1つのファイルを保存できます。例:
>>> num_shards = dataset.num_shards
>>> for index in range(num_shards):
... shard = dataset.shard(index, num_shards)
... shard.to_parquet(f"path/of/my/dataset/data-{index:05d}.parquet")