T2TDataCollator 'target_ids' key error
Author: shubhambharadwaj
Pull request: no

Hi all,

I'm facing issues with this part of the code (post making changes as suggested [here](https://github.com/huggingface/transformers/issues/5049)) in T5-Base for QA.

```
import dataclasses
import logging
import os
import sys
from dataclasses import dataclass, field
from typing import Dict, List, Optional

import numpy as np
import torch

from transformers import T5ForConditionalGeneration, T5Tokenizer, EvalPrediction
from transformers import (
    HfArgumentParser,
    DataCollator,
    Trainer,
    TrainingArguments,
    set_seed,
)


logger = logging.getLogger(__name__)

# prepares lm_labels from target_ids, returns examples with keys as expected by the forward method
# this is necessacry because the trainer directly passes this dict as arguments to the model
# so make sure the keys match the parameter names of the forward method
@dataclass
class T2TDataCollator: #(DataCollator)
    def __call__(self, batch: List) -> Dict[str, torch.Tensor]: #collate_batch
        """
        Take a list of samples from a Dataset and collate them into a batch.
        Returns:
            A dictionary of tensors
        """
        input_ids = torch.stack([example['input_ids'] for example in batch])
        lm_labels = torch.stack([example['target_ids'] for example in batch])
        lm_labels[lm_labels[:, :] == 0] = -100
        attention_mask = torch.stack([example['attention_mask'] for example in batch])
        decoder_attention_mask = torch.stack([example['target_attention_mask'] for example in batch])
        

        return {
            'input_ids': input_ids, 
            'attention_mask': attention_mask,
            'lm_labels': lm_labels, 
            'decoder_attention_mask': decoder_attention_mask
        }
```

Which is fetching this error:-

```
Exception in thread Thread-12:
Traceback (most recent call last):
  File "/usr/lib/python3.6/threading.py", line 916, in _bootstrap_inner
    self.run()
  File "/usr/lib/python3.6/threading.py", line 864, in run
    self._target(*self._args, **self._kwargs)
  File "/usr/local/lib/python3.6/dist-packages/torch_xla/distributed/parallel_loader.py", line 133, in _loader_worker
    _, data = next(data_iter)
  File "/usr/local/lib/python3.6/dist-packages/torch/utils/data/dataloader.py", line 517, in __next__
    data = self._next_data()
  File "/usr/local/lib/python3.6/dist-packages/torch/utils/data/dataloader.py", line 557, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/usr/local/lib/python3.6/dist-packages/torch/utils/data/_utils/fetch.py", line 47, in fetch
    return self.collate_fn(data)
  File "<ipython-input-7-7b8c1b4d4c9a>", line 36, in __call__
    lm_labels = torch.stack([example['target_ids'] for example in batch])
  File "<ipython-input-7-7b8c1b4d4c9a>", line 36, in <listcomp>
    lm_labels = torch.stack([example['target_ids'] for example in batch])
KeyError: 'target_ids'
```
My train and validation dataset has `'target_ids'` field (read from `datasets.Dataset.from_pandas()` method and mapped the `add_eos_to_examples` and `convert_to_features` successfully):

`train_dataset['target_ids']`

```
tensor([[ 1027,  9533,  3440,  ...,     0,     0,     0],
        [ 7327,  1387, 11597,  ...,     0,     0,     0],
        [  272,     5,  7130,  ...,     0,     0,     0],
        ...,
        [15810,     1,     0,  ...,     0,     0,     0],
        [ 7107,     1,     0,  ...,     0,     0,     0],
        [  454,     5,   134,  ...,     0,     0,     0]])
```

`valid_dataset['target_ids']`

```
tensor([[15810,     1,     0,  ...,     0,     0,     0],
        [ 4190,  4329,     1,  ...,     0,     0,     0],
        [ 4329,    11,  7107,  ...,     0,     0,     0],
        ...,
        [    3,     4,     1,  ...,     0,     0,     0],
        [    3,     4,     1,  ...,     0,     0,     0],
        [ 8642,  4425,     9,  ...,     0,     0,     0]])
```
I am unable to fetch this field using class `T2TDataCollator:`. Please assist, thank you!

Comments:
Maybe @sgugger has an idea.

Hi @maxie320,

The `Trainer` now removes unused keys from the dataset if the dataset is an instance of `datasets.Dataset`. By unused, it means all the keys which are not in the model's forward method's argument list. And since `target_ids` is not a argument expected by the forward it's getting removed by the `Trainer` and hence the `KeyError`

You can rename the `target_ids` key by `labels` and also change the collator accordingly which  should fix this issue

@patil-suraj It's now showing a key error on `target_attention_mask`. I'm guessing this name has been changed as well?

Okay, figured it out, used the same name `decoder_attention_mask` in `T2TDataCollator` as given in `forward()` method argument list, thanks for the assist @patil-suraj

Also note that you can set `remove_unused_columns=False` in your `TrainingArguments` to disable the behavior where Trainer drops the columns not in the model signature.

Sure, thank you! @sgugger

Closing this issue since it seems solved, don't hesitate to reopen if you have more problems!

Hello everyone,
I am trying to run the same notebook given in https://colab.research.google.com/github/patil-suraj/exploring-T5/blob/master/T5_on_TPU.ipynb and I have a similar problem mentioned here. I applied the possible changes mentioned here, but it could not solve my problem.

There was a problem in nightly version with import torch as mentioned https://stackoverflow.com/questions/67257008/oserror-libmkl-intel-lp64-so-1-cannot-open-shared-object-file-no-such-file-or/67479054#67479054. After I add the modifications, it throws the error below:

/usr/local/lib/python3.7/dist-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
1049 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
1050 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1051 return forward_call(*input, **kwargs)
1052 # Do not call functions when jit is used
1053 full_backward_hooks, non_full_backward_hooks = [], []

TypeError: forward() got an unexpected keyword argument 'lm_labels'

Any ideas?
Thanks

T5 expects `labels` now, not `lm_labels`. You should replace that in the return statement of your data collator
```
class T2TDataCollator(DataCollator):
    def collate_batch(self, batch: List) -> Dict[str, torch.Tensor]:
        """
        Take a list of samples from a Dataset and collate them into a batch.
        Returns:
            A dictionary of tensors
        """
        input_ids = torch.stack([example['input_ids'] for example in batch])
        lm_labels = torch.stack([example['target_ids'] for example in batch])
        lm_labels[lm_labels[:, :] == 0] = -100
        attention_mask = torch.stack([example['attention_mask'] for example in batch])
        decoder_attention_mask = torch.stack([example['target_attention_mask'] for example in batch])
        

        return {
            'input_ids': input_ids, 
            'attention_mask': attention_mask,
            'lm_labels': lm_labels, 
            'decoder_attention_mask': decoder_attention_mask
        }
```

Also that data collator should have a `__call__` method, not a `collate_batch`.

Thank you for your quick reply @sgugger, now it works!