处理过程

🤗 数据集提供了许多修改数据集结构和内容的工具。这些工具对于整理数据集、创建额外的列、在特征和格式之间转换以及更多操作非常重要。

本指南将向您展示如何

重新排序行并拆分数据集。
重命名和删除列，以及其他常见的列操作。
将处理函数应用于数据集中的每个示例。
连接数据集。
应用自定义格式转换。
保存和导出处理后的数据集。

有关特定于处理其他数据集模态的更多详细信息，请查看处理音频数据集指南、处理图像数据集指南或处理文本数据集指南。

本指南中的示例使用 MRPC 数据集，但您可以随意加载任何您选择的数据集并继续操作！

>>> from datasets import load_dataset
>>> dataset = load_dataset("glue", "mrpc", split="train")

本指南中的所有处理方法都会返回一个新的Dataset 对象。修改不会就地进行。请注意不要覆盖您之前的数据集！

排序、洗牌、选择、分割和分片

有几个函数可以重新排列数据集的结构。这些函数对于仅选择所需的行、创建训练和测试分割以及将非常大的数据集分割成较小的块很有用。

排序

使用 sort() 根据列值的数值对列值进行排序。提供的列必须与 NumPy 兼容。

>>> dataset["label"][:10]
[1, 0, 1, 0, 1, 1, 0, 1, 0, 0]
>>> sorted_dataset = dataset.sort("label")
>>> sorted_dataset["label"][:10]
[0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
>>> sorted_dataset["label"][-10:]
[1, 1, 1, 1, 1, 1, 1, 1, 1, 1]

在后台，这会创建一个根据列值排序的索引列表。然后使用此索引映射来访问底层 Arrow 表中的正确行。

洗牌

shuffle() 函数随机重新排列列值。如果希望更详细地控制用于洗牌数据集的算法，可以在此函数中指定 generator 参数以使用不同的 numpy.random.Generator。

>>> shuffled_dataset = sorted_dataset.shuffle(seed=42)
>>> shuffled_dataset["label"][:10]
[1, 1, 1, 0, 1, 1, 1, 1, 1, 0]

洗牌会获取索引列表 [0:len(my_dataset)] 并对其进行洗牌以创建索引映射。但是，一旦您的 Dataset 具有索引映射，速度可能会降低 10 倍。这是因为需要一个额外的步骤来获取使用索引映射读取的行索引，最重要的是，您不再读取连续的数据块。要恢复速度，您需要使用 Dataset.flatten_indices() 再次将整个数据集写入磁盘，这将删除索引映射。或者，您可以切换到 IterableDataset 并利用其快速的近似洗牌 IterableDataset.shuffle()

>>> iterable_dataset = dataset.to_iterable_dataset(num_shards=128)
>>> shuffled_iterable_dataset = iterable_dataset.shuffle(seed=42, buffer_size=1000)

选择和过滤

有两种选择可以在数据集中过滤行：select() 和 filter()。

select() 根据索引列表返回行

>>> small_dataset = dataset.select([0, 10, 20, 30, 40, 50])
>>> len(small_dataset)
6

filter() 返回与指定条件匹配的行

>>> start_with_ar = dataset.filter(lambda example: example["sentence1"].startswith("Ar"))
>>> len(start_with_ar)
6
>>> start_with_ar["sentence1"]
['Around 0335 GMT , Tab shares were up 19 cents , or 4.4 % , at A $ 4.56 , having earlier set a record high of A $ 4.57 .',
'Arison said Mann may have been one of the pioneers of the world music movement and he had a deep love of Brazilian music .',
'Arts helped coach the youth on an eighth-grade football team at Lombardi Middle School in Green Bay .',
'Around 9 : 00 a.m. EDT ( 1300 GMT ) , the euro was at $ 1.1566 against the dollar , up 0.07 percent on the day .',
"Arguing that the case was an isolated example , Canada has threatened a trade backlash if Tokyo 's ban is not justified on scientific grounds .",
'Artists are worried the plan would harm those who need help most - performers who have a difficult time lining up shows .'
]

如果设置 with_indices=True，filter() 也可以按索引进行过滤

>>> even_dataset = dataset.filter(lambda example, idx: idx % 2 == 0, with_indices=True)
>>> len(even_dataset)
1834
>>> len(dataset) / 2
1834.0

除非要保留的索引列表是连续的，否则这些方法也会在后台创建索引映射。

分割

如果您的数据集还没有训练和测试分割，则 train_test_split() 函数会创建训练和测试分割。这允许您调整每个分割中样本的相对比例或绝对数量。在下面的示例中，使用 test_size 参数创建一个测试分割，该分割是原始数据集的 10%

>>> dataset.train_test_split(test_size=0.1)
{'train': Dataset(schema: {'sentence1': 'string', 'sentence2': 'string', 'label': 'int64', 'idx': 'int32'}, num_rows: 3301),
'test': Dataset(schema: {'sentence1': 'string', 'sentence2': 'string', 'label': 'int64', 'idx': 'int32'}, num_rows: 367)}
>>> 0.1 * len(dataset)
366.8

默认情况下，分割会进行洗牌，但您可以设置 shuffle=False 以防止洗牌。

分片

🤗 Datasets 支持分片以将非常大的数据集划分为预定义数量的块。在 shard() 中指定 num_shards 参数以确定将数据集分成多少个分片。您还需要使用 index 参数提供要返回的分片。

例如，imdb 数据集有 25000 个示例

>>> from datasets import load_dataset
>>> dataset = load_dataset("imdb", split="train")
>>> print(dataset)
Dataset({
    features: ['text', 'label'],
    num_rows: 25000
})

将数据集分成四个块后，第一个分片将只有 6250 个示例

>>> dataset.shard(num_shards=4, index=0)
Dataset({
    features: ['text', 'label'],
    num_rows: 6250
})
>>> print(25000/4)
6250.0

重命名、删除、转换和展平

以下函数允许您修改数据集的列。这些函数对于重命名或删除列、将列更改为一组新的特征以及展平嵌套的列结构很有用。

重命名

当需要重命名数据集中的一列时，使用 rename_column()。与原始列关联的特征实际上会移动到新的列名下，而不是仅仅替换原始列。

向 rename_column() 提供原始列的名称和新列名

>>> dataset
Dataset({
    features: ['sentence1', 'sentence2', 'label', 'idx'],
    num_rows: 3668
})
>>> dataset = dataset.rename_column("sentence1", "sentenceA")
>>> dataset = dataset.rename_column("sentence2", "sentenceB")
>>> dataset
Dataset({
    features: ['sentenceA', 'sentenceB', 'label', 'idx'],
    num_rows: 3668
})

删除

当需要删除一个或多个列时，将要删除的列名提供给 remove_columns() 函数。通过提供列名列表来删除多个列

>>> dataset = dataset.remove_columns("label")
>>> dataset
Dataset({
    features: ['sentence1', 'sentence2', 'idx'],
    num_rows: 3668
})
>>> dataset = dataset.remove_columns(["sentence1", "sentence2"])
>>> dataset
Dataset({
    features: ['idx'],
    num_rows: 3668
})

相反，select_columns() 选择要保留的一个或多个列并删除其余列。此函数接受一个或一个列名列表

>>> dataset
Dataset({
    features: ['sentence1', 'sentence2', 'label', 'idx'],
    num_rows: 3668
})
>>> dataset = dataset.select_columns(['sentence1', 'sentence2', 'idx'])
>>> dataset
Dataset({
    features: ['sentence1', 'sentence2', 'idx'],
    num_rows: 3668
})
>>> dataset = dataset.select_columns('idx')
>>> dataset
Dataset({
    features: ['idx'],
    num_rows: 3668
})

转换

the cast() 函数转换一个或多个列的特征类型。此函数接受新的 Features 作为其参数。下面的示例演示了如何更改 ClassLabel 和 Value 特征

>>> dataset.features
{'sentence1': Value(dtype='string', id=None),
'sentence2': Value(dtype='string', id=None),
'label': ClassLabel(num_classes=2, names=['not_equivalent', 'equivalent'], names_file=None, id=None),
'idx': Value(dtype='int32', id=None)}

>>> 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(dtype='string', id=None),
'sentence2': Value(dtype='string', id=None),
'label': ClassLabel(num_classes=2, names=['negative', 'positive'], names_file=None, id=None),
'idx': Value(dtype='int64', id=None)}

