Extra dataset features not passing to the custom collator
Author: ArlindKadra
Pull request: no

Hey everyone, I have implemented a Dataset from a generator as follows:

```
def gen(torch_dataset):
    for idx in len(torch_dataset):
        yield torch_dataset[idx] 

train_dataset = Dataset.from_generator(gen, gen_kwargs={"torch_dataset": torch_train_dataset})
```

I am using the weighted trainer from the following tutorial:
https://github.com/huggingface/blog/blob/main/Lora-for-sequence-classification-with-Roberta-Llama-Mistral.md

Where I added an extra argument in the `compute_loss`:
`input.pop['extra_argument']`

And lastly, a custom data collator:

```
class CustomDataCollatorWithPadding(DataCollatorWithPadding):
    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
        # Extract unique_labels separately
        extra_argument = [feature.pop('extra_argument') for feature in features]

        # Use the parent class's __call__ method to handle padding
        batch = super().__call__(features)

        # Add unique_labels back into the batch
        batch['extra_argument'] = torch.tensor(extra_argument)

        return batch
```

When I check my training script and print a batch from `train_dataset`, everything is fine and it includes all the necessary features. However, when I print the features from the data collator, my `'extra_argument'` is missing and the code fails. I am using the trainer from the above tutorial and I am additionally passing my custom dataset that yields all the necessary features to the trainer along with the custom collator. 

Am I making a mistake in my setup or is there something in between that removes the extra argument?

Python 3.11.9

Environment:
```
accelerate==0.30.1
aiohttp==3.9.5
aiosignal==1.3.1
argon2-cffi==23.1.0
argon2-cffi-bindings==21.2.0
attrs==23.2.0
certifi==2024.6.2
cffi==1.16.0
charset-normalizer==3.3.2
click==8.1.7
datasets==2.19.2
dill==0.3.8
docker-pycreds==0.4.0
evaluate==0.4.2
filelock==3.14.0
frozenlist==1.4.1
fsspec==2024.3.1
gitdb==4.0.11
GitPython==3.1.43
huggingface-hub==0.23.3
idna==3.7
Jinja2==3.1.3
joblib==1.4.2
liac-arff==2.5.0
MarkupSafe==2.1.5
minio==7.2.7
mpmath==1.3.0
multidict==6.0.5
multiprocess==0.70.16
networkx==3.2.1
numpy==1.26.4
nvidia-cublas-cu12==12.1.3.1
nvidia-cuda-cupti-cu12==12.1.105
nvidia-cuda-nvrtc-cu12==12.1.105
nvidia-cuda-runtime-cu12==12.1.105
nvidia-cudnn-cu12==8.9.2.26
nvidia-cufft-cu12==11.0.2.54
nvidia-curand-cu12==10.3.2.106
nvidia-cusolver-cu12==11.4.5.107
nvidia-cusparse-cu12==12.1.0.106
nvidia-nccl-cu12==2.20.5
nvidia-nvjitlink-cu12==12.1.105
nvidia-nvtx-cu12==12.1.105
openml==0.14.2
packaging==24.0
pandas==2.2.2
peft==0.11.1
pillow==10.2.0
platformdirs==4.2.2
protobuf==4.25.3
psutil==5.9.8
pyarrow==16.1.0
pyarrow-hotfix==0.6
pycparser==2.22
pycryptodome==3.20.0
python-dateutil==2.9.0.post0
pytz==2024.1
PyYAML==6.0.1
regex==2024.5.15
requests==2.32.3
safetensors==0.4.3
scikit-learn==1.5.0
scipy==1.13.1
sentencepiece==0.2.0
sentry-sdk==2.4.0
setproctitle==1.3.3
six==1.16.0
smmap==5.0.1
sympy==1.12
threadpoolctl==3.5.0
tokenizers==0.19.1
torch==2.3.1+cu121
torchaudio==2.3.1+cu121
torchvision==0.18.1+cu121
tqdm==4.66.4
transformers==4.41.2
triton==2.3.1
typing_extensions==4.12.1
tzdata==2024.1
urllib3==2.2.1
wandb==0.17.0
xmltodict==0.13.0
xxhash==3.4.1
yarl==1.9.4
```




Comments:
Hi @ArlindKadra, thanks for raising an issue! 

As this involves questions about a custom script, this is a question best placed in our [forums](https://discuss.huggingface.co/). We try to reserve the github issues for feature requests and bug reports.

Hey @amyeroberts thanks for the reply.

I am not very familiar with the library, so I would like to know if the above behavior is supported by the library first. If it is supported by the library, in turn it could be a mistake by me or a bug.
Plus, on the other hand, on the forum, it does say you can create an issue about it :).

The forum did not help as not many questions do get any responses. However, for anyone who might stumble upon the issue in the future, passing the following to the trainer does solve it:

`remove_unused_columns=False`



@ArlindKadra  when passing the argument I get:

TypeError: Trainer.__init__() got an unexpected keyword argument 'remove_unused_columns'

During handling of the above exception, another exception occurred: