如何将train、validation、test三个DataFrame构建为DatasetDict

我有train、validation、test三个DataFrame，想要把它们构建成如下格式的DatasetDict（来自🤗 Hugging Face Datasets库）：

DatasetDict({
    train: Dataset({
        features: ['ID', 'Tweet', 'anger', 'anticipation', 'disgust', 'fear', 'joy', 'love', 'optimism', 'pessimism', 'sadness', 'surprise', 'trust'],
        num_rows: 6838
    })
    test: Dataset({
        features: ['ID', 'Tweet', 'anger', 'anticipation', 'disgust', 'fear', 'joy', 'love', 'optimism', 'pessimism', 'sadness', 'surprise', 'trust'],
        num_rows: 3259
    })
    validation: Dataset({
        features: ['ID', 'Tweet', 'anger', 'anticipation', 'disgust', 'fear', 'joy', 'love', 'optimism', 'pessimism', 'sadness', 'surprise', 'trust'],
        num_rows: 886
    })
})

其中features对应DataFrame的列名。我尝试了下面的代码但无法正常运行，请问该怎么实现这个转换？

DatasetDict = {}

dataframes = [train, validation, test]
for grp in dataframes:
    
    DatasetDict[grp] = df

你这段代码的问题在于：

• 直接把DataFrame对象作为字典键，不符合DatasetDict的要求（它需要字符串键，比如"train"）
• 循环里的df是未定义变量，应该用当前遍历的grp
• 没有将DataFrame转换为Hugging Face的Dataset对象

正确实现方式如下：

1. 导入依赖库 #

确保已安装并导入Hugging Face的datasets库：

from datasets import Dataset, DatasetDict

2. 转换并构建DatasetDict #

直接写法 #

把每个DataFrame转换成Dataset，再用字符串键组装成DatasetDict：

# 将DataFrame转为Dataset
train_dataset = Dataset.from_pandas(train)
val_dataset = Dataset.from_pandas(validation)
test_dataset = Dataset.from_pandas(test)

# 构建目标格式的DatasetDict
dataset_dict = DatasetDict({
    "train": train_dataset,
    "validation": val_dataset,
    "test": test_dataset
})

循环简化写法 #

如果想用循环批量处理，可将DataFrame与对应名称配对遍历：

dataset_dict = DatasetDict()
# 配对数据集名称和对应的DataFrame
data_pairs = [("train", train), ("validation", validation), ("test", test)]

for name, df in data_pairs:
    dataset_dict[name] = Dataset.from_pandas(df)

运行上述代码后，打印dataset_dict就能得到你需要的格式，它会自动识别DataFrame的列作为features，并统计行数num_rows。

内容的提问来源于stack exchange，提问作者Katty_one