如何通过huggingface datasets添加行来扩增数据集？

Question

我有一个数据集，里面有113287行训练数据。不过，每个'caption'字段其实是一个包含多个字符串的数组。我想把这个数组展开，添加新的行。

文档里提到，数据集的批量映射功能可以用来实现这个目标：

这意味着你可以把你的例子连接起来，分开处理，甚至可以添加更多的例子！

from datasets import load_dataset

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

def chunk_examples(entry):
    captions = [caption for caption in entry["caption"][0]]
    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']))

DatasetDict({
    train: Dataset({
        features: ['caption', 'image_id', 'image'],
        num_rows: 113287
    })
    validation: Dataset({
        features: ['caption', 'image_id', 'image'],
        num_rows: 5000
    })
    test: Dataset({
        features: ['caption', 'image_id', 'image'],
        num_rows: 5000
    })
})
DatasetDict({
    train: Dataset({
        features: ['caption'],
        num_rows: 464
    })
    validation: Dataset({
        features: ['caption'],
        num_rows: 40
    })
    test: Dataset({
        features: ['caption'],
        num_rows: 40
    })
})
464

我遇到的问题是，生成的数据集行数并没有我预期的那么多。

它显示添加了num_rows: 464行。我怀疑这可能是因为批次的原因。我该如何把它恢复成一个“正常”的数据集呢？我的映射函数是不是哪里出错了？

• datasets==2.18.0

Answer 1

我的映射函数写错了。我只访问了第一个条目，用的是 [0]。

def chunk_examples(batch):
    captions = []
    for row in batch["caption"]:
        captions += row
    return {"caption": captions}

现在它的结果是：

DatasetDict({
    train: Dataset({
        features: ['caption'],
        num_rows: 453460
    })
    validation: Dataset({
        features: ['caption'],
        num_rows: 25010
    })
    test: Dataset({
        features: ['caption'],
        num_rows: 25010
    })
})
453460