转换仅在原始特征类型和新特征类型兼容时才有效。例如，如果原始列仅包含 1 和 0，则可以将特征类型为 Value("int32") 的列转换为 Value("bool")。

使用 cast_column() 函数更改单个列的特征类型。将列名及其新的特征类型作为参数传递

>>> dataset.features
{'audio': Audio(sampling_rate=44100, mono=True, id=None)}

>>> dataset = dataset.cast_column("audio", Audio(sampling_rate=16000))
>>> dataset.features
{'audio': Audio(sampling_rate=16000, mono=True, id=None)}

扁平化

有时，一列可以是几种类型的嵌套结构。查看下面来自 SQuAD 数据集的嵌套结构

>>> from datasets import load_dataset
>>> dataset = load_dataset("squad", split="train")
>>> dataset.features
{'answers': Sequence(feature={'text': Value(dtype='string', id=None), 'answer_start': Value(dtype='int32', id=None)}, length=-1, id=None),
'context': Value(dtype='string', id=None),
'id': Value(dtype='string', id=None),
'question': Value(dtype='string', id=None),
'title': Value(dtype='string', id=None)}

the answers 字段包含两个子字段：text 和 answer_start。使用 flatten() 函数将子字段提取到它们自己的单独列中

>>> flat_dataset = dataset.flatten()
>>> flat_dataset
Dataset({
    features: ['id', 'title', 'context', 'question', 'answers.text', 'answers.answer_start'],
 num_rows: 87599
})

