基于facebook/bart-large-mnli微调模型的Pipeline部署报错排查

1. 修复"entailment"标签ID映射报错 #

用facebook/bart-large-mnli微调自定义14分类时，原模型的3分类标签映射会被覆盖，导致zero-shot pipeline找不到"entailment"标签ID。解决办法是手动在模型config中补充必要的标签映射：

from transformers import BartForSequenceClassification

# 加载预训练模型
model = BartForSequenceClassification.from_pretrained(
    "facebook/bart-large-mnli",
    num_labels=14,
    ignore_mismatched_sizes=True
)

# 手动配置label2id，必须包含entailment/contradiction/neutral（zero-shot逻辑依赖）
model.config.label2id = {
    "entailment": 0,
    "contradiction": 1,
    "neutral": 2,
    # 追加你的14个自定义标签，示例如下
    "自定义标签1": 3,
    "自定义标签2": 4,
    # ... 剩余标签依次映射到对应ID
}
# 生成反向映射id2label
model.config.id2label = {v: k for k, v in model.config.label2id.items()}

2. 解决GPU/CPU设备不匹配的RuntimeError #

报错原因是pipeline内部部分张量在CPU，模型在GPU，导致设备不一致。两种解决方式：

• 方式一：初始化pipeline时指定GPU设备
  直接让pipeline把所有组件加载到GPU：

from transformers import pipeline

# device=0对应第一块GPU，多卡环境可调整编号
classifier = pipeline(
    "zero-shot-classification",
    model=model,
    tokenizer="facebook/bart-large-mnli",
    device=0
)

• 方式二：手动强制模型和张量移至GPU
  如果方式一无效，手动将模型移到GPU后再创建pipeline：

model = model.to("cuda:0")
classifier = pipeline(
    "zero-shot-classification",
    model=model,
    tokenizer="facebook/bart-large-mnli",
    device=0
)

3. 解决save_model后加载无训练效果的问题 #

trainer.save_model默认只保存权重和基础config，会丢失自定义标签映射等关键配置。正确保存/加载流程：

保存模型 #

# 保存模型、完整config和tokenizer到本地目录
trainer.save_model("./fine_tuned_bart")
tokenizer.save_pretrained("./fine_tuned_bart")

加载模型 #

from transformers import BartForSequenceClassification, AutoTokenizer, pipeline

# 从本地目录加载完整模型和tokenizer
model = BartForSequenceClassification.from_pretrained("./fine_tuned_bart")
tokenizer = AutoTokenizer.from_pretrained("./fine_tuned_bart")

# 确认label2id包含entailment映射（若保存时已配置则可跳过）
model.config.label2id = {
    "entailment": 0,
    "contradiction": 1,
    "neutral": 2,
    # 你的自定义标签映射
    "自定义标签1": 3,
    # ...
}
model.config.id2label = {v: k for k, v in model.config.label2id.items()}

# 创建GPU环境的pipeline
classifier = pipeline(
    "zero-shot-classification",
    model=model,
    tokenizer=tokenizer,
    device=0
)

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