创建数据集加载脚本

如果你的数据集采用以下格式之一，则可能不需要数据集加载脚本：CSV、JSON、JSON 行、文本、图像、音频或 Parquet。对于这些格式，只要你的数据集仓库具有必需的结构，你就可以使用 load_dataset() 自动加载数据集。

出于安全原因，🤗 Datasets 默认不允许运行数据集加载脚本，你必须传递 trust_remote_code=True才能加载需要运行数据集脚本的数据集。

编写数据集脚本以加载和共享包含不支持格式的数据文件或需要更复杂数据准备的数据集。这是一种比使用数据集卡片中的 YAML 元数据更高级的方式来定义数据集。数据集脚本是一个 Python 文件，它定义了数据集的不同配置和拆分，以及如何下载和处理数据。

该脚本可以从任何网站或从同一数据集仓库下载数据文件。

数据集加载脚本应与数据集仓库或目录具有相同的名称。例如，名为 my_dataset 的仓库应包含 my_dataset.py 脚本。这样就可以用以下方式加载它：

my_dataset/
├── README.md
└── my_dataset.py

>>> from datasets import load_dataset
>>> load_dataset("path/to/my_dataset")

以下指南包含有关数据集脚本的说明，介绍如何

添加数据集元数据。
下载数据文件。
生成样本。
生成数据集元数据。
将数据集上传到 Hub。

打开 SQuAD 数据集加载脚本模板，以了解如何共享数据集。

为了帮助您入门，请尝试从数据集加载脚本模板开始！

添加数据集属性

第一步是在 DatasetBuilder._info() 中添加有关数据集的一些信息或属性。您应该指定的最重要的属性是

DatasetInfo.description 提供了对数据集的简要描述。该描述向用户告知数据集的内容、收集方式以及如何将其用于 NLP 任务。
DatasetInfo.features 定义了数据集每列的名称和类型。这还将为每个示例提供结构，因此如果您希望，可以在列中创建嵌套的子字段。查看 Features 以获取您可以使用的所有功能类型的完整列表。

datasets.Features(
    {
        "id": datasets.Value("string"),
        "title": datasets.Value("string"),
        "context": datasets.Value("string"),
        "question": datasets.Value("string"),
        "answers": datasets.Sequence(
            {
                "text": datasets.Value("string"),
                "answer_start": datasets.Value("int32"),
            }
        ),
    }
)

DatasetInfo.homepage 包含数据集主页面的 URL，以便用户可以找到有关数据集的更多详细信息。
DatasetInfo.citation 包含数据集的 BibTeX 引用。

在您在模板中填写完所有这些字段后，它应该看起来像 SQuAD 加载脚本中的以下示例

def _info(self):
    return datasets.DatasetInfo(
        description=_DESCRIPTION,
        features=datasets.Features(
            {
                "id": datasets.Value("string"),
                "title": datasets.Value("string"),
                "context": datasets.Value("string"),
                "question": datasets.Value("string"),
                "answers": datasets.features.Sequence(
                    {"text": datasets.Value("string"), "answer_start": datasets.Value("int32"),}
                ),
            }
        ),
        # No default supervised_keys (as we have to pass both question
        # and context as input).
        supervised_keys=None,
        homepage="https://rajpurkar.github.io/SQuAD-explorer/",
        citation=_CITATION,
    )

多个配置

在某些情况下，您的数据集可能具有多个配置。例如，SuperGLUE 数据集是一组 5 个数据集，旨在评估语言理解任务。🤗 Datasets 提供了 BuilderConfig，它允许您为用户创建不同的配置以供选择。

让我们研究一下 SuperGLUE 加载脚本，了解如何定义多个配置。

创建一个带有有关数据集属性的 BuilderConfig 子类。这些属性可以是数据集的功能、标签类以及指向数据文件的 URL。

class SuperGlueConfig(datasets.BuilderConfig):
    """BuilderConfig for SuperGLUE."""

    def __init__(self, features, data_url, citation, url, label_classes=("False", "True"), **kwargs):
        """BuilderConfig for SuperGLUE.

        Args:
        features: *list[string]*, list of the features that will appear in the
            feature dict. Should not include "label".
        data_url: *string*, url to download the zip file from.
        citation: *string*, citation for the data set.
        url: *string*, url for information about the data set.
        label_classes: *list[string]*, the list of classes for the label if the
            label is present as a string. Non-string labels will be cast to either
            'False' or 'True'.
        **kwargs: keyword arguments forwarded to super.
        """
        # Version history:
        # 1.0.2: Fixed non-nondeterminism in ReCoRD.
        # 1.0.1: Change from the pre-release trial version of SuperGLUE (v1.9) to
        #        the full release (v2.0).
        # 1.0.0: S3 (new shuffling, sharding and slicing mechanism).
        # 0.0.2: Initial version.
        super().__init__(version=datasets.Version("1.0.2"), **kwargs)
        self.features = features
        self.label_classes = label_classes
        self.data_url = data_url
        self.citation = citation
        self.url = url