注意子字段现在是如何成为它们自己的独立列的：answers.text 和 answers.answer_start。

映射

🤗 Datasets 的一些更强大的应用程序来自使用 map() 函数。 map() 的主要目的是加速处理函数。它允许您将处理函数应用于数据集中每个示例，独立或批量。此函数甚至可以创建新行和新列。

在下面的示例中，在数据集中每个 sentence1 值前面添加 'My sentence: '。

首先创建一个函数，将 'My sentence: ' 添加到每个句子的开头。该函数需要接受并输出一个 dict

>>> def add_prefix(example):
...     example["sentence1"] = 'My sentence: ' + example["sentence1"]
...     return example

现在使用 map() 将 add_prefix 函数应用于整个数据集

>>> updated_dataset = small_dataset.map(add_prefix)
>>> updated_dataset["sentence1"][:5]
['My sentence: Amrozi accused his brother , whom he called " the witness " , of deliberately distorting his evidence .',
"My sentence: Yucaipa owned Dominick 's before selling the chain to Safeway in 1998 for $ 2.5 billion .",
'My sentence: They had published an advertisement on the Internet on June 10 , offering the cargo for sale , he added .',
'My sentence: Around 0335 GMT , Tab shares were up 19 cents , or 4.4 % , at A $ 4.56 , having earlier set a record high of A $ 4.57 .',
]

让我们看另一个示例，但这次，您将使用 map() 删除一列。当您删除一列时，只有在将示例提供给映射函数后才会删除它。这允许映射函数在删除列之前使用列的内容。

在 map() 中使用 remove_columns 参数指定要删除的列

>>> updated_dataset = dataset.map(lambda example: {"new_sentence": example["sentence1"]}, remove_columns=["sentence1"])
>>> updated_dataset.column_names
['sentence2', 'label', 'idx', 'new_sentence']

🤗 Datasets 还具有一个 remove_columns() 函数，它更快，因为它不会复制其余列的数据。

如果设置 with_indices=True，您还可以将 map() 与索引一起使用。下面的示例将索引添加到每个句子的开头

>>> updated_dataset = dataset.map(lambda example, idx: {"sentence2": f"{idx}: " + example["sentence2"]}, with_indices=True)
>>> updated_dataset["sentence2"][:5]
['0: Referring to him as only " the witness " , Amrozi accused his brother of deliberately distorting his evidence .',
 "1: Yucaipa bought Dominick 's in 1995 for $ 693 million and sold it to Safeway for $ 1.8 billion in 1998 .",
 "2: On June 10 , the ship 's owners had published an advertisement on the Internet , offering the explosives for sale .",
 '3: Tab shares jumped 20 cents , or 4.6 % , to set a record closing high at A $ 4.57 .',
 '4: PG & E Corp. shares jumped $ 1.63 or 8 percent to $ 21.03 on the New York Stock Exchange on Friday .'
]

多进程

