def replace_ids(inputs):
   inputs["label"]["ids"] = [label2id[x] for x in inputs["label"]["ids"]]
   return inputs
dataset = dataset.map(replace_ids)
flat_dataset = dataset.flatten()
flat_dataset.features
{'question': Value(dtype='string', id=None),
 'image_id': Value(dtype='string', id=None),
 'label.ids': Sequence(feature=Value(dtype='int64', id=None), length=-1, id=None),
 'label.weights': Sequence(feature=Value(dtype='float64', id=None), length=-1, id=None)}