使用HuggingFace微调QA模型时KeyError问题的解决方法

错误信息 #

执行trainer.train()时触发以下异常：

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']"

完整代码 #

# 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()

数据文件内容 #

{"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!'"}

解决建议 #

1. 修正数据集引用方式 #

ds.load_dataset("json", data_files=fileName)返回的是DatasetDict对象，包含默认的train拆分，不能直接传入Trainer，必须明确指定拆分后的数据集：

# 方式1：直接取train拆分
trainDS = ds.load_dataset("json", data_files=fileName)['train']
evalDS = ds.load_dataset("json", data_files=fileName)['train']  # 建议后续拆分独立验证集

# 方式2：拆分训练/验证集（更合理）
dataset = ds.load_dataset("json", data_files=fileName)
dataset_split = dataset['train'].train_test_split(test_size=0.2)  # 按20%比例拆分验证集
trainDS = dataset_split['train']
evalDS = dataset_split['test']

将Trainer中的train_dataset和eval_dataset替换为上述处理后的数据集。

2. 适配问答模型的数据格式 #

你使用的AutoModelForQuestionAnswering是抽取式问答模型，要求数据包含question、context和answers字段（模型从context中抽取答案），当前text+label的格式不匹配：

• 若为开放域问答（无上下文），建议改用生成式模型（如T5、GPT类），对应使用AutoModelForSeq2SeqLM或AutoModelForCausalLM。
• 若坚持用抽取式模型，需调整数据格式，示例如下：

{"question": "Who wrote Charlie and the Chocolate Factory?", "context": "Charlie and the Chocolate Factory is a classic children's book written by Roald Dahl.", "answers": {"text": ["Roald Dahl"], "answer_start": [49]}}

其中answer_start是答案在context中的起始字符位置，需手动或自动标注。

3. 其他注意事项 #

• 确保fp16=True的环境支持混合精度训练（需CUDA设备），无CUDA时改为fp16=False。
• 检查remove_unused_columns=False是否必要，若模型不需要某些字段，可设为True自动过滤。

内容的提问来源于stack exchange，提问作者Jungle Zee