huggingface / huggingface/datasets

Tensor type (e.g. from `return_tensors`) ignored in map

Open
#6,688 3 comments 0 reactions 0 assignees View on GitHub
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

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.