Zenn
🤗

HuggingFace datasetsのBuilder classを使ってデータセットを自作する

2024/09/22に公開
Python
transformer
Hugging Face
tech

こんにちは。今回はHuggingFace datasetsから
呼び出せる形式でデータセットを自作する方法を紹介します。この記事では特にBuilder classを用いた方法について書いていきます。
やっていることは以下と同様です。

https://huggingface.co/docs/datasets/dataset_script

https://github.com/huggingface/datasets/tree/main/templates

前提: HuggingFaceでデータセットを作成する方法

HuggingFaceのデータセットを自作する方法は、大きく分けて以下の3つあります。

  1. あらかじめ特定の構造のファイル or ディレクトリを用意してdatasets.load_datasetを使う方法
  2. dictやgeneratorを定義してdatasets.Dataset.from_*関数を利用する方法
  3. データセットのロードをコードで定義するdatasets.DatasetBuilderクラスを利用する方法(今回紹介する方法)

1. あらかじめ特定の構造のファイル or ディレクトリを用意してdatasets.load_datasetを使う方法

まず、ファイルを読むこむ方法は以下のリンクのように、あらかじめ用意したcsv, jsonなどの形式のファイルをそのままload_dataset関数で読み込みます。

https://huggingface.co/docs/datasets/loading#local-and-remote-files

例えばcsv形式のファイルなら以下のように読み込みます。

import datasets as ds

dataset = ds.load_dataset(
    "csv",
    data_files={"train": ["my_train_file_1.csv", "my_train_file_2.csv"], "test": "my_test_file.csv"}
)

あらかじめファイルを用意しておくだけなので簡単です。

ディレクトリを読み込む方法は、以下のリンクのように特定のディレクトリ以下にtrain, validationなどのsplitディレクトリを作成して、そのディレクトリをload_dataset関数で読み込みます。

https://huggingface.co/docs/datasets/create_dataset#folder-based-builders

https://huggingface.co/docs/datasets/image_dataset#imagefolder

https://huggingface.co/docs/datasets/audio_dataset#audiofolder

hoge/dataset_root/train/label/1.png
hoge/dataset_root/train/label/2.png
...
hoge/dataset_root/test/label/1.png
hoge/dataset_root/test/label/2.png
...

import datasets as ds

dataset = ds.load_dataset("imagefolder", data_dir="hoge/dataset_root")

この方法は画像以外にもaudiofolderとすることで音声ファイルにも利用することができます。

2. datasets.Dataset.from_*関数を利用する方法

この方法では、以下のチュートリアルにあるように、Pythonの関数やdictを渡すことでデータセットを作成することができます。

https://huggingface.co/docs/datasets/create_dataset#from-python-dictionaries

例えば、dictを渡す場合は以下のように作成します。

import datasets as ds

dataset_dict = {
    "text": ["hoge", "fuga"]
    "label": [1, 3]
}

dataset = ds.Dataset.from_dict(dataset_dict)

generatorを渡す場合は以下のように作成します。引数のgen_kwargsにgenerator関数に渡す引数を、num_procに並列数を渡すことができます。

import os
import datasets as ds

dataset_data = [
    {"summary": "hoge", "class": 1},
    {"summary": "fuga", "class": 2},
]

def generator(data):
    for d in data:
        yield {
            "text": d["summary"],
            "label": d["class"]
        }

