Label casting during `map` process is canceled after the `map` process
Opened by yoshitomo-matsubara (closed)

### Describe the bug

When preprocessing a multi-label dataset, I introduced a step to convert int labels to float labels as [BCEWithLogitsLoss](https://pytorch.org/docs/stable/generated/torch.nn.BCEWithLogitsLoss.html) expects float labels and forward function of models in transformers package internally use `BCEWithLogitsLoss`

However, the casting was canceled after `.map` process and the label values still use int values, which leads to an error
```
  File "/home/yoshitomo/anaconda3/envs/torchdistill/lib/python3.10/site-packages/transformers/models/bert/modeling_bert.py", line 1711, in forward
    loss = loss_fct(logits, labels)
  File "/home/yoshitomo/anaconda3/envs/torchdistill/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/home/yoshitomo/anaconda3/envs/torchdistill/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
  File "/home/yoshitomo/anaconda3/envs/torchdistill/lib/python3.10/site-packages/torch/nn/modules/loss.py", line 819, in forward
    return F.binary_cross_entropy_with_logits(
  File "/home/yoshitomo/anaconda3/envs/torchdistill/lib/python3.10/site-packages/torch/nn/functional.py", line 3628, in binary_cross_entropy_with_logits
    return torch.binary_cross_entropy_with_logits(
RuntimeError: result type Float can't be cast to the desired output type Long
```

This seems like happening only when the original labels are int values (see examples below)

### Steps to reproduce the bug

If the original dataset uses a list of int labels, it will cancel the int->float casting

```python
from datasets import Dataset

data = {
    'text': ['text1', 'text2', 'text3', 'text4'],
    'labels': [[0, 1, 2], [3], [3, 4], [3]]
}

dataset = Dataset.from_dict(data)

label_set = set([label for labels in data['labels'] for label in labels])
label2idx = {label: idx for idx, label in enumerate(sorted(label_set))}


def multi_labels_to_ids(labels):
    ids = [0.0] * len(label2idx)
    for label in labels:
        ids[label2idx[label]] = 1.0
    return ids

def preprocess(examples):
    result = {'sentence': [[0, 3, 4] for _ in range(len(examples['labels']))]}
    print('"labels" are int', examples['labels'])
    result['labels'] = [multi_labels_to_ids(l) for l in examples['labels']]
    print('"labels" were converted to multi-label format with float values', result['labels'])
    return result


preprocessed_dataset = dataset.map(preprocess, batched=True, remove_columns=['labels', 'text'])
print(preprocessed_dataset[0]['labels'])
# Output: "[1, 1, 1, 0, 0]"
# Expected: "[1.0, 1.0, 1.0, 0.0, 0.0]"
```

If the original dataset uses non-int labels, it works as expected.

```python
from datasets import Dataset

data = {
    'text': ['text1', 'text2', 'text3', 'text4'],
    'labels': [['label1', 'label2', 'label3'], ['label4'], ['label4', 'label5'], ['label4']]
}

dataset = Dataset.from_dict(data)

label_set = set([label for labels in data['labels'] for label in labels])
label2idx = {label: idx for idx, label in enumerate(sorted(label_set))}


def multi_labels_to_ids(labels):
    ids = [0.0] * len(label2idx)
    for label in labels:
        ids[label2idx[label]] = 1.0
    return ids

def preprocess(examples):
    result = {'sentence': [[0, 3, 4] for _ in range(len(examples['labels']))]}
    print('"labels" are int', examples['labels'])
    result['labels'] = [multi_labels_to_ids(l) for l in examples['labels']]
    print('"labels" were converted to multi-label format with float values', result['labels'])
    return result


preprocessed_dataset = dataset.map(preprocess, batched=True, remove_columns=['labels', 'text'])
print(preprocessed_dataset[0]['labels'])
# Output: "[1.0, 1.0, 1.0, 0.0, 0.0]"
# Expected: "[1.0, 1.0, 1.0, 0.0, 0.0]"
```

Note that the only difference between these two examples is
>    'labels': [[0, 1, 2], [3], [3, 4], [3]]

v.s

>    'labels': [['label1', 'label2', 'label3'], ['label4'], ['label4', 'label5'], ['label4']]



### Expected behavior

Even if the original dataset uses a list of int labels, the int->float casting during `.map` process should not be canceled as shown in the above example

### Environment info

OS Ubuntu 22.04 LTS
Python 3.10.11
datasets v3.4.1

Comments:
Hi ! By default `map()` tries to keep the types of each column of the dataset, so here it reuses the int type since all your float values can be converted to integers. But I agree it would be nice to store float values as float values and don't try to reuse the same type in this case.

In the meantime, you can either store the float values in a new column, or pass the output `features=` manually to `map()`

Hi @lhoestq 

Thank you for the answer & suggestion!

Can we add some flag to `map()` function like `reuses_original_type=True` and skip reusing the original type when it's False?

Let me know if it sounds like a reasonable solution. I am happy to submit a PR for this.

In general we try to avoid adding new parameters when it's already possible to achieve the same results with existing parameters (here `features=`). But since it's not always convenient to know in advance the `features=` I'm open to contributions to adding this parameter yes

Thank you for sharing the context. Good to know that. 

I submitted a PR #7483. Could you review the PR?

Hi @lhoestq 

Let me know if there is something that I should add to [the PR](https://github.com/huggingface/datasets/pull/7483)!

Closing this issue as the PR #7483 was merged