【huggingface/datasets】複数のデータセットを組み合わせてサンプリングする

本記事はhuggingfaceのdatasetsを利用してサンプリングの割合を考慮しながら複数のデータセットを混ぜて扱う方法を紹介します。紹介する方法を知っておくと、例えば言語モデルの事前学習など、複数のテキストコーパスを組み合わせて利用する際に便利です。

環境

Python==3.10.6
datasets==2.9.0

想定する利用シーン

想定する利用シーンはサイズがそれほど大きくないWikipediaと、サイズが大きなCommon Crawlから取得したウェブテキストを合わせて言語モデルの事前学習に利用したいのだけど、品質が比較的よいWikipediaのサンプルをより学習に利用したいのでウェブテキストよりもサンプリングの割合を高くする、といったものが考えられます。

サンプルデータ

2つのデータセットのサンプルを以下のように用意します。各データセットはどのようなkeyを持っていてもいいのですが、同じkeyを持つようにします。また、ファイル形式もなんでもよいのですが、ここではjsonl形式とします。

data_source1.json

{"id": "source1_1", "text": "データソース1のテキスト1"}
{"id": "source1_2", "text": "データソース1のテキスト2"}
{"id": "source1_3", "text": "データソース1のテキスト3"}

data_source2.json

{"id": "source2_1", "text": "データソース2のテキスト1"}
{"id": "source2_2", "text": "データソース2のテキスト2"}
{"id": "source2_3", "text": "データソース2のテキスト3"}
{"id": "source2_4", "text": "データソース2のテキスト4"}
{"id": "source2_5", "text": "データソース2のテキスト5"}
{"id": "source2_6", "text": "データソース2のテキスト6"}
{"id": "source2_7", "text": "データソース2のテキスト7"}
{"id": "source2_8", "text": "データソース2のテキスト8"}
{"id": "source2_9", "text": "データソース2のテキスト9"}

プログラム例

huggingfaceが開発しているdatasetsを利用すればとても簡単に複数のデータをまとめて、しかも指定した割合でサンプリングすることができます。大事なのはinterleave_datasetsです。

probabilitiesで各データセットからサンプリングする割合を指定
stopping_strategyを"all_exhausted"にします。これを指定しないと、一番サンプル数が少ないデータをすべて読み込んだら、データの結合が終了します。そうなると、サイズが大きなデータのほうは利用されないサンプルが出る可能性があるので気を付けてください。

利用例は以下の通りです。

from datasets import interleave_datasets, load_dataset
from datasets.arrow_dataset import Dataset
from datasets.dataset_dict import DatasetDict

data_source1: DatasetDict = load_dataset(
    "json", data_files={"train": "data_source1.json"},
)
data_source2: DatasetDict = load_dataset(
    "json", data_files={"train": "data_source2.json"},
)

combined_data_source: Dataset = interleave_datasets(
    [data_source1["train"], data_source2["train"]],
    probabilities=[0.7, 0.3],
    stopping_strategy="all_exhausted",
    seed=42,
)

print(len(combined_data_source))

for sample in combined_data_source:
    print(sample)

出力結果は次のようになります。

21
{'id': 'source2_1', 'text': 'データソース2のテキスト1'}
{'id': 'source1_1', 'text': 'データソース1のテキスト1'}
{'id': 'source2_2', 'text': 'データソース2のテキスト2'}
{'id': 'source1_2', 'text': 'データソース1のテキスト2'}
{'id': 'source1_3', 'text': 'データソース1のテキスト3'}
{'id': 'source2_3', 'text': 'データソース2のテキスト3'}
{'id': 'source2_4', 'text': 'データソース2のテキスト4'}
{'id': 'source2_5', 'text': 'データソース2のテキスト5'}
{'id': 'source1_1', 'text': 'データソース1のテキスト1'}
{'id': 'source1_2', 'text': 'データソース1のテキスト2'}
{'id': 'source1_3', 'text': 'データソース1のテキスト3'}
{'id': 'source2_6', 'text': 'データソース2のテキスト6'}
{'id': 'source1_1', 'text': 'データソース1のテキスト1'}
{'id': 'source2_7', 'text': 'データソース2のテキスト7'}
{'id': 'source1_2', 'text': 'データソース1のテキスト2'}
{'id': 'source1_3', 'text': 'データソース1のテキスト3'}
{'id': 'source1_1', 'text': 'データソース1のテキスト1'}
{'id': 'source1_2', 'text': 'データソース1のテキスト2'}
{'id': 'source2_8', 'text': 'データソース2のテキスト8'}
{'id': 'source1_3', 'text': 'データソース1のテキスト3'}
{'id': 'source2_9', 'text': 'データソース2のテキスト9'}

サンプル数が少ないdata_source1.jsonですが、probabilitiesで0.7としているので比較的多くサンプリングされています (21サンプル中12件)。 stopping_strategy="all_exahsted"としているので、サンプル数が多いdata_source2.jsonのサンプルはすべてサンプリングされます (21サンプル中9件)。