from datasets import load_dataset
from transformers import Trainer, TrainingArguments, DataCollatorForSeq2Seq
​
def preprocess_for_question_generation(examples):
    if examples["answer"] is None or examples["context"] is None:
        return None
    
    if not isinstance(examples["answer"], list) or not isinstance(examples["context"], list):
        return None
    
    if len(examples["answer"]) != len(examples["context"]):
        return None
    
    prepared_input_texts = []
    for context, answer in zip(examples["context"], examples["answer"]):
        if "prepend" in examples:
            input_format = f"answer: {answer} context: {context}"
        elif "highlight" in examples:
            input_format = f"<hl> {answer} <hl> {context}"
        elif "answer_extraction" in examples:
            sentences = context.split(".")
            highlighted_sentences = []
            for sentence in sentences:
                if any(ans.lower() in sentence.lower() for ans in examples["answer"]):
                    highlighted_sentences.append(f"<hl> {sentence} <hl>")
            input_format = " ".join(highlighted_sentences)
        else:
            return None
        
        prepared_input_texts.append(input_format)
    
    tokenized_inputs = tokenizer(
        prepared_input_texts,
        max_length=128,
        padding="max_length",
        truncation=True,
        return_tensors="pt"
    )
    
    if "question" in examples: 
        tokenized_questions = tokenizer(
            examples["question"],
            max_length=64,
            padding="max_length",
            truncation=True,
            return_tensors="pt"
        )
        tokenized_inputs["labels"] = tokenized_questions["input_ids"].clone()
    
    return tokenized_inputs
​
# Define data collator
data_collator = DataCollatorForSeq2Seq(tokenizer)
​
# Define training arguments
training_args = TrainingArguments(
    output_dir="./question_generation_model",
    num_train_epochs=2,
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    save_total_limit=2,
    evaluation_strategy="epoch",
    report_to=[],  # Disable wandb logging
    learning_rate=5e-5
)
​
# Initialize Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset.map(preprocess_for_question_generation, batched=True),
    eval_dataset=val_dataset.map(preprocess_for_question_generation, batched=True),
    tokenizer=tokenizer,
    data_collator=data_collator
)
​
# Train the model
trainer.train()
​
100%
35/35 [00:00<00:00, 82.27ba/s]
100%
9/9 [00:00<00:00, 74.75ba/s]
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In[76], line 78
     68 trainer = Trainer(
     69     model=model,
     70     args=training_args,
   (...)
     74     data_collator=data_collator
     75 )
     77 # Train the model
---> 78 trainer.train()

File /opt/conda/lib/python3.10/site-packages/transformers/trainer.py:1624, in Trainer.train(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)
   1622         hf_hub_utils.enable_progress_bars()
   1623 else:
-> 1624     return inner_training_loop(
   1625         args=args,
   1626         resume_from_checkpoint=resume_from_checkpoint,
   1627         trial=trial,
   1628         ignore_keys_for_eval=ignore_keys_for_eval,
   1629     )

File /opt/conda/lib/python3.10/site-packages/transformers/trainer.py:1928, in Trainer._inner_training_loop(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)
   1925     rng_to_sync = True
   1927 step = -1
-> 1928 for step, inputs in enumerate(epoch_iterator):
   1929     total_batched_samples += 1
   1931     if self.args.include_num_input_tokens_seen:

File /opt/conda/lib/python3.10/site-packages/accelerate/data_loader.py:452, in DataLoaderShard.__iter__(self)
    450 # We iterate one batch ahead to check when we are at the end
    451 try:
--> 452     current_batch = next(dataloader_iter)
    453 except StopIteration:
    454     yield

File /opt/conda/lib/python3.10/site-packages/torch/utils/data/dataloader.py:630, in _BaseDataLoaderIter.__next__(self)
    627 if self._sampler_iter is None:
    628     # TODO(https://github.com/pytorch/pytorch/issues/76750)
    629     self._reset()  # type: ignore[call-arg]
--> 630 data = self._next_data()
    631 self._num_yielded += 1
    632 if self._dataset_kind == _DatasetKind.Iterable and \
    633         self._IterableDataset_len_called is not None and \
    634         self._num_yielded > self._IterableDataset_len_called:

File /opt/conda/lib/python3.10/site-packages/torch/utils/data/dataloader.py:674, in _SingleProcessDataLoaderIter._next_data(self)
    672 def _next_data(self):
    673     index = self._next_index()  # may raise StopIteration
--> 674     data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
    675     if self._pin_memory:
    676         data = _utils.pin_memory.pin_memory(data, self._pin_memory_device)

File /opt/conda/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py:51, in _MapDatasetFetcher.fetch(self, possibly_batched_index)
     49         data = self.dataset.__getitems__(possibly_batched_index)
     50     else:
---> 51         data = [self.dataset[idx] for idx in possibly_batched_index]
     52 else:
     53     data = self.dataset[possibly_batched_index]

File /opt/conda/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py:51, in <listcomp>(.0)
     49         data = self.dataset.__getitems__(possibly_batched_index)
     50     else:
---> 51         data = [self.dataset[idx] for idx in possibly_batched_index]
     52 else:
     53     data = self.dataset[possibly_batched_index]

File /opt/conda/lib/python3.10/site-packages/datasets/arrow_dataset.py:1764, in Dataset.__getitem__(self, key)
   1762 def __getitem__(self, key):  # noqa: F811
   1763     """Can be used to index columns (by string names) or rows (by integer index or iterable of indices or bools)."""
-> 1764     return self._getitem(
   1765         key,
   1766     )

File /opt/conda/lib/python3.10/site-packages/datasets/arrow_dataset.py:1748, in Dataset._getitem(self, key, decoded, **kwargs)
   1746 format_kwargs = format_kwargs if format_kwargs is not None else {}
   1747 formatter = get_formatter(format_type, features=self.features, decoded=decoded, **format_kwargs)
-> 1748 pa_subtable = query_table(self._data, key, indices=self._indices if self._indices is not None else None)
   1749 formatted_output = format_table(
   1750     pa_subtable, key, formatter=formatter, format_columns=format_columns, output_all_columns=output_all_columns
   1751 )
   1752 return formatted_output

File /opt/conda/lib/python3.10/site-packages/datasets/formatting/formatting.py:486, in query_table(table, key, indices)
    484 else:
    485     size = indices.num_rows if indices is not None else table.num_rows
--> 486     _check_valid_index_key(key, size)
    487 # Query the main table
    488 if indices is None:

File /opt/conda/lib/python3.10/site-packages/datasets/formatting/formatting.py:429, in _check_valid_index_key(key, size)
    427 if isinstance(key, int):
    428     if (key < 0 and key + size < 0) or (key >= size):
--> 429         raise IndexError(f"Invalid key: {key} is out of bounds for size {size}")
    430     return
    431 elif isinstance(key, slice):

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