向 HuggingFace 数据集添加新列

Question

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

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

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

但我收到这个错误：

ArrowInvalidTraceback（最近一次调用最后一次）位于 ----> 1 dataset = dataset.add_column('embeddings', embeddings)

/opt/conda/lib/python3.8/site-packages/datasets/arrow_dataset.py inwrapper(*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 inwrapper(*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 :class: 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 @classmethod 368 def from_pydict(cls, *args, **kwargs): - -> 369 返回 cls(pa.Table.from_pydict(*args, **kwargs)) 370 371 @inject_arrow_table_documentation(pa.Table.from_batches)

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

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

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

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

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

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

ArrowInvalid：仅处理一维数组

由于嵌入数组不适合 RAM，我该如何解决（可能以有效的方式）？

Answer 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
})