如何将HuggingFace DatasetDict中的数据导出为Pandas DataFrame？

我从HuggingFace下载了amazon_reviews_multi数据集，使用datasets.Dataset.load_dataset()方法加载后，得到一个基于Apache Arrow表的DatasetDict对象。现在需要把其中的train、validation、test三个子集分别导出为独立的Pandas DataFrame。

该DatasetDict的结构如下：

DatasetDict({
    train: Dataset({
        features: ['review_id', 'product_id', 'reviewer_id', 'stars', 'review_body', 'review_title', 'language', 'product_category'],
        num_rows: 1200000
    })
    validation: Dataset({
        features: ['review_id', 'product_id', 'reviewer_id', 'stars', 'review_body', 'review_title', 'language', 'product_category'],
        num_rows: 30000
    })
    test: Dataset({
        features: ['review_id', 'product_id', 'reviewer_id', 'stars', 'review_body', 'review_title', 'language', 'product_category'],
        num_rows: 30000
    })
})

其中dataset['train'].features的信息为：

{'review_id': Value(dtype='string', id=None),
 'product_id': Value(dtype='string', id=None),
 'reviewer_id': Value(dtype='string', id=None),
 'stars': Value(dtype='int32', id=None),
 'review_body': Value(dtype='string', id=None),
 'review_title': Value(dtype='string', id=None),
 'language': Value(dtype='string', id=None),
 'product_category': Value(dtype='string', id=None)}

解决方案 #

Hugging Face Datasets库的Dataset对象内置了to_pandas()方法，可直接将Arrow格式数据集转换为Pandas DataFrame。针对每个子集调用该方法即可完成导出：

import datasets
import pandas as pd

# 加载数据集（若未完成加载执行此步骤）
dataset = datasets.load_dataset("amazon_reviews_multi")

# 分别导出三个子集为独立DataFrame
train_df = dataset['train'].to_pandas()
validation_df = dataset['validation'].to_pandas()
test_df = dataset['test'].to_pandas()

验证转换结果 #

可通过以下方式确认转换成功：

• 查看对象类型：print(type(train_df))，输出应为 <class 'pandas.core.frame.DataFrame'>
• 查看数据前几行：print(train_df.head())，会显示数据集的前5条记录

内容的提问来源于stack exchange，提问作者M.og.op.gpt