Huggingface数据集采样之后得到的arrow文件无法用Dataset.load_from_disk加载解决方法

问题起源

我现在有数据集

我想要采样1/10构成一个新的数据集。起初我只是简单的使用save_to_disk保存，如下代码：

from datasets import Dataset
import os

# 原数据集路径
source_dataset_path = "~/.cache/huggingface/datasets/allenai___tulu-3-sft-mixture-1024-test/train"
# 新数据集保存路径
output_dataset_path = "~/.cache/huggingface/datasets/allenai___tulu-3-sft-mixture-1024-73k/train"

# 加载原数据集
dataset = Dataset.load_from_disk(source_dataset_path)

# 设置随机种子，确保每次运行结果相同
seed = 123
dataset = dataset.shuffle(seed=seed)

# 随机采样1/10的数据
sampled_dataset = dataset.select(range(len(dataset)//10))

# 输出采样后的数据集大小
print(f"采样后的数据集包含 {len(sampled_dataset)} 条数据")

# 保存采样后的数据集到新目录
sampled_dataset.save_to_disk(output_dataset_path)

print(f"数据集已保存到: {output_dataset_path}") 

但是后面使用的时候遇到问题，当执行下面的代码，load_from_disk会报错：

from datasets import load_dataset, load_from_disk

# 加载数据集
input_dir="~/.cache/huggingface/datasets/allenai___tulu-3-sft-mixture-1024-73k"
dataset = load_from_disk(input_dir)

# 指定保存路径
output_dir = "xxxx/data/tulu-3-sft-mixture-seqlen-1024-json-73k"

# 将数据集转换为 JSON 格式
for split in dataset.keys():  # 处理所有划分 (train, validation 等)
    dataset[split].to_json(f"{output_dir}/{split}.json", orient="records", lines=True)

报错如下：我里面有.arrow文件，却显示neither a Dataset directory nor a DatasetDict directory.

 (base)@test:~$ python arrow2json.py 
Traceback (most recent call last):
lib/python3.10/site-packages/datasets/load.py", line 2218, in load_from_disk
    raise FileNotFoundError(
FileNotFoundError: Directory ~/.cache/huggingface/datasets/allenai___tulu-3-sft-mixture-1024-73k is neither a Dataset directory nor a DatasetDict directory.

提问GPT4o：

它是通过什么来判断不是Dataset directory nor a DatasetDict directory的，是通过~/.cache/huggingface/datasets/allenai___tulu-3-sft-mixture-1024-73k/dataset_dict.json这个嘛

是的，Hugging Face 的 datasets 库会通过检查目录中的特定文件来确定一个目录是否是一个有效的 Dataset 或 DatasetDict 目录。主要通过以下几个文件来进行判断：

1. dataset_dict.json：这个文件是用于 DatasetDict 类型的数据集。DatasetDict 是一个包含多个数据集分割（如 train、validation 等）的字典。如果这个文件存在，Hugging Face 会认为该目录是一个 DatasetDict。

2. dataset_info.json：这个文件存储了数据集的元数据，包括数据集的特征、版本信息等。如果这个文件存在，Hugging Face 会根据它来加载数据集的结构。

3. state.json 指定arrow文件和split切分等。

这是上面截图中的state.json文件内容：看一下它的具体内容

{
  "_data_files": [
    {
      "filename": "data-00000-of-00005.arrow"
    },
    {
      "filename": "data-00001-of-00005.arrow"
    },
    {
      "filename": "data-00002-of-00005.arrow"
    },
    {
      "filename": "data-00003-of-00005.arrow"
    },
    {
      "filename": "data-00004-of-00005.arrow"
    }
  ],
  "_fingerprint": "3f23ea1c8451dc7f",
  "_format_columns": null,
  "_format_kwargs": {},
  "_format_type": null,
  "_output_all_columns": false,
  "_split": "train"
}

GPT4o继续回答：
错误原因：报错提示 FileNotFoundError: Directory ~/.cache/huggingface/datasets/allenai___tulu-3-sft-mixture-1024-73k is neither a Dataset directory nor a DatasetDict directory，这意味着该目录下缺少 dataset_dict.json 或其他标识该目录为有效 Hugging Face 数据集目录的文件。

继续debug：

sampled_dataset.save_to_disk(output_dataset_path)
为什么保存的文件中没有dataset_info.json等信息，如何做才能有

GPT4o的回答：

Dataset.save_to_disk() 方法仅会保存数据集的实际内容（即数据文件），但它不会自动保存 dataset_info.json、state.json 等元数据文件。为了确保在保存采样后的数据集时包括这些元数据，你需要手动创建并保存 dataset_info.json 和 state.json 文件。以下是如何做到这一点。

解决方案：手动生成并保存 dataset_info.json和 state.json 文件

在你保存采样后的数据集时，你需要确保保存相关的元数据文件。你可以通过以下步骤：

1. 创建 dataset_info.json 文件：这可以通过复制原始数据集的元数据并修改其中的一些信息（例如样本数和文件大小）来完成。
2. 保存 state.json 文件：需要模仿样本的格式来完成。

你可以通过以下方式修改你的代码来保存这些元数据文件：（这里笔者和GPT反复对线，个中过程只有经历过的人才懂）

重新采样和保存 dataset_info.json和 state.json的核心代码

import os
import json
from datasets import Dataset, DatasetInfo
from pathlib import Path

# 原数据集路径
source_dataset_path = "~/.cache/huggingface/datasets/allenai___tulu-3-sft-mixture-1024/train"
# 新数据集保存路径
output_dataset_path = "~/.cache/huggingface/datasets/allenai___tulu-3-sft-mixture-1024-73k/train"

# 加载原数据集
dataset = Dataset.load_from_disk(source_dataset_path)

# 设置随机种子，确保每次运行结果相同
seed = 123
dataset = dataset.shuffle(seed=seed)

# 随机采样1/10的数据
sampled_dataset = dataset.select(range(len(dataset)//10))

# 输出采样后的数据集大小
print(f"采样后的数据集包含 {len(sampled_dataset)} 条数据")

# 保存采样后的数据集到新目录
sampled_dataset.save_to_disk(output_dataset_path)

# 获取原始数据集的信息
original_dataset_info = DatasetInfo.from_directory(source_dataset_path)

# 计算目录下所有 Arrow 文件的总大小
def get_total_size(directory_path):
    total_size = 0
    for file_path in Path(directory_path).rglob("*.arrow"):
        total_size += file_path.stat().st_size
    return total_size

# 获取新数据集的大小（所有 Arrow 文件的总大小）
total_size_in_bytes = get_total_size(output_dataset_path)

# 创建dataset_info.json
dataset_info = {
    "builder_name": original_dataset_info.builder_name,
    "config_name": original_dataset_info.config_name,
    "dataset_name": original_dataset_info.dataset_name,
    "dataset_size": original_dataset_info.dataset_size,
    "description": original_dataset_info.description,
    "features": original_dataset_info.features.to_dict(),
    "size_in_bytes": total_size_in_bytes,
    "splits": {
        "train": {
            "name": "train",
            "num_bytes": total_size_in_bytes,
            "num_examples": len(sampled_dataset),
            "shard_lengths": [len(sampled_dataset)],
            "dataset_name": original_dataset_info.dataset_name
        }
    },
    "version": {
        "version_str": f"{original_dataset_info.version.major}.{original_dataset_info.version.minor}.{original_dataset_info.version.patch}",
        "major": original_dataset_info.version.major,
        "minor": original_dataset_info.version.minor,
        "patch": original_dataset_info.version.patch
    }
}

# 保存dataset_info.json
dataset_info_path = os.path.join(output_dataset_path, "dataset_info.json")
with open(dataset_info_path, "w") as f:
    json.dump(dataset_info, f, indent=4)
print(f"dataset_info.json 已保存至: {dataset_info_path}")

# 创建state.json文件
state_json = {
    "_data_files": [
        {
            "filename": os.path.basename(file)
        }
        for file in sorted(Path(output_dataset_path).glob("*.arrow"))
    ],
    "_fingerprint": sampled_dataset._fingerprint,
    "_format_columns": None,
    "_format_kwargs": {},
    "_format_type": None,
    "_output_all_columns": False,
    "_split": "train"
}

# 保存state.json
state_json_path = os.path.join(output_dataset_path, "state.json")
with open(state_json_path, "w") as f:
    json.dump(state_json, f, indent=4)
print(f"state.json 已保存至: {state_json_path}")

代码说明

1. 加载原始数据集：首先，加载原始数据集并执行数据采样。

2. 保存采样后的数据集：保存采样后的数据集内容到 output_dataset_path。

3. 获取原始数据集的元信息：使用 DatasetInfo.from_directory() 方法获取原始数据集的元信息。通过这个信息，你可以得到原始数据集的结构、特征、分割等。

4. 创建并保存 dataset_info.json：

   • 创建 dataset_info.json 文件并填写相关信息。
   • 其中包括数据集的 features（字段结构）、splits（分割信息）、dataset_size（文件大小）等信息。

5. 创建并保存 state.json：

   • 创建 state.json 文件并包含必要信息。

6. 保存文件：将 dataset_info.json 和 state.json 文件保存到新数据集目录中。

重要注意事项

• features：这部分来自原始数据集的 DatasetInfo 对象，保留了字段和类型的信息。
• splits：你需要手动为采样后的数据集定义分割信息。采样了训练集，splits 部分需要包括 train 分割。
• size_in_bytes 和 num_bytes：这些值应基于采样后的数据集大小和文件大小进行更新。

通过这些操作，你就能确保采样后的数据集包括了 dataset_info.json 和 state.json 文件，从而能够正常加载和使用。

创建state.json文件的代码解释

# 创建state.json文件
state_json = {
    "_data_files": [
        {
            "filename": os.path.basename(file)
        }
        for file in sorted(Path(output_dataset_path).glob("*.arrow"))
    ],
    "_fingerprint": sampled_dataset._fingerprint,
    "_format_columns": None,
    "_format_kwargs": {},
    "_format_type": None,
    "_output_all_columns": False,
    "_split": "train"
}

# 保存state.json
state_json_path = os.path.join(output_dataset_path, "state.json")
with open(state_json_path, "w") as f:
    json.dump(state_json, f, indent=4)
print(f"state.json 已保存至: {state_json_path}")

说明

1. _data_files 字段：

   • 使用 Path(output_dataset_path).glob("*.arrow") 遍历目标目录下的所有 .arrow 文件，并将其文件名作为 filename 的值填入。
   • 文件名需要用 os.path.basename(file) 获取，因为文件路径不应该包含绝对路径。

2. _fingerprint 字段：

   • 使用 sampled_dataset._fingerprint 获取数据集的唯一指纹值，确保其与实际数据内容匹配。

3. _format_columns、_format_kwargs 和 _format_type 字段：

   • 设置为 None，与 Hugging Face 数据集的默认值保持一致。

4. _split 字段：

   • 设置为 "train"，因为你的数据集是训练集。

5. 排序：

   • 使用 sorted() 确保文件按照文件名顺序排列，避免 _data_files 顺序不一致。

运行上面的代码，可以得到想要的两个文件

state.json如下：

{
    "_data_files": [
        {
            "filename": "data-00000-of-00001.arrow"
        }
    ],
    "_fingerprint": "fda1608541b4fb67",
    "_format_columns": null,
    "_format_kwargs": {},
    "_format_type": null,
    "_output_all_columns": false,
    "_split": "train"
}

dataset_info.json如下：

{
    "builder_name": "arrow",
    "config_name": "default",
    "dataset_name": "55e9fd6d41c3cd1a98270dff07557bc2a1e1ba91",
    "dataset_size": 2914253735,
    "description": "",
    "features": {
        "id": {
            "dtype": "string",
            "_type": "Value"
        },
        "messages": [
            {
                "content": {
                    "dtype": "string",
                    "_type": "Value"
                },
                "role": {
                    "dtype": "string",
                    "_type": "Value"
                }
            }
        ],
        "source": {
            "dtype": "string",
            "_type": "Value"
        }
    },
    "size_in_bytes": 121605560,
    "splits": {
        "train": {
            "name": "train",
            "num_bytes": 121605560,
            "num_examples": 73752,
            "shard_lengths": [
                73752
            ],
            "dataset_name": "55e9fd6d41c3cd1a98270dff07557bc2a1e1ba91"
        }
    },
    "version": {
        "version_str": "0.0.0",
        "major": 0,
        "minor": 0,
        "patch": 0
    }
}

万事俱备，只欠东风

还需要如下图所示的dataset_dict.json文件，这个比较简单

对于我的情况，直接复制过来就行
dataset_dict.json内容如下：

{"splits": ["train"]}

重新load_from_disk正常运行

from datasets import load_dataset, load_from_disk

# 加载数据集
input_dir="~/.cache/huggingface/datasets/allenai___tulu-3-sft-mixture-1024-73k"
dataset = load_from_disk(input_dir)

# 指定保存路径
output_dir = "xxxx/data/tulu-3-sft-mixture-seqlen-1024-json-73k"

# 将数据集转换为 JSON 格式
for split in dataset.keys():  # 处理所有划分 (train, validation 等)
    dataset[split].to_json(f"{output_dir}/{split}.json", orient="records", lines=True)

生成的结果如下：

成功解决数据集采样之后只得到.arrow文件，但是无法load_from_disk的问题。

后记

2024年12月29日22点08分于上海， 在GPT4o大模型辅助下完成。