huggingface / huggingface/datasets
Tensor type (e.g. from `return_tensors`) ignored in map
- Dominant language
- Python
- Stars
- 22k
- Forks
- 3.4k
- Avg merge
- 5d 7h
- Merged PRs (30d)
- 17
Description
### Describe the bug
I don't know if it is a bug or an expected behavior, but the tensor type seems to be ignored after applying map. For example, mapping over to tokenize text with a transformers' tokenizer always returns lists and it ignore the `return_tensors` argument.
If this is an expected behaviour (e.g., for caching/Arrow compatibility/etc.) it should be clearly documented. For example, current documentation (see [here](https://huggingface.co/docs/datasets/v2.17.1/en/nlp_process#map)) clearly state to "set `return_tensors="np"` when you tokenize your text" to have Numpy arrays.
### Steps to reproduce the bug
```py
# %%%
import datasets
import numpy as np
import tensorflow as tf
import torch
from transformers import AutoTokenizer
# %%
ds = datasets.load_dataset("cnn_dailymail", "1.0.0", split="train[:1%]")
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
#%%
for return_tensors in [None, "np", "pt", "tf", "jax"]:
print(f"********** no map, return_tensors={return_tensors} **********")
_ds = tokenizer(ds["article"], return_tensors=return_tensors, truncation=True, padding=True)
print('Type :', type(_ds["input_ids"]))
# %%
for return_tensors in [None, "np", "pt", "tf", "jax"]:
print(f"********** map, return_tensors={return_tensors} **********")
_ds = ds.map(
lambda examples: tokenizer(examples["article"], return_tensors=return_tensors, truncation=True, padding=True),
batched=True,
remove_columns=["article"],
)
print('Type :', type(_ds[0]["input_ids"]))
```
### Expected behavior
The output from the script above. I would expect the second half to be the same.
```
********** no map, return_tensors=None **********
Type :
********** no map, return_tensors=np **********
Type :
********** no map, return_tensors=pt **********
Type :
********** no map, return_tensors=tf **********
Type :
********** no map, return_tensors=jax **********
Type :
********** map, return_tensors=None **********
Type :
********** map, return_tensors=np **********
Type :
********** map, return_tensors=pt **********
Type :
********** map, return_tensors=tf **********
Type :
********** map, return_tensors=jax **********
Type :
```
### Environment info
- `datasets` version: 2.17.1
- Platform: Redacted (linux)
- Python version: 3.10.12
- `huggingface_hub` version: 0.20.3
- PyArrow version: 15.0.0
- Pandas version: 2.1.3
- `fsspec` version: 2023.10.0
Contributor guide
Assessment
This issue has not been assessed yet.