如何从 Hugging Face 数据集中随机采样数据并保存为新的 Arrow 文件

在使用 Hugging Face 的数据集进行模型训练时，有时我们并不需要整个数据集，尤其是当数据集非常大时。为了节省存储空间和提高训练效率，我们可以从数据集中随机采样一部分数据，并将其保存为新的 Arrow 文件格式。本文将介绍如何通过代码实现这一过程，并解释如何计算文件大小，以便在 dataset_info.json 文件中记录文件信息，方便后续训练使用。

1. 背景介绍

Hugging Face 提供的 datasets 库支持直接加载和操作 Arrow 格式的数据集。Arrow 是一个高效的列式数据格式，适用于大规模数据处理和分析。其高效性体现在对内存的友好支持和读取速度上，这使得它在深度学习中得到广泛应用。

然而，整个数据集可能会非常庞大，尤其是在进行大规模模型训练时。为了提高效率和减少内存占用，通常我们只需要数据集的一部分。在这种情况下，随机采样并保存为一个新的 Arrow 文件是一个很好的解决方案。

2. 代码实现

以下是从 Hugging Face 数据集中随机采样 1000 条数据，并将其保存为新的 Arrow 文件的代码：

from datasets import Dataset, DatasetDict
import os# 加载原始 Arrow 文件
dataset = Dataset.from_file("/.cache/huggingface/datasets/allenai___tulu-3-sft-mixture/default/0.0.0/55e9fd6d41c3cd1a98270dff07557bc2a1e1ba91/tulu-3-sft-mixture-train-00000-of-00001.arrow"
)# 采样 1000 条数据（随机采样）
sampled_dataset = dataset.shuffle(seed=42).select(range(1000))# 保存为新的 Arrow 文件
output_path = "/.cache/huggingface/datasets/allenai_test"
sampled_dataset.save_to_disk(output_path)# 计算文件大小
file_size = sum(os.path.getsize(os.path.join(dp, f)) for dp, dn, filenames in os.walk(output_path) for f in filenames
)# 打印采样结果和大小
print(f"采样数据集保存路径: {output_path}")
print(f"文件大小: {file_size / (1024 ** 2):.2f} MB")  # 转换为 MB
print(file_size)

代码步骤解释：

加载原始 Arrow 文件：我们通过 Dataset.from_file() 方法加载原始的 Arrow 文件。这个文件通常较大，包含了整个数据集的内容。
随机采样数据：使用 dataset.shuffle(seed=42).select(range(1000)) 随机采样出 1000 条数据。shuffle() 方法随机打乱数据集，select() 方法选择数据集的前 1000 条记录。
保存为新文件：通过 save_to_disk() 方法将采样后的数据保存为新的 Arrow 文件。这时，我们可以将这个小型的数据集用于模型训练，而不需要加载整个大数据集。
计算文件大小：通过遍历文件夹的方式，使用 os.path.getsize() 获取保存的 Arrow 文件的大小。计算结果以 MB 为单位输出，便于理解文件的存储需求。

3. 文件大小和 `dataset_info.json`

保存采样数据时，计算文件大小是非常重要的。这是因为在 Hugging Face 的数据集格式中，dataset_info.json 文件记录了数据集的基本信息，包括数据集的大小、特征、列数等。在训练时，Hugging Face 会根据 dataset_info.json 文件的信息来进行数据加载和管理。确保文件大小准确，可以帮助在加载数据集时正确管理内存和硬盘空间。

下面是新的dataset_info.json文件内容，需要改的地方有
“num_bytes”: 3781998,
“num_examples”: 1000,
“download_size”: 3781998,
“dataset_size”: 3781998,
“size_in_bytes”: 3781998
这些，这里的3781998就是上面的file_size，num_examples是上面提到的采样1000条数据。

{"description": "A sampled version of tulu-3-sft-mixture dataset with 1000 examples.","citation": "","homepage": "","license": "","features": {"id": {"dtype": "string","_type": "Value"},"messages": [{"content": {"dtype": "string","_type": "Value"},"role": {"dtype": "string","_type": "Value"}}],"source": {"dtype": "string","_type": "Value"}},"builder_name": "parquet","dataset_name": "tulu-3-sft-mixture","config_name": "default","version": {"version_str": "0.0.0","major": 0,"minor": 0,"patch": 0},"splits": {"train": {"name": "train","num_bytes": 3781998,"num_examples": 1000,"shard_lengths": [1000],"dataset_name": "tulu-3-sft-mixture"}},"download_checksums": {},"download_size": 3781998,"dataset_size": 3781998,"size_in_bytes": 3781998
}

然后记得把之前的数据集文件夹改名为其他，比如改成这里的allenai___tulu-3-sft-mixture1，然后将新的数据集放到/.cache/huggingface/datasets/allenai___tulu-3-sft-mixture/default/0.0.0/55e9fd6d41c3cd1a98270dff07557bc2a1e1ba91/这个路径下，请注意，这个路径应该与hf下载下来的路径完全相同。这个是新建的，可以使用LInux命令来新建：mkdir命令后面记得加-p参数

mkdir -p /.cache/huggingface/datasets/allenai___tulu-3-sft-mixture/default/0.0.0/55e9fd6d41c3cd1a98270dff07557bc2a1e1ba91/

在这里插入图片描述
如下图所示：需要在这里放入新的arrow文件（只需要放自己的那1000条数据的arrow即可，其他cache开头的arrow文件是系统自己生成的，不用管）和dataset_info.json文件（这个是需要按照上面更改后的，不能用之前的），然后arrow文件记得改名：tulu-3-sft-mixture-train-00000-of-00001.arrow，这里记得按照原始文件中的arrow文件命名格式，比如原来的是tulu-3-sft-mixture-train-00000-of-00006.arrow，tulu-3-sft-mixture-train-00001-of-00006.arrow这样，后面的00006是分块的个数，由于我们只有一个arrow文件，后面的00006应该改为00001。

在这里插入图片描述

4. 如何使用新的 Arrow 文件进行训练

在 Hugging Face 上使用数据集时，我们通常指定一个数据集路径，比如：

--dataset_mixer_list allenai/tulu-3-sft-mixture 1.0

这个参数指定了使用某个数据集进行训练。当我们使用采样的 Arrow 文件时，文件路径应该指向我们保存的采样文件（这里由于我们用新的arrow覆盖掉了原来的文件，所以不用指定新的路径，默认即可），而无需更改 --dataset_mixer_list 参数。这样，我们就可以利用数据集的一部分进行训练，而不需要更改 Hugging Face 数据集的整体配置。

5. 下载后的文件为何变成 Arrow 格式

在使用 Hugging Face 的数据集时，很多时候我们会下载数据集并看到它是以 .arrow 格式存储的。这是因为 Arrow 格式在性能和存储上优于其他格式，尤其是在大规模数据集的处理过程中，能够提供更高效的内存和磁盘使用。下载到本地后，文件会以 Arrow 格式存储，便于后续使用和处理。