使用Hugging Face Datasets展开caption数组扩充数据集时行数异常排查

问题根源 #

你的映射函数存在错误：当batched=True时，传入函数的entry是批次级别的数据字典，其中entry["caption"]是包含当前批次所有样本caption数组的列表。而你只处理了批次中的第一个样本（entry["caption"][0]），丢弃了批次内其他所有样本，最终导致生成的总行数远低于预期。

修正后的映射函数 #

根据需求（将每个样本的caption数组展开为独立行），需要遍历批次内的所有样本，将每个样本的caption数组逐个展开并合并为一个大列表：

from datasets import load_dataset

dataset_name = "Jotschi/coco-karpathy-opus-de"
coco_dataset = load_dataset(dataset_name)

def chunk_examples(entry):
    # 遍历批次中每个样本的caption数组，展开所有字符串
    captions = [cap for caps in entry["caption"] for cap in caps]
    return {"caption": captions}

print(coco_dataset)
chunked_dataset = coco_dataset.map(
    chunk_examples,
    batched=True,
    num_proc=4,
    remove_columns=["image_id", "caption", "image"]
)
print(chunked_dataset)
print(len(chunked_dataset['train']))

特殊情况处理 #

如果数据集里每个样本的caption字段是嵌套数组（比如格式为[[cap1, cap2, cap3]]），则需要先取出内层数组再展开：

def chunk_examples(entry):
    captions = [cap for caps in entry["caption"] for cap in caps[0]]
    return {"caption": captions}

验证结果 #

修正后，训练集的行数应该等于所有样本的caption字符串总数（比如原数据集每条样本平均有4个caption，总行数会接近113287 * 4 = 453148），符合预期。

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