IndexError Not Solving -> IndexError: Invalid key: ?? is out of bounds for size 0 or ??

### Describe the bug

in :1 │
│ │
│ /usr/local/lib/python3.10/dist-packages/transformers/trainer.py:1537 in train │
│ │
│ 1534 │ │ inner_training_loop = find_executable_batch_size( │
│ 1535 │ │ │ self._inner_training_loop, self._train_batch_size, args.auto_find_batch_size │
│ 1536 │ │ ) │
│ ❱ 1537 │ │ return inner_training_loop( │
│ 1538 │ │ │ args=args, │
│ 1539 │ │ │ resume_from_checkpoint=resume_from_checkpoint, │
│ 1540 │ │ │ trial=trial, │
│ │
│ /usr/local/lib/python3.10/dist-packages/transformers/trainer.py:1789 in _inner_training_loop │
│ │
│ 1786 │ │ │ │ rng_to_sync = True │
│ 1787 │ │ │ │
│ 1788 │ │ │ step = -1 │
│ ❱ 1789 │ │ │ for step, inputs in enumerate(epoch_iterator): │
│ 1790 │ │ │ │ total_batched_samples += 1 │
│ 1791 │ │ │ │ if rng_to_sync: │
│ 1792 │ │ │ │ │ self._load_rng_state(resume_from_checkpoint) │
│ │
│ /usr/local/lib/python3.10/dist-packages/accelerate/data_loader.py:377 in __iter__ │
│ │
│ 374 │ │ dataloader_iter = super().__iter__() │
│ 375 │ │ # We iterate one batch ahead to check when we are at the end │
│ 376 │ │ try: │
│ ❱ 377 │ │ │ current_batch = next(dataloader_iter) │
│ 378 │ │ except StopIteration: │
│ 379 │ │ │ yield │
│ 380 │
│ │
│ /usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py:633 in __next__ │
│ │
│ 630 │ │ │ if self._sampler_iter is None: │
│ 631 │ │ │ │ # TODO(https://github.com/pytorch/pytorch/issues/76750) │
│ 632 │ │ │ │ self._reset() # type: ignore[call-arg] │
│ ❱ 633 │ │ │ data = self._next_data() │
│ 634 │ │ │ self._num_yielded += 1 │
│ 635 │ │ │ if self._dataset_kind == _DatasetKind.Iterable and \ │
│ 636 │ │ │ │ │ self._IterableDataset_len_called is not None and \ │
│ │
│ /usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py:677 in _next_data │
│ │
│ 674 │ │
│ 675 │ def _next_data(self): │
│ 676 │ │ index = self._next_index() # may raise StopIteration │
│ ❱ 677 │ │ data = self._dataset_fetcher.fetch(index) # may raise StopIteration │
│ 678 │ │ if self._pin_memory: │
│ 679 │ │ │ data = _utils.pin_memory.pin_memory(data, self._pin_memory_device) │
│ 680 │ │ return data │
│ │
│ /usr/local/lib/python3.10/dist-packages/torch/utils/data/_utils/fetch.py:49 in fetch │
│ │
│ 46 │ def fetch(self, possibly_batched_index): │
│ 47 │ │ if self.auto_collation: │
│ 48 │ │ │ if hasattr(self.dataset, "__getitems__") and self.dataset.__getitems__: │
│ ❱ 49 │ │ │ │ data = self.dataset.__getitems__(possibly_batched_index) │
│ 50 │ │ │ else: │
│ 51 │ │ │ │ data = [self.dataset[idx] for idx in possibly_batched_index] │
│ 52 │ │ else: │
│ │
│ /usr/local/lib/python3.10/dist-packages/datasets/arrow_dataset.py:2782 in __getitems__ │
│ │
│ 2779 │ │
│ 2780 │ def __getitems__(self, keys: List) -> List: │
│ 2781 │ │ """Can be used to get a batch using a list of integers indices.""" │
│ ❱ 2782 │ │ batch = self.__getitem__(keys) │
│ 2783 │ │ n_examples = len(batch[next(iter(batch))]) │
│ 2784 │ │ return [{col: array[i] for col, array in batch.items()} for i in range(n_example │
│ 2785 │
│ │
│ /usr/local/lib/python3.10/dist-packages/datasets/arrow_dataset.py:2778 in __getitem__ │
│ │
│ 2775 │ │
│ 2776 │ def __getitem__(self, key): # noqa: F811 │
│ 2777 │ │ """Can be used to index columns (by string names) or rows (by integer index or i │
│ ❱ 2778 │ │ return self._getitem(key) │
│ 2779 │ │
│ 2780 │ def __getitems__(self, keys: List) -> List: │
│ 2781 │ │ """Can be used to get a batch using a list of integers indices.""" │
│ │
│ /usr/local/lib/python3.10/dist-packages/datasets/arrow_dataset.py:2762 in _getitem │
│ │
│ 2759 │ │ format_kwargs = kwargs["format_kwargs"] if "format_kwargs" in kwargs else self._ │
│ 2760 │ │ format_kwargs = format_kwargs if format_kwargs is not None else {} │
│ 2761 │ │ formatter = get_formatter(format_type, features=self._info.features, **format_kw │
│ ❱ 2762 │ │ pa_subtable = query_table(self._data, key, indices=self._indices if self._indice │
│ 2763 │ │ formatted_output = format_table( │
│ 2764 │ │ │ pa_subtable, key, formatter=formatter, format_columns=format_columns, output │
│ 2765 │ │ ) │
│ │
│ /usr/local/lib/python3.10/dist-packages/datasets/formatting/formatting.py:578 in query_table │
│ │
│ 575 │ │ _check_valid_column_key(key, table.column_names) │
│ 576 │ else: │
│ 577 │ │ size = indices.num_rows if indices is not None else table.num_rows │
│ ❱ 578 │ │ _check_valid_index_key(key, size) │
│ 579 │ # Query the main table │
│ 580 │ if indices is None: │
│ 581 │ │ pa_subtable = _query_table(table, key) │
│ │
│ /usr/local/lib/python3.10/dist-packages/datasets/formatting/formatting.py:531 in │
│ _check_valid_index_key │
│ │
│ 528 │ │ │ _check_valid_index_key(min(key), size=size) │
│ 529 │ elif isinstance(key, Iterable): │
│ 530 │ │ if len(key) > 0: │
│ ❱ 531 │ │ │ _check_valid_index_key(int(max(key)), size=size) │
│ 532 │ │ │ _check_valid_index_key(int(min(key)), size=size) │
│ 533 │ else: │
│ 534 │ │ _raise_bad_key_type(key) │
│ │
│ /usr/local/lib/python3.10/dist-packages/datasets/formatting/formatting.py:521 in │
│ _check_valid_index_key │
│ │
│ 518 def _check_valid_index_key(key: Union[int, slice, range, Iterable], size: int) -> None: │
│ 519 │ if isinstance(key, int): │
│ 520 │ │ if (key < 0 and key + size < 0) or (key >= size): │
│ ❱ 521 │ │ │ raise IndexError(f"Invalid key: {key} is out of bounds for size {size}") │
│ 522 │ │ return │
│ 523 │ elif isinstance(key, slice): │
│ 524 │ │ pass

### Steps to reproduce the bug

``
import json
import os
from pprint import pprint

import bitsandbytes as bnb
import pandas as pd
import torch
import torch.nn as nn
import transformers
from datasets import Dataset,load_dataset

from peft import (
LoraConfig,
PeftConfig,
PeftModel,
get_peft_model,
prepare_model_for_kbit_training
)

from transformers import (
AutoConfig,
AutoModelForCausalLM,
AutoTokenizer,
BitsAndBytesConfig,

)

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

def print_trainable_parameters(model):
"""
Prints the number of trainable parameters in the model.
"""
trainable_params = 0
all_param = 0
for _, param in model.named_parameters():
all_param += param.numel()
if param.requires_grad:
trainable_params += param.numel()
print(
f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}"
)

MODEL_NAME = "tiiuae/falcon-7b"

bnb_config = BitsAndBytesConfig(
load_in_4bit = True,
bnb_4bit_use_double_quant=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.bfloat16,
)

model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
device_map = "auto",
trust_remote_code = True,
quantization_config = bnb_config
)

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
tokenizer.pad_token = tokenizer.eos_token

model.gradient_checkpointing_enable()
model = prepare_model_for_kbit_training(model)

config = LoraConfig(
r = 16,
lora_alpha = 32,
target_modules = ["query_key_value"],
lora_dropout = 0.05,
bias = "none",
task_type = "CASUAL_LM"
)

model = get_peft_model(model,config)
print_trainable_parameters(model)

def generate_prompt(data_point):
return f"""
: {data_point["question"]}
: {data_point["answer"]}
""".strip()

def generate_and_tokenize_prompt(data_point):
full_prompt = generate_prompt(data_point)
tokenized_full_prompt = tokenizer(full_prompt, padding = True, truncation = True,return_tensors = None)
return dict({
"input_ids" : tokenized_full_prompt["input_ids"],
"attention_mask" : tokenized_full_prompt["attention_mask"]

})

data = data["train"].shuffle().map(generate_and_tokenize_prompt, batched = False)

OUTPUT_DIR = "experiments"

trainings_args = transformers.TrainingArguments(
per_device_train_batch_size = 1,
gradient_accumulation_steps = 4,
num_train_epochs = 1,
learning_rate = 2e-4,
fp16 = True,
save_total_limit = 3,
logging_steps = 1,
output_dir = OUTPUT_DIR,
max_steps = 80,
optim = "paged_adamw_8bit",
lr_scheduler_type = "cosine",
warmup_ratio = 0.05,
#remove_unused_columns=True
)

trainer = transformers.Trainer(
model = model,
train_dataset = data,
args = trainings_args,
data_collator = transformers.DataCollatorForLanguageModeling(tokenizer, mlm=False),

)

model.config.use_cache = False

trainer.train()

IndexError: Invalid key: 32 is out of bounds for size 0

DataSet Format is like :
[{"question": "How can I create an account?", "answer": "To create an account, click on the 'Sign Up' button on the top right corner of our website and follow the instructions to complete the registration process."}, .... ]

### Expected behavior

-

### Environment info

!pip install -q pip
!pip install -q bitsandbytes==0.39.0
!pip install -q torch==2.0.1

!pip install -q git+https://github.com/huggingface/transformers.git
!pip install -q git+https://github.com/huggingface/peft.git
!pip install -q git+https://github.com/huggingface/accelerate.git

!pip install -q datasets
!pip install -q loralib==0.1.1
!pip install -q einops==0.6.1

import json
import os
from pprint import pprint

import bitsandbytes as bnb
import pandas as pd
import torch
import torch.nn as nn
import transformers
from datasets import Dataset,load_dataset

from peft import (
LoraConfig,
PeftConfig,
PeftModel,
get_peft_model,
prepare_model_for_kbit_training
)

from transformers import (
AutoConfig,
AutoModelForCausalLM,
AutoTokenizer,
BitsAndBytesConfig,

)

os.environ["CUDA_VISIBLE_DEVICES"] = "0"