通过并行化 CPU 上的进程，多进程可以显著加快处理速度。在 map() 中设置 num_proc 参数以设置要使用的进程数

>>> updated_dataset = dataset.map(lambda example, idx: {"sentence2": f"{idx}: " + example["sentence2"]}, with_indices=True, num_proc=4)

如果设置 with_rank=True，则 map() 也适用于进程的排名。这类似于 with_indices 参数。如果 with_indices 参数已存在，则映射函数中的 with_rank 参数位于其之后。

>>> import torch
>>> from multiprocess import set_start_method
>>> from transformers import AutoTokenizer, AutoModelForCausalLM 
>>> from datasets import load_dataset
>>> 
>>> # Get an example dataset
>>> dataset = load_dataset("fka/awesome-chatgpt-prompts", split="train")
>>> 
>>> # Get an example model and its tokenizer 
>>> model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen1.5-0.5B-Chat").eval()
>>> tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen1.5-0.5B-Chat")
>>>
>>> def gpu_computation(batch, rank):
...     # Move the model on the right GPU if it's not there already
...     device = f"cuda:{(rank or 0) % torch.cuda.device_count()}"
...     model.to(device)
...     
...     # Your big GPU call goes here, for example:
...     chats = [[
...         {"role": "system", "content": "You are a helpful assistant."},
...         {"role": "user", "content": prompt}
...     ] for prompt in batch["prompt"]]
...     texts = [tokenizer.apply_chat_template(
...         chat,
...         tokenize=False,
...         add_generation_prompt=True
...     ) for chat in chats]
...     model_inputs = tokenizer(texts, padding=True, return_tensors="pt").to(device)
...     with torch.no_grad():
...         outputs = model.generate(**model_inputs, max_new_tokens=512)
...     batch["output"] = tokenizer.batch_decode(outputs, skip_special_tokens=True)
...     return batch
>>>
>>> if __name__ == "__main__":
...     set_start_method("spawn")
...     updated_dataset = dataset.map(
...         gpu_computation,
...         batched=True,
...         batch_size=16,
...         with_rank=True,
...         num_proc=torch.cuda.device_count(),  # one process per GPU
...     )

排名的主要用例是在多个 GPU 上并行化计算。这需要设置 multiprocess.set_start_method("spawn")。如果不这样做，您将收到以下 CUDA 错误

RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method.

批量处理

the map() 函数支持使用示例批次。通过设置 batched=True 对批次进行操作。默认批次大小为 1000，但您可以使用 batch_size 参数对其进行调整。批量处理支持有趣的应用程序，例如将长句子拆分为较短的块和数据增强。

拆分长示例

当示例过长时，您可能希望将其拆分为几个较小的块。首先创建一个函数，该函数

将 sentence1 字段拆分为 50 个字符的块。
将所有块堆叠在一起以创建新的数据集。

>>> def chunk_examples(examples):
...     chunks = []
...     for sentence in examples["sentence1"]:
...         chunks += [sentence[i:i + 50] for i in range(0, len(sentence), 50)]
...     return {"chunks": chunks}

使用 map() 应用该函数

>>> chunked_dataset = dataset.map(chunk_examples, batched=True, remove_columns=dataset.column_names)
>>> chunked_dataset[:10]
{'chunks': ['Amrozi accused his brother , whom he called " the ',
            'witness " , of deliberately distorting his evidenc',
            'e .',
            "Yucaipa owned Dominick 's before selling the chain",
            ' to Safeway in 1998 for $ 2.5 billion .',
            'They had published an advertisement on the Interne',
            't on June 10 , offering the cargo for sale , he ad',
            'ded .',
            'Around 0335 GMT , Tab shares were up 19 cents , or',
            ' 4.4 % , at A $ 4.56 , having earlier set a record']}

注意句子现在是如何拆分为较短的块的，并且数据集中有更多行。

>>> dataset
Dataset({
 features: ['sentence1', 'sentence2', 'label', 'idx'],
 num_rows: 3668
})
>>> chunked_dataset
Dataset({
    features: ['chunks'],
    num_rows: 10470
})

数据增强

the map() 函数也可用于数据增强。以下示例为句子中掩码标记生成其他单词。

加载并在 🤗 Transformers 的 FillMaskPipeline 中使用 RoBERTA 模型

