GitHub issue #5602: Return dict structure if columns are lists - to_tf_dataset
State: open | Created: 2023-03-02T15:51:12Z

This PR introduces new logic to `to_tf_dataset` affecting the returned data structure, enabling a dictionary structure to be returned, even if only one feature column is selected. 

If the passed in `columns` or `label_cols` to `to_tf_dataset` are a list, they are returned as a dictionary, respectively. If they are a string, the tensor is returned. 

An outline of the behaviour:
```
dataset,to_tf_dataset(columns=["col_1"], label_cols="col_2")
# ({'col_1': col_1}, col_2}

dataset,to_tf_dataset(columns="col1", label_cols="col_2")
# (col1, col2)

dataset,to_tf_dataset(columns="col1")
# col1

dataset,to_tf_dataset(columns=["col_1"], labels=["col_2"])
# ({'col1': tensor}, {'col2': tensor}}

dataset,to_tf_dataset(columns="col_1", labels=["col_2"])
# (col1, {'col2': tensor}}
```

## Motivation

Currently, when calling `to_tf_dataset`, the returned dataset will always return a tuple structure if a single feature column is used. This can cause issues when calling `model.fit` on models which train without labels e.g. [TFVitMAEForPreTraining](https://github.com/huggingface/transformers/blob/b6f47b539377ac1fd845c7adb4ccaa5eb514e126/src/transformers/models/vit_mae/modeling_vit_mae.py#L849).  Specifically, [this line](https://github.com/huggingface/transformers/blob/d9e28d91a8b2d09b51a33155d3a03ad9fcfcbd1f/src/transformers/modeling_tf_utils.py#L1521) where it's assumed the input `x` is a dictionary if there is no label. 


## Example

Previous behaviour
```python
In [1]: import tensorflow as tf
   ...: from datasets import load_dataset
   ...: 
   ...: 
   ...: def transform(batch):
   ...:     def _transform_img(img):
   ...:         img = img.convert("RGB")
   ...:         img = tf.keras.utils.img_to_array(img)
   ...:         img = tf.image.resize(img, (224, 224))
   ...:         img /= 255.0
   ...:         img = tf.transpose(img, perm=[2, 0, 1])
   ...:         return img
   ...:     batch['pixel_values'] = [_transform_img(pil_img) for pil_img in batch['img']]
   ...:     return batch
   ...: 
   ...: 
   ...: def collate_fn(examples):
   ...:     pixel_values = tf.stack([example["pixel_values"] for example in examples])
   ...:     return {"pixel_values": pixel_values}
   ...: 
   ...: 
   ...: dataset = load_dataset('cifar10')['train']
   ...: dataset = dataset.with_transform(transform)
   ...: dataset.to_tf_dataset(batch_size=8, columns=['pixel_values'], collate_fn=collate_fn)

Out[1]: <PrefetchDataset element_spec=TensorSpec(shape=(None, 3, 224, 224), dtype=tf.float32, name=None)>
```

New behaviour
```python
In [1]: import tensorflow as tf
   ...: from datasets import load_dataset
   ...: 
   ...: 
   ...: def transform(batch):
   ...:     def _transform_img(img):
   ...:         img = img.convert("RGB")
   ...:         img = tf.keras.utils.img_to_array(img)
   ...:         img = tf.image.resize(img, (224, 224))
   ...:         img /= 255.0
   ...:         img = tf.transpose(img, perm=[2, 0, 1])
   ...:         return img
   ...:     batch['pixel_values'] = [_transform_img(pil_img) for pil_img in batch['img']]
   ...:     return batch
   ...: 
   ...: 
   ...: def collate_fn(examples):
   ...:     pixel_values = tf.stack([example["pixel_values"] for example in examples])
   ...:     return {"pixel_values": pixel_values}
   ...: 
   ...: 
   ...: dataset = load_dataset('cifar10')['train']
   ...: dataset = dataset.with_transform(transform)
   ...: dataset.to_tf_dataset(batch_size=8, columns=['pixel_values'], collate_fn=collate_fn)

Out[1]: <PrefetchDataset element_spec={'pixel_values': TensorSpec(shape=(None, 3, 224, 224), dtype=tf.float32, name=None)}>
```