Huggingface数据集采样之后得到的arrow文件无法用Dataset.load_from

问题起源

我现在有数据集

在这里插入图片描述

我想要采样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_diskraise 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 目录。主要通过以下几个文件来进行判断：

dataset_dict.json：这个文件是用于 DatasetDict 类型的数据集。DatasetDict 是一个包含多个数据集分割（如 train、validation 等）的字典。如果这个文件存在，Hugging Face 会认为该目录是一个 DatasetDict。
dataset_info.json：这个文件存储了数据集的元数据，包括数据集的特征、版本信息等。如果这个文件存在，Hugging Face 会根据它来加载数据集的结构。
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` 文件

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

创建 dataset_info.json 文件：这可以通过复制原始数据集的元数据并修改其中的一些信息（例如样本数和文件大小）来完成。
保存 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)# 创建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": os.path.getsize(output_dataset_path),"splits": {"train": {"name": "train","num_bytes": os.path.getsize(output_dataset_path),"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}")

代码说明

加载原始数据集：首先，加载原始数据集并执行数据采样。
保存采样后的数据集：保存采样后的数据集内容到 output_dataset_path。
获取原始数据集的元信息：使用 DatasetInfo.from_directory() 方法获取原始数据集的元信息。通过这个信息，你可以得到原始数据集的结构、特征、分割等。
创建并保存 dataset_info.json：
- 创建 dataset_info.json 文件并填写相关信息。
- 其中包括数据集的 features（字段结构）、splits（分割信息）、dataset_size（文件大小）等信息。
创建并保存 state.json：
- 创建 state.json 文件并包含必要信息。
保存文件：将 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}")

说明

_data_files 字段：
- 使用 Path(output_dataset_path).glob("*.arrow") 遍历目标目录下的所有 .arrow 文件，并将其文件名作为 filename 的值填入。
- 文件名需要用 os.path.basename(file) 获取，因为文件路径不应该包含绝对路径。
_fingerprint 字段：
- 使用 sampled_dataset._fingerprint 获取数据集的唯一指纹值，确保其与实际数据内容匹配。
_format_columns、_format_kwargs 和 _format_type 字段：
- 设置为 None，与 Hugging Face 数据集的默认值保持一致。
_split 字段：
- 设置为 "train"，因为你的数据集是训练集。
排序：
- 使用 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": 4096,"splits": {"train": {"name": "train","num_bytes": 4096,"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)