IndexError: Invalid key is out of bounds for size 0 despite having a populated dataset
Opened by ZYM66 (closed)

### Describe the bug

I am encountering an `IndexError` when trying to access data from a DataLoader which wraps around a dataset I've loaded using the `datasets` library. The error suggests that the dataset size is `0`, but when I check the length and print the dataset, it's clear that it has `1166` entries.

### Steps to reproduce the bug

1. Load a dataset with `1166` entries.
2. Create a DataLoader using this dataset.
3. Try iterating over the DataLoader.

code:
```python
def get_train_dataloader(self) -> DataLoader:
        if self.train_dataset is None:
            raise ValueError("Trainer: training requires a train_dataset.")

        train_dataset = self.train_dataset
        data_collator = self.data_collator
        print(len(train_dataset))
        print(train_dataset)
        if is_datasets_available() and isinstance(train_dataset, datasets.Dataset):
            train_dataset = self._remove_unused_columns(train_dataset, description="training")
        else:
            data_collator = self._get_collator_with_removed_columns(data_collator, description="training")
        train_sampler = self._get_train_sampler()
        dl = DataLoader(
            train_dataset,
            batch_size=self._train_batch_size,
            sampler=train_sampler,
            collate_fn=data_collator,
            drop_last=self.args.dataloader_drop_last,
            num_workers=self.args.dataloader_num_workers,
            pin_memory=self.args.dataloader_pin_memory,
            worker_init_fn=seed_worker,
        )
        print(dl)
        print(len(dl))
        for i in dl:
            print(i)
            break
        return dl
```

output : 
```
1166
Dataset({
    features: ['input_ids', 'special_tokens_mask'],
    num_rows: 1166
})
<torch.utils.data.dataloader.DataLoader object ...>
146
```
Error:
```
Traceback (most recent call last):
  File "/home/dl/zym/llamaJP/TestUseContinuePretrainLlama.py", line 266, in <module>
    train()
  File "/home/dl/zym/llamaJP/TestUseContinuePretrainLlama.py", line 260, in train
    trainer.train()
  File "/root/miniconda3/envs/LLM/lib/python3.10/site-packages/transformers/trainer.py", line 1506, in train
    return inner_training_loop(
  File "/root/miniconda3/envs/LLM/lib/python3.10/site-packages/transformers/trainer.py", line 1520, in _inner_training_loop
    train_dataloader = self.get_train_dataloader()
  File "/home/dl/zym/llamaJP/TestUseContinuePretrainLlama.py", line 80, in get_train_dataloader
    for i in dl:
  File "/root/miniconda3/envs/LLM/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 630, in __next__
    data = self._next_data()
  File "/root/miniconda3/envs/LLM/lib/python3.10/site-packages/torch/utils/data/dataloader.py", line 674, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/root/miniconda3/envs/LLM/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 49, in fetch
    data = self.dataset.__getitems__(possibly_batched_index)
  File "/root/miniconda3/envs/LLM/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 2807, in __getitems__
    batch = self.__getitem__(keys)
  File "/root/miniconda3/envs/LLM/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 2803, in __getitem__
    return self._getitem(key)
  File "/root/miniconda3/envs/LLM/lib/python3.10/site-packages/datasets/arrow_dataset.py", line 2787, in _getitem
    pa_subtable = query_table(self._data, key, indices=self._indices if self._indices is not None else None)
  File "/root/miniconda3/envs/LLM/lib/python3.10/site-packages/datasets/formatting/formatting.py", line 583, in query_table
    _check_valid_index_key(key, size)
  File "/root/miniconda3/envs/LLM/lib/python3.10/site-packages/datasets/formatting/formatting.py", line 536, in _check_valid_index_key
    _check_valid_index_key(int(max(key)), size=size)
  File "/root/miniconda3/envs/LLM/lib/python3.10/site-packages/datasets/formatting/formatting.py", line 526, in _check_valid_index_key
    raise IndexError(f"Invalid key: {key} is out of bounds for size {size}")
IndexError: Invalid key: 1116 is out of bounds for size 0
```

### Expected behavior

I expect to be able to iterate over the DataLoader without encountering an IndexError since the dataset is populated.

### Environment info

- `datasets` library version: [2.14.5]
- Platform: [Linux] 
- Python version: 3.10
- Other libraries involved: HuggingFace Transformers

Comments:
It looks to be the same issue as the one reported in https://discuss.huggingface.co/t/indexerror-invalid-key-16-is-out-of-bounds-for-size-0.

Can you check the length of `train_dataset` before the `train_sampler = self._get_train_sampler()` (and after `_remove_unused_columns`) line?