>>> from random import randint
>>> from transformers import pipeline

>>> fillmask = pipeline("fill-mask", model="roberta-base")
>>> mask_token = fillmask.tokenizer.mask_token
>>> smaller_dataset = dataset.filter(lambda e, i: i<100, with_indices=True)

创建一个函数，在句子中随机选择一个词进行掩码。该函数还应返回原始句子和 RoBERTA 生成的前两个替换。

>>> def augment_data(examples):
...     outputs = []
...     for sentence in examples["sentence1"]:
...         words = sentence.split(' ')
...         K = randint(1, len(words)-1)
...         masked_sentence = " ".join(words[:K]  + [mask_token] + words[K+1:])
...         predictions = fillmask(masked_sentence)
...         augmented_sequences = [predictions[i]["sequence"] for i in range(3)]
...         outputs += [sentence] + augmented_sequences
...
...     return {"data": outputs}

使用 map() 将函数应用于整个数据集

>>> augmented_dataset = smaller_dataset.map(augment_data, batched=True, remove_columns=dataset.column_names, batch_size=8)
>>> augmented_dataset[:9]["data"]
['Amrozi accused his brother , whom he called " the witness " , of deliberately distorting his evidence .',
 'Amrozi accused his brother, whom he called " the witness ", of deliberately withholding his evidence.',
 'Amrozi accused his brother, whom he called " the witness ", of deliberately suppressing his evidence.',
 'Amrozi accused his brother, whom he called " the witness ", of deliberately destroying his evidence.',
 "Yucaipa owned Dominick 's before selling the chain to Safeway in 1998 for $ 2.5 billion .",
 'Yucaipa owned Dominick Stores before selling the chain to Safeway in 1998 for $ 2.5 billion.',
 "Yucaipa owned Dominick's before selling the chain to Safeway in 1998 for $ 2.5 billion.",
 'Yucaipa owned Dominick Pizza before selling the chain to Safeway in 1998 for $ 2.5 billion.'
]

对于每个原始句子，RoBERTA 使用三个备选方案增强了一个随机单词。原始单词 distorting 由 withholding、suppressing 和 destroying 补充。

处理多个拆分

许多数据集都具有可以同时使用 DatasetDict.map() 处理的拆分。例如，通过以下方式对训练和测试拆分中的 sentence1 字段进行标记化

>>> from datasets import load_dataset

# load all the splits
>>> dataset = load_dataset('glue', 'mrpc')
>>> encoded_dataset = dataset.map(lambda examples: tokenizer(examples["sentence1"]), batched=True)
>>> encoded_dataset["train"][0]
{'sentence1': 'Amrozi accused his brother , whom he called " the witness " , of deliberately distorting his evidence .',
'sentence2': 'Referring to him as only " the witness " , Amrozi accused his brother of deliberately distorting his evidence .',
'label': 1,
'idx': 0,
'input_ids': [  101,  7277,  2180,  5303,  4806,  1117,  1711,   117,  2292, 1119,  1270,   107,  1103,  7737,   107,   117,  1104,  9938, 4267, 12223, 21811,  1117,  2554,   119,   102],
'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]
}

分布式使用

当您在分布式环境中使用 map() 时，也应该使用 torch.distributed.barrier。这可以确保主进程执行映射，而其他进程加载结果，从而避免重复工作。

以下示例演示了如何使用 torch.distributed.barrier 来同步进程

>>> from datasets import Dataset
>>> import torch.distributed

>>> dataset1 = Dataset.from_dict({"a": [0, 1, 2]})

>>> if training_args.local_rank > 0:
...     print("Waiting for main process to perform the mapping")
...     torch.distributed.barrier()

>>> dataset2 = dataset1.map(lambda x: {"a": x["a"] + 1})

>>> if training_args.local_rank == 0:
...     print("Loading results from main process")
...     torch.distributed.barrier()

批处理

batch() 方法允许您将数据集中的样本分组为批次。当您想要为训练或评估创建数据批次时，这尤其有用，尤其是在使用深度学习模型时。

以下是如何使用 batch() 方法的示例

