如何将两个 pandas 数据帧加载到 Hugginface 的数据集对象中？

Question

我正在尝试将训练和测试数据帧加载到数据集对象中。将 pandas 数据帧加载到数据集对象中的常用方法是：

from datasets import Dataset
import pandas as pd
df = pd.DataFrame({"a": [1, 2, 3]})
dataset = Dataset.from_pandas(df)

我的问题是如何将训练和测试两个 pandas 数据帧加载到数据集中？

例如，如果我有两个数据框：

from datasets import Dataset
import pandas as pd
df_train = pd.DataFrame({"a": [1, 2, 3]})
df_test = pd.DataFrame({"ab": [1, 2, 3]})

如何加载这两个框架？

Answer 1

您可以使用DatasetDict和将两个 pandas 数据帧加载到数据集中Dataset.from_pandas()。您可以通过向 DatasetDict 添加键来保持训练和测试的分离。

from datasets import Dataset, DatasetDict

import datasets
import pandas as pd


df_train = pd.DataFrame({"a": [1, 2, 3]})
df_test = pd.DataFrame({"ab": [1, 2, 3]})

datasets_train_test = DatasetDict({
    "train": Dataset.from_pandas(df_train),
    "test": Dataset.from_pandas(df_test)
    })

这导致

DatasetDict({
    train: Dataset({
        features: ['a'],
        num_rows: 3
    })
    test: Dataset({
        features: ['ab'],
        num_rows: 3
    })
})

您可以用来concatenate_datasets()连接数据集列表。

dataset_train = Dataset.from_pandas(df_train)
dataset_test = Dataset.from_pandas(df_test)
datasets_all_in_one = datasets.concatenate_datasets([dataset_train, dataset_test])

这导致

Dataset({
    features: ['a', 'ab'],
    num_rows: 6
})