Huggingface Dataset Map将列表转为NumPy数组，如何保留列表格式？

问题重现 #

我编写了一个生成嵌入向量列表的函数，通过map映射到数据集后，特意将Tensor转为Python列表，但转为DataFrame或CSV时，列表却变成了NumPy数组：

代码实现 #

def fetch_embedding(data):
    text = data["text"]
    out = trainer.predict(text)
    embeddings = out[0][1][-1][:,0,:]
    embeddings = embeddings.tolist()  # 手动转为Python列表
    return {"embeddings" : embeddings}

映射到数据集：

dataset = dataset.map(fetch_embedding)

检查类型：

df = dataset.to_pandas()
A = df.iloc[0].loc["embeddings"]
print(type(A))

输出结果：

<class 'numpy.ndarray'>

解决方法 #

原因分析 #

🤗 Datasets库会自动推断特征类型，当检测到嵌套的数值列表时，会将其标记为Array类型而非List类型，转为Pandas DataFrame时就会自动解析为NumPy数组。

方案1：手动指定特征类型为List #

在map操作后，显式将embeddings列的类型改为Python列表类型：

from datasets import Sequence, Value

# 定义特征类型：嵌套的float32列表
dataset = dataset.cast_column("embeddings", Sequence(feature=Value(dtype="float32")))

# 再转为DataFrame
df = dataset.to_pandas()
print(type(df.iloc[0]["embeddings"]))  # 输出 <class 'list'>

方案2：在map时强制保留Python列表类型 #

提前指定数据集特征，避免自动推断将列表识别为数组类型：

from datasets import Sequence, Value

# 复制原有特征并修改embeddings的类型定义
features = dataset.features.copy()
features["embeddings"] = Sequence(feature=Value(dtype="float32"))

# 映射时指定features参数
dataset = dataset.map(fetch_embedding, features=features)

# 转为DataFrame后检查类型
df = dataset.to_pandas()
print(type(df.iloc[0]["embeddings"]))

方案3：转为DataFrame后手动转回列表 #

如果不想修改数据集特征，也可以在转为DataFrame后批量将NumPy数组转回列表：

import numpy as np

df = dataset.to_pandas()
df["embeddings"] = df["embeddings"].apply(lambda x: x.tolist() if isinstance(x, np.ndarray) else x)
print(type(df.iloc[0]["embeddings"]))

内容的提问来源于stack exchange，提问作者Ricardo Muñoz Sánchez