将 pandas 数据框转换为 datasetDict

Question

我无法在任何地方找到如何将 pandas 数据帧转换为 type datasets.dataset_dict.DatasetDict，以便在带有 Huggingface 模型的 BERT 工作流程中得到最佳使用。以这些简单的数据框为例。

train_df = pd.DataFrame({
     "label" : [1, 2, 3],
     "text" : ["apple", "pear", "strawberry"]
})

test_df = pd.DataFrame({
     "label" : [2, 2, 1],
     "text" : ["banana", "pear", "apple"]
})

将它们转换为上述类型的最有效方法是什么？

Answer 1

一种可能性是首先创建两个数据集，然后连接它们：

import datasets
import pandas as pd


train_df = pd.DataFrame({
     "label" : [1, 2, 3],
     "text" : ["apple", "pear", "strawberry"]
})

test_df = pd.DataFrame({
     "label" : [2, 2, 1],
     "text" : ["banana", "pear", "apple"]
})

train_dataset = Dataset.from_dict(train_df)
test_dataset = Dataset.from_dict(test_df)
my_dataset_dict = datasets.DatasetDict({"train":train_dataset,"test":test_dataset})

结果是：

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