创建数据集加载脚本

如果你的数据集是以下格式之一，则可能不需要数据集加载脚本：CSV、JSON、JSON Lines、文本、图像、音频或 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"

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

下载数据文件并组织拆分

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

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

_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 存档或 JSONL 文件（例如 oscar），这将很有帮助。

为了使其正常工作，我们将 gen_kwargs 中的文件列表视为分片 (shards)。因此，🤗 Datasets 可以自动生成多个 worker 并行运行 _generate_examples，并且每个 worker 会被分配一个分片子集进行处理。


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= 来指定用作 worker 的进程数量。

ArrowBasedBuilder

对于某些数据集，批量 (batches) 生成数据可能比逐个生成 example 快得多。您可以通过直接生成 Arrow 表格 (tables) 而不是 example 来加速数据集的生成。如果您的数据来自 Pandas DataFrames，这将特别有用，因为从 Pandas 转换为 Arrow 非常简单。

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

要生成 Arrow 表格而不是单个 example，请使您的数据集构建器 (builder) 继承自 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

请记住保持您的脚本内存高效，以防用户在 RAM 较低的机器上运行它们。

< > 在 GitHub 上更新