>>> from datasets import load_dataset
>>> dataset = load_dataset("rotten_tomatoes", split="train")
>>> batched_dataset = dataset.batch(batch_size=4)
>>> batched_dataset[0]
{'text': ['the rock is destined to be the 21st century\'s new " conan " and that he\'s going to make a splash even greater than arnold schwarzenegger , jean-claud van damme or steven segal .',
        'the gorgeously elaborate continuation of " the lord of the rings " trilogy is so huge that a column of words cannot adequately describe co-writer/director peter jackson\'s expanded vision of j . r . r . tolkien\'s middle-earth .',
        'effective but too-tepid biopic',
        'if you sometimes like to go to the movies to have fun , wasabi is a good place to start .'],
'label': [1, 1, 1, 1]}

batch() 方法接受以下参数

batch_size (int)：每个批次中的样本数。
drop_last_batch (bool，默认为 False)：如果数据集大小不能被批次大小整除，是否丢弃最后一个不完整的批次。
num_proc (int，可选，默认为 None)：用于多处理的进程数。如果为 None，则不使用多处理。这可以显着加快大型数据集的批处理速度。

请注意，Dataset.batch() 返回一个新的 Dataset，其中每个项目都是来自原始数据集的多个样本的批次。如果要分批处理数据，则应直接使用批处理的 map()，它将函数应用于批次，但输出数据集是非批处理的。

连接

如果单独的数据集共享相同的列类型，则可以将它们连接起来。使用 concatenate_datasets() 连接数据集

>>> from datasets import concatenate_datasets, load_dataset

>>> bookcorpus = load_dataset("bookcorpus", split="train")
>>> wiki = load_dataset("wikipedia", "20220301.en", split="train")
>>> wiki = wiki.remove_columns([col for col in wiki.column_names if col != "text"])  # only keep the 'text' column

>>> assert bookcorpus.features.type == wiki.features.type
>>> bert_dataset = concatenate_datasets([bookcorpus, wiki])

只要数据集的行数相同，您也可以通过将 axis=1 设置为 1 来水平连接两个数据集

>>> from datasets import Dataset
>>> bookcorpus_ids = Dataset.from_dict({"ids": list(range(len(bookcorpus)))})
>>> bookcorpus_with_ids = concatenate_datasets([bookcorpus, bookcorpus_ids], axis=1)

交错

您还可以通过从每个数据集交替获取示例来混合多个数据集，以创建一个新的数据集。这称为交错，由 interleave_datasets() 函数启用。 interleave_datasets() 和 concatenate_datasets() 都适用于常规 Dataset 和 IterableDataset 对象。请参阅流指南，了解如何交错 IterableDataset 对象的示例。

您可以为每个原始数据集定义采样概率，以指定如何交错数据集。在这种情况下，新数据集是通过从随机数据集一个接一个地获取示例构建的，直到其中一个数据集用完样本。

>>> from datasets import Dataset, interleave_datasets
>>> seed = 42
>>> probabilities = [0.3, 0.5, 0.2]
>>> d1 = Dataset.from_dict({"a": [0, 1, 2]})
>>> d2 = Dataset.from_dict({"a": [10, 11, 12, 13]})
>>> d3 = Dataset.from_dict({"a": [20, 21, 22]})
>>> dataset = interleave_datasets([d1, d2, d3], probabilities=probabilities, seed=seed)
>>> dataset["a"]
[10, 11, 20, 12, 0, 21, 13]

您还可以指定 stopping_strategy。默认策略 first_exhausted 是一种子采样策略，即一旦其中一个数据集用完样本，数据集构建就会停止。您可以指定 stopping_strategy=all_exhausted 来执行过采样策略。在这种情况下，一旦每个数据集中的每个样本至少被添加一次，数据集构建就会停止。在实践中，这意味着如果一个数据集用尽，它将返回到该数据集的开头，直到达到停止条件。请注意，如果未指定采样概率，则新数据集将具有 max_length_datasets*nb_dataset 个样本。

>>> d1 = Dataset.from_dict({"a": [0, 1, 2]})
>>> d2 = Dataset.from_dict({"a": [10, 11, 12, 13]})
>>> d3 = Dataset.from_dict({"a": [20, 21, 22]})
>>> dataset = interleave_datasets([d1, d2, d3], stopping_strategy="all_exhausted")
>>> dataset["a"]
[0, 10, 20, 1, 11, 21, 2, 12, 22, 0, 13, 20]

格式

