Issue: Return dict structure if columns are lists - to_tf_dataset
State: open
Author: amyeroberts

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)}>
```

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/datasets/pr_5602). All of your documentation changes will be reflected on that endpoint.

This is a great PR! Thinking about the UX though, maybe we could do it without the extra argument? Before this PR, the logic in `to_tf_dataset` was that if the user passed a single column name in either `columns` or `label_cols`, we converted it to a length-1 list. Then, later in the code, we convert output dicts with only one key to naked Tensors.

Would it be easier if we removed the argument, but instead treated the cases differently? Passing a column name as a string could yield a single naked Tensor in the output as before, but passing a list of length 1 would yield a full dict? That way if you wanted dict output with a single key you could just say `columns=[col_name]`.

(I'm not totally convinced this is a good idea yet, it just seems like it might be more intuitive)

@Rocketknight1  Happy to implement it that way - it's certainly cleaner to not have another arg. In this case, am I right in saying we'd effectively set `return_dict` [here](https://github.com/huggingface/datasets/blob/6569014a9948eab7d031a3587405e64ba92d6c59/src/datasets/arrow_dataset.py#L410) - where columns are made into a list if they were a string? 

There only concern I have is this changes the default behaviour, which might break things for people who were happily using `columns=["my_col_str"]` before. 




@amyeroberts That's correct! Probably the simplest way to implement it would be to just add the flag there.

And yeah, I'm aware this might be a slightly breaking change, but we've mostly tried to move users to `prepare_tf_dataset` in `transformers` at this point, so hopefully as long as that method doesn't break then most users won't be negatively affected by the change.

@lhoestq @Rocketknight1  - I've remove the `return_dict` argument and implemented @Rocketknight1 's suggestion. LMK what you think :) 

@lhoestq Of course :) I've opened a draft PR here for the updates needed in transformers examples and docs to keep the returned data structure consistent: https://github.com/huggingface/transformers/pull/21935. Note: even with the different structure, `model.fit` can still successfully be called. 

For the [link you shared](https://github.com/huggingface/datasets/pull/url) - for me it returns a 404 error. Is there another link I could follow to see how to run the transformers CI with this branch? 

Currently looking into the failing tests 😭 

Oh sorry - I fixed the URL: https://github.com/huggingface/transformers/commit/4eb55bbd593adf2e49362613ee32a11ddc4a854d

The error shows `There appear to be 80 leaked shared_memory objects to clean up at shutdown`. IIRC to_tf_dataset does some shared memory stuff for multiprocessing - maybe @Rocketknight1 you know what's going on ?

@lhoestq That warning appears anytime you interrupt a process using Python `SharedMemory` objects - it's only a problem if you still get the error when the process finishes normally! Our implementation of `to_tf_dataset` should clean things up properly.

Ok, not sure why it fails then :/

Hmm, will investigate! Sorry, I misread - I thought that warning was coming up in the context of another error

IMO outputing different types based on nuances in the input could confuse users.

Also, in the ideal scenario,`to_tf_function` should return a `tf.data.Dataset` that iterates over the underlying Arrow data and yields (unprocessed) dicts of TF tensors, and all the model-specific code should live in Transformers (e.g., in `prepare_tf_dataset`). So the goal would be to make `to_tf_dataset` more user-friendly, not more complex :).

I think we agree @mariosasko  :) 

> Also, in the ideal scenario,to_tf_function should return a tf.data.Dataset that iterates over the underlying Arrow data and yields (unprocessed) dicts of TF tensors

This I'll leave for another PR as it's outside the scope of this one and @Rocketknight1 will have far more knowledge and ideas about what is possible

> all the model-specific code should live in Transformers (e.g., in prepare_tf_dataset

Agreed! This PR isn't really a model specific change - although it was highlighted when trying to train a model. We definitely want to move model specific things out of datasets as much as possible. 

> IMO outputing different types based on nuances in the input could confuse users.
> So the goal would be to make to_tf_dataset more user-friendly, not more complex :).

The aim was to move more towards being able to return the dict of TF tensors you suggest, whilst maintaining backwards compatibility. Personally, I found it surprising to be returned a tuple structure when I was using `to_tf_dataset`. The aim was to make `to_tf_dataset` more user friendly, but I agree that it has the potential to be confusing. 

For context, the thought process behind this design was to: 
* Not add even more arguments to `to_tf_dataset`. 
* Have a feature selection -> return type logic in keeping with `datasets` e.g. `dataset['train'][:10]['feat1']` returns a list of values, whereas `dataset['train'][:10]['feat1', 'feat2']` returns a dictionary. 

Very happy to add any suggestions or changes you might have about how to make this design better! :) 


Hi ! Anything blocking here ? I'b be happy to help

Hi @lhoestq - sorry this hasn't been very active for the past ~1.5 weeks. There's nothing specific blocking, other than not being able to replicate without running on CI, and still need to test a bit more to narrow down the issue. I should have time tomorrow to pick it up again :) 

@lhoestq @Rocketknight1 Friendly ping for a review :) 

Awesome ! What about showing a warning that this change is about to happen in the next version of `datasets`, and then apply this change in a subsequent major release ? This way folks at twitter won't hate us: https://github.com/twitter/the-algorithm/blob/138bb519975407d4ea0dc1478d897d451ef05dab/trust_and_safety_models/toxicity/data/mb_generator.py#L142-L148

@lhoestq Sounds good! How would you like this warning to happen? I could open a PR to add a warning message within `to_tf_dataset`?

Yup sounds good :)