创建配置实例以指定每个配置的属性值。这使您可以灵活地指定每个配置的所有名称和描述。这些子类实例应列在 DatasetBuilder.BUILDER_CONFIGS 下

class SuperGlue(datasets.GeneratorBasedBuilder):
    """The SuperGLUE benchmark."""

    BUILDER_CONFIG_CLASS = SuperGlueConfig

    BUILDER_CONFIGS = [
        SuperGlueConfig(
            name="boolq",
            description=_BOOLQ_DESCRIPTION,
            features=["question", "passage"],
            data_url="https://dl.fbaipublicfiles.com/glue/superglue/data/v2/BoolQ.zip",
            citation=_BOOLQ_CITATION,
            url="https://github.com/google-research-datasets/boolean-questions",
        ),
        ...
        ...
        SuperGlueConfig(
            name="axg",
            description=_AXG_DESCRIPTION,
            features=["premise", "hypothesis"],
            label_classes=["entailment", "not_entailment"],
            data_url="https://dl.fbaipublicfiles.com/glue/superglue/data/v2/AX-g.zip",
            citation=_AXG_CITATION,
            url="https://github.com/rudinger/winogender-schemas",
        ),

现在，用户可以使用配置 name 加载数据集的特定配置

>>> from datasets import load_dataset
>>> dataset = load_dataset('super_glue', 'boolq')

此外，用户可以通过将构建器配置参数传递给 load_dataset() 来实例化自定义构建器配置

>>> from datasets import load_dataset
>>> dataset = load_dataset('super_glue', data_url="https://custom_url")

默认配置

用户在加载具有多个配置的数据集时必须指定配置名称。否则，🤗 Datasets 将引发 ValueError，并提示用户选择配置名称。您可以通过使用 DEFAULT_CONFIG_NAME 属性设置默认数据集配置来避免这种情况

class NewDataset(datasets.GeneratorBasedBuilder):

VERSION = datasets.Version("1.1.0")

BUILDER_CONFIGS = [
    datasets.BuilderConfig(name="first_domain", version=VERSION, description="This part of my dataset covers a first domain"),
    datasets.BuilderConfig(name="second_domain", version=VERSION, description="This part of my dataset covers a second domain"),
]

DEFAULT_CONFIG_NAME = "first_domain"

仅在有意义的情况下才使用默认配置。不要设置一个，因为用户可能更方便地在加载数据集时不指定配置。例如，多语言数据集通常为每种语言都有单独的配置。如果用户没有请求特定语言，则一个合适的默认值可能是加载数据集所有语言的聚合配置。

下载数据文件并组织拆分

在您定义了数据集的属性后，下一步是下载数据文件并根据其拆分进行组织。

在加载脚本中创建一个指向原始 SQuAD 数据文件的 URL 字典

_URL = "https://rajpurkar.github.io/SQuAD-explorer/dataset/"
_URLS = {
    "train": _URL + "train-v1.1.json",
    "dev": _URL + "dev-v1.1.json",
}

如果数据文件位于与数据集脚本相同的文件夹或存储库中，则可以传递指向文件的相对路径，而不是 URL。

DownloadManager.download_and_extract() 使用此字典下载数据文件。文件下载完毕后，请使用 SplitGenerator 来组织数据集中的每个拆分。这是一个简单的类，其中包含
- 每个拆分的 name。您应该使用标准拆分名称：Split.TRAIN、Split.TEST 和 Split.VALIDATION。
- gen_kwargs 提供了要为每个拆分加载的数据文件的路径。

您的 DatasetBuilder._split_generator() 现在应该如下所示

def _split_generators(self, dl_manager: datasets.DownloadManager) -> List[datasets.SplitGenerator]:
    urls_to_download = self._URLS
    downloaded_files = dl_manager.download_and_extract(urls_to_download)

    return [
        datasets.SplitGenerator(name=datasets.Split.TRAIN, gen_kwargs={"filepath": downloaded_files["train"]}),
        datasets.SplitGenerator(name=datasets.Split.VALIDATION, gen_kwargs={"filepath": downloaded_files["dev"]}),
    ]

生成样本

此时，您有

添加了数据集属性。
提供了有关如何下载数据文件的说明。
组织了拆分。

下一步是实际生成每个拆分的样本。

DatasetBuilder._generate_examples 使用 gen_kwargs 提供的文件路径来读取和解析数据文件。您需要编写一个函数来加载数据文件并提取列。
您的函数应生成一个包含 id_ 和数据集示例的元组。

def _generate_examples(self, filepath):
    """This function returns the examples in the raw (text) form."""
    logger.info("generating examples from = %s", filepath)
    with open(filepath) as f:
        squad = json.load(f)
        for article in squad["data"]:
            title = article.get("title", "").strip()
            for paragraph in article["paragraphs"]:
                context = paragraph["context"].strip()
                for qa in paragraph["qas"]:
                    question = qa["question"].strip()
                    id_ = qa["id"]

                    answer_starts = [answer["answer_start"] for answer in qa["answers"]]
                    answers = [answer["text"].strip() for answer in qa["answers"]]

                    # Features currently used are "context", "question", and "answers".
                    # Others are extracted here for the ease of future expansions.
                    yield id_, {
                        "title": title,
                        "context": context,
                        "question": question,
                        "id": id_,
                        "answers": {"answer_start": answer_starts, "text": answers,},
                    }

(可选) 生成数据集元数据

添加数据集元数据是包含有关数据集信息的绝佳方式。元数据存储在数据集卡片 README.md 中的 YAML 中。它包含一些信息，例如确认数据集是否正确生成所需的示例数量，以及有关数据集的信息，例如其 features。

运行以下命令在 README.md 中生成数据集元数据，并确保您的新数据集加载脚本正常工作

datasets-cli test path/to/<your-dataset-loading-script> --save_info --all_configs

如果您的数据集加载脚本通过了测试，您现在应该在数据集文件夹中有一个 README.md 文件，其中包含一个包含一些元数据的 dataset_info 字段。

上传到 Hub

脚本准备就绪后，创建数据集卡片，然后将其上传到 Hub。

恭喜你，你现在可以从 Hub 加载你的数据集了！🥳

>>> from datasets import load_dataset
>>> load_dataset("<username>/my_dataset")

高级功能

分片

如果你的数据集包含很多大文件， 🤗 Datasets 会自动并行运行你的脚本，使之超级快！如果你有数百或数千个 TAR 存档，或者像 oscar 这样的 JSONL 文件，这将非常有用。

为了使之生效，我们认为 gen_kwargs 中的文件列表是分片。因此 🤗 Datasets 可以自动生成多个工作进程来并行运行 _generate_examples，并且每个工作进程都会被分配到一个分片子集进行处理。


class MyShardedDataset(datasets.GeneratorBasedBuilder):

    def _split_generators(self, dl_manager: datasets.DownloadManager) -> List[datasets.SplitGenerator]:
        downloaded_files = dl_manager.download([f"data/shard_{i}.jsonl" for i in range(1024)])
        return [
            datasets.SplitGenerator(name=datasets.Split.TRAIN, gen_kwargs={"filepaths": downloaded_files}),
        ]

    def _generate_examples(self, filepaths):
        # Each worker can be given a slice of the original `filepaths` list defined in the `gen_kwargs`
        # so that this code can run in parallel on several shards at the same time
        for filepath in filepaths:
            ...

用户也可以在 load_dataset() 中指定 num_proc= 来指定用作工作进程的进程数量。

ArrowBasedBuilder

对于某些数据集，生成数据批次而不是逐个生成样本可能会快得多。你可以通过直接生成 Arrow 表而不是生成样本，来加速数据集生成。这在你的数据来自 Pandas DataFrames 时特别有用，因为从 Pandas 到 Arrow 的转换就像这样：

import pyarrow as pa
pa_table = pa.Table.from_pandas(df)

要生成 Arrow 表而不是单个样本，请让你的数据集构建器继承自 ArrowBasedBuilder 而不是 GeneratorBasedBuilder，并使用 _generate_tables 而不是 _generate_examples

class MySuperFastDataset(datasets.ArrowBasedBuilder):

    def _generate_tables(self, filepaths):
        idx = 0
        for filepath in filepaths:
            ...
            yield idx, pa_table
            idx += 1

不要忘记让你的脚本内存效率高，以防用户在内存少的机器上运行它们。

< > 在 GitHub 上更新