如何在Hugging Face Dataset中指定单一特征为numpy数组？

使用load_dataset()加载Mozilla Common Voice（v11）数据集时，得到的数据集ds里audio.array是numpy数组，但不知道怎么复现这个效果。请问如何仅将单个特征设置为ndarray类型？

查看Common Voice数据集的示例：

>>> tt = load_dataset(
    data_args.dataset_name,
    data_args.dataset_config_name,
    split=f'{data_args.train_split_name}[:15%]',  # 只加载前15%数据
    cache_dir=model_args.cache_dir,
    token=model_args.token,
)
>>> type(tt.select([0])['audio'][0]['path'])
<class 'str'>
>>> type(tt.select([0])['audio'][0]['array'])
<class 'numpy.ndarray'>
>>> type(tt.select([0])['path'][0]) # 数据集里重复存了路径作为顶级特征
<class 'str'>

但在自定义代码里没法单独存储numpy数组，只发现用ds = ds.with_format('np')时，重新加载的数据集所有顶级特征都会变成numpy类型（测试结果如下）：

>>> type(test_ds['path'][0])
<class 'numpy.str_'>

需求是仅把“audio -> array”设为一维numpy数组。以下是创建并重新加载数据集以检查类型的测试代码：

#!/usr/bin/env python
# 尝试在huggingface dataset中保存和重新加载numpy数组
# 加载后的数组类型必须是numpy array()
from datasets import Dataset, Features, Array2D, Sequence, Value
import numpy as np

audio_arrays = [np.random.rand(16000), np.random.rand(16000)] 

features = Features({
  # 每个audio包含音频数据的np数组和源音频文件路径
  'audio': Sequence({
    #'array': Sequence(feature=Array2D(shape=(None,), dtype="float32")),
    'array': Sequence(feature=Value('float32')),
    'path': Value('string'),
  }),
  'path': Value('string'), # Common Voice数据集里也重复存了路径
})

ddata = {
    'path': [],        # 字符串列表
    'audio': [],       # 字典列表
}

ddata['path'] = ['/foo0/', '/bar0/'] # 确保能看到存储差异
ddata['audio'] = [
        {'array': audio_arrays[0], 'path': '/foo1/' },
        {'array': audio_arrays[1], 'path': '/bar1/', },
]
ds = Dataset.from_dict(ddata)
ds = ds.with_format('np')
ds.save_to_disk('/tmp/ds.ds') 

loaded_dataset = Dataset.load_from_disk('/tmp/ds.ds')
ld = loaded_dataset
au = ld['audio'][0]
ar = ld['audio'][0]['array']
print("Type of audio array:", type(ar))
print("Type of path:", type(ld['path'][0]))
print("Type of au path:", type(ld['audio'][0]['path']))
import ipdb; ipdb.set_trace(context=16); pass

要实现仅让audio.array保持numpy数组类型，其他特征保留原生Python类型，不需要全局设置with_format('np')，正确做法是在定义Features时用Array类型而非Sequence指定音频数组，并且无需全局格式化整个数据集。

步骤1：修正Features定义 #

使用Array类型（而非Sequence(Value('float32'))）定义audio.array，数据集会自动将其识别为numpy数组，同时其他特征保持原生类型：

from datasets import Dataset, Features, Array, Value

features = Features({
    'audio': {
        'array': Array(dtype='float32', shape=(None,)),  # 一维可变长度float32数组
        'path': Value('string'),
    },
    'path': Value('string'),
})

shape=(None,)表示数组长度可变，适配音频数据的特性。

步骤2：创建并保存数据集 #

无需调用with_format('np')，直接创建数据集并保存：

audio_arrays = [np.random.rand(16000), np.random.rand(16000)] 

ddata = {
    'path': ['/foo0/', '/bar0/'],
    'audio': [
        {'array': audio_arrays[0], 'path': '/foo1/' },
        {'array': audio_arrays[1], 'path': '/bar1/', },
    ]
}

ds = Dataset.from_dict(ddata, features=features)
ds.save_to_disk('/tmp/ds.ds') 

步骤3：加载并验证类型 #

重新加载后，audio.array会是numpy数组，而path等特征保持原生Python字符串类型：

loaded_dataset = Dataset.load_from_disk('/tmp/ds.ds')

print("Type of audio array:", type(loaded_dataset['audio'][0]['array']))  # <class 'numpy.ndarray'>
print("Type of top-level path:", type(loaded_dataset['path'][0]))  # <class 'str'>
print("Type of audio path:", type(loaded_dataset['audio'][0]['path']))  # <class 'str'>

原理说明 #

Common Voice数据集能自动让audio.array成为numpy数组，核心是它的特征定义使用了Array类型描述音频数组，而非通用的Sequence。全局with_format('np')会强制所有特征转为numpy类型，而精准定义特征类型，可以实现仅指定单个字段为numpy数组，其他字段保持原生类型。

内容的提问来源于stack exchange，提问作者Jaggz