Getting KeyError when trying to train my model

I wrote a code that will fine-tine my QnA model using huggingface, usin a dataset from a JSON file containing questions and answers. However, things went awry when I ran this traceback:

Traceback (most recent call last):
  File "C:\Users\Philip Chen\OneDrive\Documents\ML\machineLearning.py", line 45, in <module>
    trainer.train()
  File "C:\Users\Philip Chen\AppData\Local\Programs\Python\Python310\lib\site-packages\transformers\trainer.py", line 1859, in train
    return inner_training_loop(
  File "C:\Users\Philip Chen\AppData\Local\Programs\Python\Python310\lib\site-packages\transformers\trainer.py", line 2165, in _inner_training_loop
    for step, inputs in enumerate(epoch_iterator):
  File "C:\Users\Philip Chen\AppData\Local\Programs\Python\Python310\lib\site-packages\accelerate\data_loader.py", line 454, in __iter__
    current_batch = next(dataloader_iter)
  File "C:\Users\Philip Chen\AppData\Local\Programs\Python\Python310\lib\site-packages\torch\utils\data\dataloader.py", line 631, in __next__
    data = self._next_data()
  File "C:\Users\Philip Chen\AppData\Local\Programs\Python\Python310\lib\site-packages\torch\utils\data\dataloader.py", line 675, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "C:\Users\Philip Chen\AppData\Local\Programs\Python\Python310\lib\site-packages\torch\utils\data\_utils\fetch.py", line 51, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "C:\Users\Philip Chen\AppData\Local\Programs\Python\Python310\lib\site-packages\torch\utils\data\_utils\fetch.py", line 51, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "C:\Users\Philip Chen\AppData\Local\Programs\Python\Python310\lib\site-packages\datasets\dataset_dict.py", line 81, in __getitem__
    raise KeyError(
KeyError: "Invalid key: 0. Please first select a split. For example: `my_dataset_dictionary['train'][0]`. Available splits: ['train']"

From this offending line:

trainer.train()

And this is my full code:

# https://www.mlexpert.io/blog/alpaca-fine-tuning
# https://wellsr.com/python/fine-tuning-huggingface-models-in-tensorflow-keras/
# https://learnopencv.com/fine-tuning-bert/
# https://medium.com/@karary/nlp-fine-tune-question-answering-model-%E5%AF%A6%E4%BD%9C-3-model-training-%E5%84%B2%E5%AD%98%E8%88%87-inference-13d2a5bf5c32
import transformers as tf
import datasets as ds
import pandas as pd
import numpy as np
import torch
import json
 
############## Check if CUDA is enabled. ################
hasCUDA=torch.cuda.is_available()
print(f"CUDA Enabled? {hasCUDA}")
device="cuda" if hasCUDA else "cpu"      
 
############## Loading file and populating data ################
fileName="qna.json"
trainDS=ds.load_dataset("json", data_files=fileName)
evalDS=ds.load_dataset("json", data_files=fileName)
# rawDS=ds.load_dataset('squad')
############## Model ##########################################
modelName="./distilbert-base-cased"     #or replace the model name with whatever you feel like.
config=tf.AutoConfig.from_pretrained(modelName+"/config.json")
model=tf.AutoModelForQuestionAnswering.from_pretrained(modelName,config=config)
tokenizer=tf.AutoTokenizer.from_pretrained(modelName)
############## Training #######################################
trnArgs=tf.TrainingArguments(
    output_dir="./",
    evaluation_strategy="epoch",
    save_strategy="epoch",
    learning_rate=2e-5,
    num_train_epochs=3,
    remove_unused_columns=False,
    fp16=True
)
 
trainer=tf.Trainer(
    model=model,
    args=trnArgs,
    train_dataset=trainDS,
    eval_dataset=evalDS,
    tokenizer=tokenizer
)
trainer.train()

The JSON file:

{"text": "Who wrote Charlie and the Chocolate Factory?", "label": "Roald Dahl"}
{"text": "Name a few ways to treat constipation naturally.", "label": "Exercise regularly, eat more fibers, and drink more water."}
{"text": "Where is the longest roller coaster located?", "label": "Nagashima, Japan. The name of the coaster is Steel Dragon 2000."}
{"text": "What are the 11 herbs and spices that Colonel Sanders used in KFC?", "label": "Nobody knows, as it's a secret."}
{"text": "Who wrote Les Miserables?", "label": "Victor Hugo"}
{"text": "What is the Watergate Scandal?", "label": "The Watergate scandal was a significant political controversy in the United States during the presidency of Richard Nixon from 1972 to 1974, ultimately resulting in Nixon's resignation. It originated from attempts by the Nixon administration to conceal its involvement in the June 17, 1972, break-in at the Democratic National Committee headquarters located in the Watergate Office Building in Washington, D.C."}
{"text": "What is Obama's most famous quote?", "label": "'Yes we can!'"}

Please suggest how I should fix this problem. Thanks!

Have you tried doing what it told you to do?

So I managed to format the dataset to the trainer’s desired format from this code (Thanks, @m_hutley for pointing it out)

############## Encoding and Tokenizing #######################################
sampleDS2=sampleDS2.map(lambda batch: tokenizer(sampleDS2["text"], truncation=True, padding='max_length', max_length=512), batched=True)
sampleDS2.set_format("torch", columns=["input_ids", "attention_mask", "question"])
print(sampleDS2.column_names)

Here is the updated version:

# https://www.mlexpert.io/blog/alpaca-fine-tuning
# https://wellsr.com/python/fine-tuning-huggingface-models-in-tensorflow-keras/
# https://learnopencv.com/fine-tuning-bert/
# https://medium.com/@karary/nlp-fine-tune-question-answering-model-%E5%AF%A6%E4%BD%9C-3-model-training-%E5%84%B2%E5%AD%98%E8%88%87-inference-13d2a5bf5c32
# https://medium.com/@anyuanay/fine-tuning-the-pre-trained-bert-model-in-hugging-face-for-question-answering-8edc76890ce0
import transformers as tf
import datasets as ds
import pandas as pd
import numpy as np
import torch
import json
 
############## Check if CUDA is enabled. ################
hasCUDA=torch.cuda.is_available()
print(f"CUDA Enabled? {hasCUDA}")
device="cuda" if hasCUDA else "cpu"      
 
############## Loading file and populating data ################
fileName="qna.json"
sampleDS2=ds.load_dataset("json", data_files=fileName, split="train")

############## Model ##########################################
modelName="distilbert/distilbert-base-cased"     #or replace the model name with whatever you feel like.
# config=tf.AutoConfig.from_pretrained(modelName+"/config.json")
model=tf.AutoModelForQuestionAnswering.from_pretrained(modelName)
tokenizer=tf.AutoTokenizer.from_pretrained(modelName)

############## Encoding and Tokenizing #######################################
sampleDS2=sampleDS2.map(lambda batch: tokenizer(sampleDS2["text"], truncation=True, padding='max_length', max_length=512), batched=True)
sampleDS2.set_format("torch", columns=["input_ids", "attention_mask", "text"])

############## Training #######################################
trnArgs=tf.TrainingArguments(
    output_dir="./",
    evaluation_strategy="epoch",
    save_strategy="epoch",
    learning_rate=2e-5,
    num_train_epochs=3,
    remove_unused_columns=True,
    fp16=True
)
 
trainer=tf.Trainer(
    model=model,
    args=trnArgs,
    train_dataset=sampleDS2, 
    eval_dataset=sampleDS2,
    tokenizer=tokenizer
)
trainer.train()

but now I get this error



Traceback (most recent call last):
  File "C:\Users\chenp\Documents\ML\machineLearning.py", line 52, in <module>
    trainer.train()
  File "C:\Users\chenp\AppData\Local\Programs\Python\Python310\lib\site-packages\transformers\trainer.py", line 1859, in train
    return inner_training_loop(
  File "C:\Users\chenp\AppData\Local\Programs\Python\Python310\lib\site-packages\transformers\trainer.py", line 2203, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs)
  File "C:\Users\chenp\AppData\Local\Programs\Python\Python310\lib\site-packages\transformers\trainer.py", line 3138, in training_step
    loss = self.compute_loss(model, inputs)
  File "C:\Users\chenp\AppData\Local\Programs\Python\Python310\lib\site-packages\transformers\trainer.py", line 3179, in compute_loss
    raise ValueError(
ValueError: The model did not return a loss from the inputs, only the following keys: start_logits,end_logits. For reference, the inputs it received are input_ids,attention_mask.
  0%|          | 0/9 [00:00<?, ?it/s]

Please, I need this to work. SUggestions are appreciated!

https://discuss.huggingface.co/t/the-model-did-not-return-a-loss-from-the-inputs-only-the-following-keys-logits-for-reference-the-inputs-it-received-are-input-values/25420/3 ?