HuggingFace 数据集：加载具有不同列集的数据集

Question

这就是我使用 HF 加载训练和测试数据集的方式：

dataset = {name.replace('/', '.'): f'{name}/*.parquet' for name in ["train", "test"]}
dataset = load_dataset("parquet", data_files=dataset)

但是，存在一个问题，因为我的

train

数据集具有

test

没有的额外列（并且我需要加载它们才能准备训练数据。

由于存在列不匹配，因此我在加载时收到错误：

datasets.table.CastError: Couldn't cast
column1: string
column2: list<item: string>
  child 0, item: string
column3: list<item: int32>
  child 0, item: int32
__index_level_0__: int64
-- schema metadata --
pandas: '{"index_columns": ["__index_level_0__"], "column_indexes": [{"na' + 1438
to
{'column1': Value(dtype='string', id=None), 'column2': Sequence(feature=Value(dtype='string', id=None), length=-1, id=None), '__index_level_0__': Value(dtype='int64', id=None)}

因为列名称不匹配（缺少

column3

）。

使用不同的列组加载训练和测试的正确方法是什么？

Answer 1

from datasets import load_dataset

def load_datasets_with_varying_columns(train_path, test_path):

  train_dataset = load_dataset("parquet", data_files={"train": train_path})
  test_dataset = load_dataset("parquet", data_files={"test": test_path})

  common_columns = set(train_dataset["train"].column_names).intersection(set(test_dataset["test"].column_names))

  test_dataset = test_dataset.map(lambda x: {col: x[col] for col in common_columns})

  return {"train": train_dataset["train"], "test": test_dataset["test"]}

train_path = "train/*.parquet"
test_path = "test/*.parquet"
datasets = load_datasets_with_varying_columns(train_path, test_path)

希望这有帮助