如何解决Hugging Face Transformer train()函数中的KeyError: 0错误

问题描述 #

在使用Hugging Face Transformer框架微调gpt2-medium模型时，执行trainer.train()触发KeyError: 0错误，错误栈如下：

KeyError                                  Traceback (most recent call last)
 in <cell line: 1>()
----> 1 trainer.train()

5 frames
/usr/local/lib/python3.9/dist-packages/torch/utils/data/_utils/fetch.py in (.0)
49                 data = self.dataset.__getitems__(possibly_batched_index)
50             else:
---&gt; 51                 data = [self.dataset[idx] for idx in possibly_batched_index]
52         else:
53             data = self.dataset[possibly_batched_index]

KeyError: 0

完整实现代码：

import pandas as pd 
import numpy as np

dataset = pd.read_csv('Train_rev1.csv',error_bad_lines=False, engine='python')

def replace_string(row):
    row['FullDescription'] = row['FullDescription'].replace('****', str(row['SalaryNormalized']))
    return row

dataset = dataset.apply(replace_string, axis=1)
dataset = dataset.drop(['ContractType','ContractTime','LocationRaw','SalaryRaw','SourceName','Id','Title', 'LocationNormalized', 'Company', 'Category',
       'SalaryNormalized'], axis=1)

! pip install -q transformers
from transformers import GPT2Tokenizer, GPT2LMHeadModel, Trainer, TrainingArguments
tokenizer = GPT2Tokenizer.from_pretrained('gpt2-medium')
tokenizer.add_special_tokens({'pad_token': '[PAD]'})
tokenized_data = tokenizer(dataset['FullDescription'].tolist(), truncation=True, padding=True)

# Split data into training and validation sets
train_size = int(0.8 * len(tokenized_data['input_ids']))
val_size = len(tokenized_data['input_ids']) - train_size

train_dataset = {'input_ids': tokenized_data['input_ids'][:train_size],
                 'attention_mask': tokenized_data['attention_mask'][:train_size]}
val_dataset = {'input_ids': tokenized_data['input_ids'][train_size:],
               'attention_mask': tokenized_data['attention_mask'][train_size:]}

from transformers import GPT2Config
# Define model configuration and instantiate model
model_config = GPT2Config.from_pretrained('gpt2-medium')
model_config.output_hidden_states = True
model = GPT2LMHeadModel.from_pretrained('gpt2-medium', config=model_config)

# Train model using Huggingface Trainer API
training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=1,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='./logs',
    logging_steps=10,
    evaluation_strategy='steps',
    eval_steps=50,
    load_best_model_at_end=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset
)

trainer.train()

问题原因 #

Hugging Face的Trainer依赖PyTorch的DataLoader加载数据，而普通Python字典无法满足DataLoader的要求——它需要数据集实现__getitem__和__len__方法来通过数字索引（如0、1）获取样本。直接用字典作为数据集时，DataLoader尝试用数字索引访问字典，而字典的键并非这些数字，因此触发KeyError。

解决方案 #

方法1：使用Hugging Face Dataset类（推荐） #

Hugging Face的datasets.Dataset类专门为Transformer训练设计，能直接被Trainer识别，只需将字典格式的数据集转换为Dataset对象即可。

步骤：

1. 安装datasets库（如果未安装）：

!pip install -q datasets

2. 导入Dataset并转换数据集：

from datasets import Dataset

# 将字典转换为Dataset对象
train_dataset = Dataset.from_dict(train_dataset)
val_dataset = Dataset.from_dict(val_dataset)

3. 额外注意：由于给GPT2添加了新的pad token，需要调整模型的embedding层大小，否则会触发维度不匹配错误：

# 在实例化模型后添加这行代码
model.resize_token_embeddings(len(tokenizer))

修改后的完整代码（关键部分） #

! pip install -q transformers datasets
from transformers import GPT2Tokenizer, GPT2LMHeadModel, Trainer, TrainingArguments
from datasets import Dataset

tokenizer = GPT2Tokenizer.from_pretrained('gpt2-medium')
tokenizer.add_special_tokens({'pad_token': '[PAD]'})
tokenized_data = tokenizer(dataset['FullDescription'].tolist(), truncation=True, padding=True)

# Split data into training and validation sets
train_size = int(0.8 * len(tokenized_data['input_ids']))
val_size = len(tokenized_data['input_ids']) - train_size

train_dataset = {'input_ids': tokenized_data['input_ids'][:train_size],
                 'attention_mask': tokenized_data['attention_mask'][:train_size]}
val_dataset = {'input_ids': tokenized_data['input_ids'][train_size:],
               'attention_mask': tokenized_data['attention_mask'][train_size:]}

# 转换为Dataset对象
train_dataset = Dataset.from_dict(train_dataset)
val_dataset = Dataset.from_dict(val_dataset)

from transformers import GPT2Config
model_config = GPT2Config.from_pretrained('gpt2-medium')
model_config.output_hidden_states = True
model = GPT2LMHeadModel.from_pretrained('gpt2-medium', config=model_config)
# 调整embedding层大小以适配新添加的pad token
model.resize_token_embeddings(len(tokenizer))

# 后续TrainingArguments和Trainer初始化不变
training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=1,
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='./logs',
    logging_steps=10,
    evaluation_strategy='steps',
    eval_steps=50,
    load_best_model_at_end=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=val_dataset
)

trainer.train()

方法2：自定义PyTorch Dataset类 #

如果不想使用datasets库，可以自定义一个继承自torch.utils.data.Dataset的类，实现__getitem__和__len__方法：

import torch
from torch.utils.data import Dataset

class TextDataset(Dataset):
    def __init__(self, input_ids, attention_mask):
        self.input_ids = input_ids
        self.attention_mask = attention_mask
    
    def __len__(self):
        return len(self.input_ids)
    
    def __getitem__(self, idx):
        return {
            'input_ids': torch.tensor(self.input_ids[idx]),
            'attention_mask': torch.tensor(self.attention_mask[idx])
        }

# 替换原来的数据集定义
train_dataset = TextDataset(
    tokenized_data['input_ids'][:train_size],
    tokenized_data['attention_mask'][:train_size]
)
val_dataset = TextDataset(
    tokenized_data['input_ids'][train_size:],
    tokenized_data['attention_mask'][train_size:]
)

同样需要添加model.resize_token_embeddings(len(tokenizer))这行代码。

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