如何使用facebook/bart-large-mnli模型完成三分类NLI任务？

问题出在你手动筛选了logits的索引，只保留了第0和第2位（矛盾、蕴含），然后对这两个值做softmax，自然只能得到两个概率。要输出三分类（矛盾、中立、蕴含）的完整概率，只需要直接对模型输出的全部三个logits做softmax即可。

BART-large-MNLI的输出logits对应关系是：

• 索引0：矛盾（contradiction）
• 索引1：中立（neutral）
• 索引2：蕴含（entailment）
  刚好和你的标签[0,1,2]对应，不需要调整索引映射。

修正后的代码如下：

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer  # 补充缺失的导入

device    = torch.device("cuda" if torch.cuda.is_available() else "cpu")
nli_model = AutoModelForSequenceClassification.from_pretrained('facebook/bart-large-mnli')
tokenizer = AutoTokenizer.from_pretrained('facebook/bart-large-mnli')

nli_model.to(device)

i          = 0 # 验证第一个样本
premise    = tokenized_datasets['TRAIN'][i]['premise']
hypothesis = tokenized_datasets['TRAIN'][i]['hypothesis']

# 新版transformers已弃用truncation_strategy，改用truncation=True自动处理截断
x      = tokenizer.encode(premise, hypothesis, return_tensors='pt', truncation=True)
logits = nli_model(x.to(device))[0]

# 直接对全部三个logits做softmax，得到三分类概率
probs = logits.softmax(dim=1)
print(probs)

补充说明：

• 运行后会得到形状为[1,3]的张量，三个值分别对应矛盾、中立、蕴含的概率
• 如果需要获取预测的类别标签，可使用torch.argmax(probs, dim=1)得到概率最大的类别索引
• 注意代码中补充了transformers相关类的导入，避免运行报错

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