使用Huggingface Trainer训练时如何在模型forward()中获取自定义列？

核心原因 #

Huggingface Trainer默认只会自动传递预训练模型（比如Llama）原生forward方法包含的参数（input_ids、labels等），自定义的字段不会被自动传入，必须显式告诉Trainer要处理这些字段。

具体解决方案 #

1. 确认字段名完全匹配

确保数据集里的自定义列名就是interact_ids和candidate_ids，和模型forward方法里的参数名完全一致（大小写、拼写都不能错）。可以打印数据集样本验证：

print(dataset[0])  # 检查输出里是否存在目标自定义字段

2. 自定义DataCollator，包含自定义字段

默认的DataCollatorForLanguageModeling只会处理标准字段，你需要扩展它来收集自定义字段：

from transformers import DataCollatorForLanguageModeling
import torch

class CustomDataCollator(DataCollatorForLanguageModeling):
    def __call__(self, features):
        # 先处理input_ids、labels等标准字段
        batch = super().__call__(features)
        # 添加自定义字段，将样本列表转为tensor
        batch["interact_ids"] = torch.tensor([f["interact_ids"] for f in features])
        batch["candidate_ids"] = torch.tensor([f["candidate_ids"] for f in features])
        return batch

初始化Trainer时指定这个自定义collator：

data_collator = CustomDataCollator(tokenizer=tokenizer, mlm=False)  # mlm参数根据你的任务需求设置

trainer = Trainer(
    model=model,
    train_dataset=train_dataset,
    data_collator=data_collator,
    # 其他训练参数...
)

3. 修复forward方法的语法错误

你代码里的interact_ids.shape(0)是错误写法，应该用方括号interact_ids.shape[0]，否则即使拿到字段值也会抛出AttributeError。

额外处理（如果是变长序列）

如果interact_ids或candidate_ids是长度不一致的序列，需要在collator里做padding：

from torch.nn.utils.rnn import pad_sequence

batch["interact_ids"] = pad_sequence(
    [torch.tensor(f["interact_ids"]) for f in features],
    batch_first=True,
    padding_value=0  # 替换为你数据集对应的padding值
)

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