如何使用基于STS模型微调的Textual Entailment模型进行分类预测？

我正在微调STS模型，用于文本蕴含的三类分类任务，类别对应关系为：0(neutral/中立)、1(entailment/蕴含)、2(contradiction/矛盾)。参考NLI训练逻辑完成模型训练后，需要实现三类分类的预测功能，以下是对原参考代码的验证和修正方案：

原代码问题分析 #

你提供的参考代码仅保留了蕴含和矛盾类的logits，丢弃了中立类，不符合三类分类的需求；同时使用了已废弃的truncation_strategy参数，且未明确类别索引的对应关系，需要针对性调整。

修正后的预测代码 #

from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

# 替换为你训练后模型的保存路径
model_path = "你的模型保存路径"
nli_model = AutoModelForSequenceClassification.from_pretrained(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)

# 设备配置，自动使用GPU（如果可用）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
nli_model.to(device)
nli_model.eval()  # 切换到评估模式

# 示例输入
premise = "我口渴了"
hypothesis = "我想要喝水"

# 编码输入：使用truncation=True替代已废弃的truncation_strategy
inputs = tokenizer(
    premise,
    hypothesis,
    return_tensors="pt",
    truncation=True,
    max_length=512  # 可根据模型最大输入长度调整
)

# 将输入转移到指定设备
inputs = {k: v.to(device) for k, v in inputs.items()}

# 前向传播，获取logits
with torch.no_grad():  # 评估阶段关闭梯度计算
    outputs = nli_model(**inputs)
logits = outputs.logits

# 计算三类的概率
probs = torch.softmax(logits, dim=1)
# 获取预测类别索引（0=中立，1=蕴含，2=矛盾）
predicted_class_idx = torch.argmax(probs, dim=1).item()

# 类别映射
label_map = {0: "neutral(中立)", 1: "entailment(蕴含)", 2: "contradiction(矛盾)"}
predicted_label = label_map[predicted_class_idx]

# 输出结果
print(f"前提：{premise}")
print(f"假设：{hypothesis}")
print(f"预测类别：{predicted_label}")
print(f"各类概率：中立={probs[0][0]:.4f}, 蕴含={probs[0][1]:.4f}, 矛盾={probs[0][2]:.4f}")

关键注意事项 #

• 类别索引对应：确保预测时的类别索引和训练时的标签映射完全一致（训练时需将SNLI数据集的标签按需求映射为0=中立、1=蕴含、2=矛盾）
• 参数废弃处理：truncation_strategy已被Hugging Face Transformers库弃用，改用truncation=True会自动处理过长序列
• 评估模式切换：预测前必须调用model.eval()，关闭Dropout等训练专属层，保证预测结果稳定
• 梯度关闭：使用torch.no_grad()避免不必要的内存消耗，加快预测速度

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