set_format() 函数将列的格式更改为与某些常见数据格式兼容。在 type 参数中指定您想要的输出以及要格式化的列。格式化是即时应用的。

例如，通过设置 type="torch" 创建 PyTorch 张量

>>> import torch
>>> dataset.set_format(type="torch", columns=["input_ids", "token_type_ids", "attention_mask", "label"])

with_format() 函数也更改列的格式，除了它返回一个新的 Dataset 对象

>>> dataset = dataset.with_format(type="torch", columns=["input_ids", "token_type_ids", "attention_mask", "label"])

🤗 Datasets 还支持其他常见的数据格式，例如 NumPy、Pandas 和 JAX。查看将数据集与 TensorFlow 一起使用指南，以获取有关如何有效创建 TensorFlow 数据集的更多详细信息。

如果需要将数据集重置为其原始格式，请使用 reset_format() 函数

>>> dataset.format
{'type': 'torch', 'format_kwargs': {}, 'columns': ['label'], 'output_all_columns': False}
>>> dataset.reset_format()
>>> dataset.format
{'type': 'python', 'format_kwargs': {}, 'columns': ['idx', 'label', 'sentence1', 'sentence2'], 'output_all_columns': False}

格式转换

set_transform() 函数即时应用自定义格式转换。此函数将替换任何先前指定的格式。例如，您可以使用此函数即时分词和填充标记。分词仅在访问示例时应用

>>> from transformers import AutoTokenizer

>>> tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
>>> def encode(batch):
...     return tokenizer(batch["sentence1"], batch["sentence2"], padding="longest", truncation=True, max_length=512, return_tensors="pt")
>>> dataset.set_transform(encode)
>>> dataset.format
{'type': 'custom', 'format_kwargs': {'transform': <function __main__.encode(batch)>}, 'columns': ['idx', 'label', 'sentence1', 'sentence2'], 'output_all_columns': False}

您还可以使用 set_transform() 函数解码 Features 不支持的格式。例如，Audio 特征使用 soundfile（一个快速且易于安装的库），但它不支持不太常见的音频格式。在这里，您可以使用 set_transform() 即时应用自定义解码转换。您可以自由使用任何您喜欢的库来解码音频文件。

下面的示例使用 pydub 包打开 soundfile 不支持的音频格式

>>> import numpy as np
>>> from pydub import AudioSegment

>>> audio_dataset_amr = Dataset.from_dict({"audio": ["audio_samples/audio.amr"]})

>>> def decode_audio_with_pydub(batch, sampling_rate=16_000):
...     def pydub_decode_file(audio_path):
...         sound = AudioSegment.from_file(audio_path)
...         if sound.frame_rate != sampling_rate:
...             sound = sound.set_frame_rate(sampling_rate)
...         channel_sounds = sound.split_to_mono()
...         samples = [s.get_array_of_samples() for s in channel_sounds]
...         fp_arr = np.array(samples).T.astype(np.float32)
...         fp_arr /= np.iinfo(samples[0].typecode).max
...         return fp_arr
...
...     batch["audio"] = [pydub_decode_file(audio_path) for audio_path in batch["audio"]]
...     return batch

>>> audio_dataset_amr.set_transform(decode_audio_with_pydub)

保存

完成数据集处理后，您可以使用 save_to_disk() 保存并稍后重用它。

通过提供要保存到的目录路径来保存数据集

>>> encoded_dataset.save_to_disk("path/of/my/dataset/directory")

使用 load_from_disk() 函数重新加载数据集

>>> from datasets import load_from_disk
>>> reloaded_dataset = load_from_disk("path/of/my/dataset/directory")

想要将您的数据集保存到云存储提供商吗？请阅读我们的云存储指南，了解如何将数据集保存到 AWS 或 Google Cloud Storage。

导出

🤗 Datasets 也支持导出，因此您可以在其他应用程序中使用您的数据集。下表显示了当前支持的导出文件格式。

文件类型	导出方法
CSV	Dataset.to_csv()
JSON	Dataset.to_json()
Parquet	Dataset.to_parquet()
SQL	Dataset.to_sql()
内存中的 Python 对象	Dataset.to_pandas() 或 Dataset.to_dict()

例如，您可以像这样将数据集导出到 CSV 文件

>>> encoded_dataset.to_csv("path/of/my/dataset.csv")

< > 在 GitHub 上更新