自学内容网 自学内容网
自学内容网 > 正文

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

🕗 发布于 2024-11-30 11:15 huggingface  

如何从 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)
代码步骤解释：

  1. 加载原始 Arrow 文件：我们通过 Dataset.from_file() 方法加载原始的 Arrow 文件。这个文件通常较大，包含了整个数据集的内容。

  2. 随机采样数据：使用 dataset.shuffle(seed=42).select(range(1000)) 随机采样出 1000 条数据。shuffle() 方法随机打乱数据集，select() 方法选择数据集的前 1000 条记录。

  3. 保存为新文件：通过 save_to_disk() 方法将采样后的数据保存为新的 Arrow 文件。这时，我们可以将这个小型的数据集用于模型训练，而不需要加载整个大数据集。

  4. 计算文件大小：通过遍历文件夹的方式，使用 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 格式存储，便于后续使用和处理。

6. 结论

通过从 Hugging Face 数据集中随机采样一部分数据并保存为新的 Arrow 文件，我们可以更高效地进行模型训练，特别是当数据集庞大时。通过计算文件大小并更新 dataset_info.json 文件，我们可以确保训练过程中数据管理的准确性。

这种方法不仅适用于大数据集，也为需要快速原型设计或进行小规模实验的研究人员提供了便利。


原文地址：https://blog.csdn.net/shizheng_Li/article/details/144132851

免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！

相关文章

自学内容网
Copyright © 2015-2024