dataset = ds.Dataset.from_generator(generator,
                                    gen_kwargs={"data": dataset_data},
                                    num_proc=os.cpu_count() // 2)

個人的に最も手軽なのは、generatorとdatasets.Dataset.from_generatorを使う方法かと思います。スクリプト上で元データの整形をしながら並列にデータセットを作成可能なので手軽でおすすめです。

3. datasets.DatasetBuilderクラスを利用する方法

今回紹介する方法です。後に詳しく書きます。

本題: datasets.DatasetBuilderクラスを利用したデータセットの作成の概要

本記事の本題である、datasets.DatasetBuilderなどのBuilder classを利用したデータセット作成方法を紹介します。今回はBuilder classの中でも、datasets.GeneratorBasedBuilderを使います。まずはBuilder classを利用するとどんな利点があるのかをまとめてから、実際に作成するコードを紹介します。

メリット

Builder classを使うと以下のような利点があります。

  • データセットのファイル自体を入手し、前準備をする方法を定義できる
  • HuggingFace datasetsでサポートされていない形式や複雑な準備が必要なデータセットにも対応できる
  • 複数のデータセットを束ねることができる

前述した1,2のload_datasetやfrom_generator関数を利用した方法では、あらかじめダウンロードやディレクトリの準備が必要でしたが、Builder classを使うと準備の操作もコードとして定義することが可能なため、Pythonスクリプトのみを共有するだけでHuggingFace datasetsを共有することが可能です。

HuggingFaceで既に公開されているJMTEBやJGLUEといった複数のデータセットを用いるモデルの評価指標リポジトリはBuilder classで定義されています。言語モデルの事前学習に用いる多言語データセットのwikimedia/wikipediaは以下のように呼び出すことで、言語の切り替えを行うことが可能ですが、Builder classを用いると、このような複雑な分割についても扱うことができます。

from datasets import load_dataset

ds = load_dataset("wikimedia/wikipedia", name="20231101.en")

関連class, object

データセットを定義するに当たって、以下のいくつかの関連クラスやオブジェクトを利用します。

  • datasets.BuilderConfig
    Builder classでデータセットを呼び出すのに必要なデータセットの名前、バージョンやdata_dirなどを定義します。複数のデータセットを束ねる場合は複数定義することになります。
  • datasets.DatasetInfo
    データセットの情報を定義します。citationやhomepage、データセットの要素の型などを定義します。
  • datasets.SplitGenerator
    データセットのsplitの定義をします。例えばデータセットにtrain, val, testといったsplitが含まれる場合はその数だけ定義することになります。
  • datasets.GeneratorBasedBuilder
    上記3つのclassを元に、Datasetを作成するClassを定義します。

実装

それでは実際にdatasets.GeneratorBasedBuilderを使って、日本語コーパスのlivedoor ニュースコーパスとllm-book/japanese-wikipediaを束ねたJapaneseTextDatasetを作成してみます。
livedoor ニュースコーパスはWebから入手し、llm-book/japanese-wikipediaはHuggingFaceにアップロードされているものをそのまま利用します。

ディレクトリ構成は以下とします。

JapaneseTextDataset/
└── JapaneseTextDataset.py

以下がJapaneseTextDataset.pyの中身です。順に説明していきます。

from dataclasses import dataclass
from pathlib import Path
from typing import Literal, Optional

import datasets as ds


@dataclass
class LivedoorCorpusConfig(ds.BuilderConfig):
    """BuilderConfig for LivedoorCorpus"""

    name: str = "livedoor"


@dataclass
class WikiJPConfig(ds.BuilderConfig):
    """BuilderConfig for WikiJP"""
    name: str = "wiki"


class JapaneseTextDataset(ds.GeneratorBasedBuilder):
    """Japanese Text Dataset"""

    VERSION = "1.0.0"
    DEFAULT_CONFIG_NAME = "wiki"
    BUILDER_CONFIGS = [
        WikiJPConfig(
            version=ds.Version(version_str=VERSION), description="Japanese Wikipedia"
        ),
        LivedoorCorpusConfig(
            version=ds.Version(version_str=VERSION), description="Livedoor News Corpus"
        ),
    ]

    def _info(self) -> ds.DatasetInfo:
        if self.config.name == "wiki":
            return ds.DatasetInfo(
                homepage="https://huggingface.co/datasets/llm-book/japanese-wikipedia",
                features=ds.Features(
                    {
                        "text": ds.Value("string"),
                    }
                ),
            )
        elif self.config.name == "livedoor":
            return ds.DatasetInfo(
                homepage="http://www.rondhuit.com/download.html#ldcc",
                features=ds.Features(
                    {
                        "url": ds.Value("string"),
                        "title": ds.Value("string"),
                        "text": ds.Value("string"),
                    }
                ),
            )

    def _split_generators(self, dl_manager: ds.DownloadManager):
        if self.config.name == "wiki":
            dataset = ds.load_dataset(
                "llm-book/japanese-wikipedia",
                trust_remote_code=True,
            )
            return [
                ds.SplitGenerator(
                    name=ds.Split.TRAIN,
                    gen_kwargs={"data": dataset["train"]},
                )
            ]
        elif self.config.name == "livedoor":
            file_path = dl_manager.download_and_extract(
                "http://www.rondhuit.com/download/ldcc-20140209.tar.gz"
            )
            return [
                ds.SplitGenerator(
                    name=ds.Split.TRAIN,
                    gen_kwargs={"file_path": file_path, "split": "train"},
                ),
                ds.SplitGenerator(
                    name=ds.Split.TEST,
                    gen_kwargs={"file_path": file_path, "split": "test"},
                ),
            ]

    def _generate_examples(
        self,
        data: Optional[ds.Dataset] = None,
        file_path: Optional[str] = None,
        split: Literal["train", "test"] | None = None,
    ):
        if self.config.name == "wiki":
            if data is None:
                raise ValueError("data must be specified")
            for i, example in enumerate(data):
                yield i, {"text": example["text"]}
        elif self.config.name == "livedoor":
            if file_path is None:
                raise ValueError("file_path must be specified")
            if split is None:
                raise ValueError("split must be specified")

            paths = [
                p
                for p in (Path(file_path) / "text").glob("*/*.txt")
                if p.name != "LICENSE.txt"
            ]
            if split == "train":
                paths = paths[: int(len(paths) * 0.8)]
            else:
                paths = paths[int(len(paths) * 0.8) :]
            for i, path in enumerate(paths):
                with open(path, "r", encoding="utf-8") as f:
                    url = f.readline().strip()
                    title = f.readline().strip()
                    text = f.read()
                yield i, {"url": url, "title": title, "text": text}

datasets.BuilderConfig

上記のコードの以下の部分で、LivedoorとWikipediaのBuilderConfigを定義しています。

@dataclass
class LivedoorCorpusConfig(ds.BuilderConfig):
    """BuilderConfig for LivedoorCorpus"""

    name: str = "livedoor"


@dataclass
class WikiJPConfig(ds.BuilderConfig):
    """BuilderConfig for WikiJP"""
    name: str = "wiki"

今回は特別なことはしていませんが、load_datasetを呼ぶときに引数を追加する場合は以下のようにします。

@dataclass
class WikiJPConfig(ds.BuilderConfig):
    """BuilderConfig for WikiJP"""
    name: str = "wiki"
    language: str = "ja"

上記の例ではlanguageという引数を足しています。こうすることで、データセットを読み込む時に以下のような方法でlanguage引数を渡すことができるようになります。

dataset = ds.load_dataset("JapaneseTextDataset", name="wiki", language="ja")

datasets.GeneratorBasedBuilder

datasets.GeneratorBasedBuilderを継承したJapaneseTextDatasetクラスにはVERSION, DEFAULT_CONFIG_NAME, BUILDER_CONFIGSを設定します。
複数データセットを呼べるようにする場合はBUILDER_CONFIGSに複数のdatasets.BuilderConfigをリストで列挙しておきます。
DEFAULT_CONFIG_NAMEではload_datasetで何も指定しなかった場合に、どのconfig_nameが呼ばれるかを設定します。デフォルト値は"default"です。

class JapaneseTextDataset(ds.GeneratorBasedBuilder):
    """Japanese Text Dataset"""

    VERSION = "1.0.0"
    DEFAULT_CONFIG_NAME = "wiki"
    BUILDER_CONFIGS = [
        WikiJPConfig(
            version=ds.Version(version_str=VERSION), description="Japanese Wikipedia"
        ),
        LivedoorCorpusConfig(
            version=ds.Version(version_str=VERSION), description="Livedoor News Corpus"
        ),
    ]

datasets.GeneratorBasedBuilder::_info

datasets.GeneratorBasedBuilderの_info関数では、datasets.DatasetInfo形式でデータセットの情報を返してあげます。このとき、self.configで、BUILDER_CONFIGSの内容を参照できるので、呼び出すデータセットの情報を動的に変更できます。

def _info(self) -> ds.DatasetInfo:
    if self.config.name == "wiki":
        return ds.DatasetInfo(
            homepage="https://huggingface.co/datasets/llm-book/japanese-wikipedia",
            features=ds.Features(
                {
                    "text": ds.Value("string"),
                }
            ),
        )
    elif self.config.name == "livedoor":
        return ds.DatasetInfo(
            homepage="http://www.rondhuit.com/download.html#ldcc",
            features=ds.Features(
                {
                    "url": ds.Value("string"),
                    "title": ds.Value("string"),
                    "text": ds.Value("string"),
                }
            ),
        )

datasets.GeneratorBasedBuilder::_split_generators

datasets.GeneratorBasedBuilderの_split_generators(self, dl_manager: ds.DownloadManager)関数では、データセットにどのようなsplitがあるのかをds.SplitGeneratorで定義してリストで返します。SplitGeneratorのgen_kwargsは後ほど出てくる、データ1行を返す_generate_examplesに渡す引数を定義します。

今回は、"llm-book/japanese-wikipedia"はHuggingFaceから取得するので内部でload_datasetを呼んでいます。
Livedoor News Corpusはウェブからダウンロードして利用するので、ダウンロードをスクリプトで書き、保存先をgen_kwargsに渡してあります。
dl_manager.download_and_extractを使うと、URLからダウンロードと回答を一気に行ってくれるのでこれを利用します。

def _split_generators(self, dl_manager: ds.DownloadManager):
    if self.config.name == "wiki":
        dataset = ds.load_dataset(
            "llm-book/japanese-wikipedia",
            trust_remote_code=True,
        )
        return [
            ds.SplitGenerator(
                name=ds.Split.TRAIN,
                gen_kwargs={"data": dataset["train"]},
            )
        ]
    elif self.config.name == "livedoor":
        file_path = dl_manager.download_and_extract(
            "http://www.rondhuit.com/download/ldcc-20140209.tar.gz"
        )
        return [
            ds.SplitGenerator(
                name=ds.Split.TRAIN,
                gen_kwargs={"file_path": file_path, "split": "train"},
            ),
            ds.SplitGenerator(
                name=ds.Split.TEST,
                gen_kwargs={"file_path": file_path, "split": "test"},
            ),
        ]

datasets.GeneratorBasedBuilder::_generate_examples

datasets.GeneratorBasedBuilderの_generate_examples関数では、データセットの1データをどのように返すかを定義します。引数には先ほどの_split_generatorsで定義したgen_kwargsが渡されます。generatorで定義するので、yield構文を使っています。

def _generate_examples(
        self,
        data: Optional[ds.Dataset] = None,
        file_path: Optional[str] = None,
        split: Literal["train", "test"] | None = None,
    ):
        if self.config.name == "wiki":
            if data is None:
                raise ValueError("data must be specified")
            for i, example in enumerate(data):
                yield i, {"text": example["text"]}
        elif self.config.name == "livedoor":
            if file_path is None:
                raise ValueError("file_path must be specified")
            if split is None:
                raise ValueError("split must be specified")

            paths = [
                p
                for p in (Path(file_path) / "text").glob("*/*.txt")
                if p.name != "LICENSE.txt"
            ]
            if split == "train":
                paths = paths[: int(len(paths) * 0.8)]
            else:
                paths = paths[int(len(paths) * 0.8) :]
            for i, path in enumerate(paths):
                with open(path, "r", encoding="utf-8") as f:
                    url = f.readline().strip()
                    title = f.readline().strip()
                    text = f.read()
                yield i, {"url": url, "title": title, "text": text}

データセットの呼び出し

以上でデータセットの構築は終了です。実際にデータセットを呼び出してみます。
以下のようなスクリプトで呼び出せます。

import datasets as ds

# Load the dataset
dataset = ds.load_dataset("./JapaneseTextDataset", "wiki", trust_remote_code=True)
print(dataset)

# Load the dataset
dataset = ds.load_dataset(
    "./JapaneseTextDataset", "livedoor", trust_remote_code=True
)
print(dataset)

実行結果は以下です。

ja_wiki.jsonl: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 6.37G/6.37G [00:53<00:00, 38.7MB/s]
Generating train split: 1363395 examples [00:38, 35532.80 examples/s]
Generating train split: 1363395 examples [03:22, 6731.14 examples/s]
DatasetDict({
    train: Dataset({
        features: ['text'],
        num_rows: 1363395
    })
})
Downloading data: 31.6MB [00:00, 50.9MB/s]
Generating train split: 5893 examples [00:01, 3227.36 examples/s]
Generating test split: 1474 examples [00:00, 2580.49 examples/s]
DatasetDict({
    train: Dataset({
        features: ['url', 'title', 'text'],
        num_rows: 5893
    })
    test: Dataset({
        features: ['url', 'title', 'text'],
        num_rows: 1474
    })
})

無事データセットが呼び出せました。

HuggingFace Hubへのアップロード

作成したデータセットをHuggingFaceにアップロードするには、先ほど作成したJapaneseTextDatasetディレクトリをそのままpushすれば良いです。

以下のTemplateにあるように、README.mdにDatasetの情報を書いておくとそのまま使われます。

https://github.com/huggingface/datasets/tree/main/templates

また、以下のリポジトリはJGLUEをHuggingFaceにアップロードしているGitHubのリポジトリなのですが、CI/CDまで組んであり、大変参考になります。

https://github.com/shunk031/huggingface-datasets_JGLUE/tree/main

以上、簡単にですがBuilder classを使ったHuggingFace datasetsの作成方法でした。

参考

https://github.com/shunk031/huggingface-datasets_JGLUE

https://qiita.com/oyahiroki/items/662ed10ab3f8d4bf8f61

https://github.com/huggingface/datasets/tree/main/templates

https://huggingface.co/docs/datasets/package_reference/builder_classes

https://huggingface.co/datasets/sbintuitions/JMTEB/blob/main/JMTEB.py

Discussion