如何计算HuggingFace T5模型生成MNLI分类标签的概率？

问题场景 #

使用HuggingFace的T5-Large（或T5-Small）进行MNLI任务推理：给定premise（前提）和hypothesis（假设），判断二者关系为entailment（蕴含）、neutral（中立）或contradiction（矛盾）。当前已能得到分类结果，但无法获取生成该结果对应的概率（如生成entailment时的P(entailment)）。

当前代码可生成分类结果，但不知如何利用generate()返回的scores计算概率，也不清楚softmax的具体操作方式。

解决方案 #

T5在MNLI任务中会生成单个类别词，我们可以通过generate()返回的scores（每一步的预测logits）结合softmax函数，计算每个类别的概率。具体步骤如下：

核心逻辑 #

1. 获取生成步骤的logits：generate()开启output_scores=True后，会返回每个生成步骤的logits张量，对应下一个token的预测得分。
2. 映射类别到token id：将三个MNLI类别词转换为T5 tokenizer对应的id。
3. softmax计算概率：对logits做softmax处理，将得分转换为0-1区间的概率值，再提取对应类别的概率。

修改后的完整代码 #

from transformers import T5Tokenizer, T5ForConditionalGeneration
import torch

def get_entailment_with_prob(premise, hypothesis):
    # 构建符合T5 MNLI任务要求的输入文本
    input_text = f"mnli premise: {premise} hypothesis: {hypothesis}"
    token_output = tokenizer(input_text, return_tensors="pt")
    input_ids = token_output.input_ids

    # 生成结果，开启scores输出，设置max_new_tokens=1（MNLI类别为单个词）
    output = model.generate(
        input_ids,
        output_scores=True,
        return_dict_in_generate=True,
        max_new_tokens=1,
        num_beams=1  # 保持贪心解码，与默认行为一致
    )

    # 解码生成的类别结果
    predicted_label = tokenizer.decode(output.sequences[0], skip_special_tokens=True)

    # 映射三个MNLI类别到对应的token id
    label_to_id = {
        "entailment": tokenizer.convert_tokens_to_ids("entailment"),
        "neutral": tokenizer.convert_tokens_to_ids("neutral"),
        "contradiction": tokenizer.convert_tokens_to_ids("contradiction")
    }

    # 获取生成第一步的logits（仅生成1个token，scores列表仅含1个元素）
    step_logits = output.scores[0]  # 形状: (1, 词汇表大小)

    # 对logits应用softmax，转换为概率分布
    prob_distribution = torch.nn.functional.softmax(step_logits, dim=-1)

    # 提取三个类别的概率值
    label_probs = {
        label: prob_distribution[0, token_id].item()
        for label, token_id in label_to_id.items()
    }

    return predicted_label, label_probs

# 加载模型和tokenizer
tokenizer = T5Tokenizer.from_pretrained('t5-small')
model = T5ForConditionalGeneration.from_pretrained('t5-small', return_dict=True)

# 测试输入
premise = "This game will NOT open unless you agree to them sharing your information to advertisers."
hypothesis = "Personal data disclosure is discussed."

# 获取结果和概率
result, probs = get_entailment_with_prob(premise, hypothesis)
print(f"预测类别: {result}")
print(f"各类别概率: {probs}")

代码说明 #

• max_new_tokens=1：MNLI的三个类别均为单个单词，设置该参数可限制生成长度，减少不必要的计算。
• label_to_id字典：直接关联类别词与token id，方便从logits中提取对应得分。
• softmax转换：将模型输出的logits（未归一化的得分）转换为符合概率分布的数值，dim=-1指定对词汇维度进行归一化。
• 返回结果：同时返回预测的类别和三个类别的概率字典，既得到分类结果，也能查看每个类别的置信度。

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