训练Hugging Face Transformer自定义数据集时触发随机KeyError报错

报错核心是Pandas索引不匹配，和Transformer本身逻辑无关：

• 你用train_test_split切分得到的X_train/X_test、y_train/y_test默认是Pandas Series对象，切分后会保留原DataFrame的行标签索引，不是从0开始连续递增的位置索引。
• 你后续手动遍历编码结果组装data_to_add时，是按遍历顺序往列表追加元素，相当于把编码结果按0、1、2…的连续位置存储，但传入Dataset的y_train/y_test仍保留原Pandas索引。
• PyTorch DataLoader取数据时会传入连续位置的idx值（比如0、1、2…19），当你用self.labels[idx]取值时，Pandas会按标签索引匹配而非位置匹配，若当前切分结果的索引中不存在传入的idx值，就会抛出KeyError。因为每次运行train_test_split的切分结果随机，所以报错的索引数字每次都会变化。
• 额外问题：你原先单条文本调用tokenizer加padding的写法有误，单条文本padding只会pad到当前文本自身长度，会导致后续训练时批次内张量维度不统一报错；且你设置的max_length=10过短，会截断绝大多数有效文本信息。

1. 切分数据集时直接将文本、标签转为原生Python列表，彻底规避Pandas索引匹配问题
2. 修正tokenizer调用逻辑，直接传入文本列表做批量编码，统一padding规则，将max_length调整为BERT支持的合理长度（建议512，可根据你的文本长度调整）
3. 删除冗余的手动重组encodings字典的代码，Hugging Face tokenizer批量处理后直接输出符合Dataset要求的字典格式

修正后的tokenize_dataset方法代码如下：

def tokenize_dataset(self):
    tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
    # 切分后直接转原生列表，消除Pandas索引影响
    X_train, X_test, y_train, y_test = train_test_split(
        self.pdfs["text_clean"].tolist(), 
        self.pdfs["labels"].tolist(),
        test_size=0.2,
        random_state=42 # 加随机种子固定切分结果，方便复现
    )

    # 批量编码文本，自动统一padding长度
    train_encodings = tokenizer(X_train, truncation=True, padding=True, max_length=512)
    test_encodings = tokenizer(X_test, truncation=True, padding=True, max_length=512)
    
    # 直接实例化数据集，不需要后续手动重组编码字段
    self.train_dataset = PDFsDataset(train_encodings, y_train)
    self.eval_dataset = PDFsDataset(test_encodings, y_test)

你原先写的PDFsDataset自定义类逻辑正确，不需要修改。

• 你提供的JSON样例存在语法错误（字符串缺引号、数组元素缺逗号），实际使用时需要保证JSON格式合法，正确格式参考：

{
    "text_clean": [
        "article with a few hundred words",
        "another article with a lot of words",
        "yet another article"
    ],
    "most_similar_label":[
        "Quantum",
        "Artificial intelligence",
        "Materials"
    ]
}

• 你的import_from_json方法中用到了new_tags字段，但给出的JSON样例中不存在该字段，实际运行时需要保证输入JSON包含该字段，或增加字段存在性判断逻辑，避免触发字段不存在的KeyError。

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