向 HuggingFace 数据集添加新列

Posted

技术标签:

【中文标题】向 HuggingFace 数据集添加新列【英文标题】:Add new column to a HuggingFace dataset 【发布时间】:2021-12-31 23:44:00 【问题描述】:

在我有 5000000 行的数据集中,我想在我的数据集中添加一个名为“嵌入”的列。

dataset = dataset.add_column('embeddings', embeddings)

变量 embeddings 是一个大小为 (5000000, 512) 的 numpy memmap 数组。

但我收到此错误:

ArrowInvalidTraceback(最后一次调用) 在 ----> 1 个数据集 = dataset.add_column('embeddings', embeddings)

/opt/conda/lib/python3.8/site-packages/datasets/arrow_dataset.py 在 wrapper(*args, **kwargs) 第486章 第487章 --> 488 out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs) 489 个数据集: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out] 490 # 重新应用格式到输出

/opt/conda/lib/python3.8/site-packages/datasets/fingerprint.py in wrapper(*args, **kwargs) 404 # 调用实际函数 405 --> 406 out = func(self, *args, **kwargs) 407 408 # 更新原地转换指纹+更新原地转换历史

/opt/conda/lib/python3.8/site-packages/datasets/arrow_dataset.py in add_column(self, name, column, new_fingerprint) 3346:类:Dataset 第3347章 -> 3348 column_table = InMemoryTable.from_pydict(name: column) 3349 # 水平连接表 3350 table = ConcatenationTable.from_tables([self._data, column_table],axis=1)

/opt/conda/lib/python3.8/site-packages/datasets/table.py in from_pydict(cls, *args, **kwargs) 第367章 第368章 --> 369 返回 cls(pa.Table.from_pydict(*args, **kwargs)) 370 第371章

/opt/conda/lib/python3.8/site-packages/pyarrow/table.pxi in pyarrow.lib.Table.from_pydict()

/opt/conda/lib/python3.8/site-packages/pyarrow/table.pxi in pyarrow.lib._from_pydict()

/opt/conda/lib/python3.8/site-packages/pyarrow/array.pxi in pyarrow.lib.asarray()

/opt/conda/lib/python3.8/site-packages/pyarrow/array.pxi in pyarrow.lib.array()

/opt/conda/lib/python3.8/site-packages/pyarrow/array.pxi in pyarrow.lib._ndarray_to_array()

/opt/conda/lib/python3.8/site-packages/pyarrow/error.pxi in pyarrow.lib.check_status()

ArrowInvalid:只处理一维数组

由于嵌入数组不适合 RAM,我该如何以一种有效的方式解决?

【问题讨论】:

【参考方案1】:
from datasets import load_dataset

ds = load_dataset("cosmos_qa", split="train")

new_column = ["foo"] * len(ds)
ds = ds.add_column("new_column", new_column)

你得到一个数据集

Dataset(
    features: ['id', 'context', 'question', 'answer0', 'answer1', 'answer2', 'answer3', 'label', 'new_column'],
    num_rows: 25262
)

【讨论】:

【参考方案2】:

这里的问题是您正在尝试添加一列,但您传递的数据是一个 2d numpy 数组。 arrow(用来表示数据集的库)只支持一维numpy数组。

您可以尝试将 2d numpy 数组的每一列一一添加:

for i, column in enumerate(embeddings.T):
    ds = ds.add_column('embeddings_' + str(i), column)

由于嵌入数组不适合 RAM,我该如何以一种有效的方式解决?

我认为没有办法解决内存问题。 huggingface 数据集由箭头表支持,它必须适合内存。

【讨论】:

以上是关于向 HuggingFace 数据集添加新列的主要内容,如果未能解决你的问题,请参考以下文章

Transformers学习笔记2. HuggingFace数据集Datasets

Transformers学习笔记2. HuggingFace数据集Datasets

如何在 Huggingface 中从 CSV 加载自定义数据集

根据其他列向 Panda 数据框添加新列

KeyError:使用 Huggingface Transformers 使用 BioASQ 数据集时出现“答案”错误

如何将具有值的列添加到 Spark Java 中的新数据集?