让小模型做选择题：限定候选项并读取 logits 概率的入门实践（附 Python 代码）

适合人群：会基本 Python，希望用本地小模型做分类、打标签、单选判断等任务，又苦于模型输出格式不稳定的同学。

一、背景

X 上有中文用户分享了一个思路：让大模型"吐 JSON"时，模型要逐个生成 token，不仅慢，格式还可能出错；而对于"从几个固定选项里选一个"的任务，可以只让 1.7B 级别的小模型做一次前向计算，然后只比较几个候选选项的概率，屏蔽词表里的其他词。作者还提到，在常识问答数据集上测试时，模型"嘴上说的确定程度"和实际准确率并不一致，需要做校准后才能信任置信度。

这个方法在工程上很实用：情感判断、工单分类、是否违规审核等任务，本质上都是选择题。本文用 Hugging Face Transformers 实现一个最小可运行的版本，并解释背后的原理。

二、原理：模型其实早就"算好了"所有答案的概率

语言模型每生成一个 token，都会先输出一个长度等于词表大小的向量，称为 logits。对它做 softmax，就得到"下一个 token 是词表中每个词"的概率。

平时我们让模型"生成"，是从这个分布中采样或取最大值，然后继续下一步。但如果答案只可能是 A、B、C、D 之一，我们完全可以：

构造好提示词，让下一个 token 恰好就是答案；
取出这个位置的 logits；
只看 A、B、C、D 对应的那几个值，在它们之间做 softmax。

这样做有三个好处：

快：只需要一次前向计算，不需要逐字生成；
格式稳定：结果必然是候选项之一，不会出现"我认为答案可能是......"这样的废话；
能拿到概率：可以根据概率设置阈值，低置信度的样本交给人工处理。

三、准备工作

bash 复制代码

pip install torch transformers

示例使用一个小型开源指令模型，CPU 也可以运行，只是速度慢一些。

四、操作步骤

步骤 1：加载模型

python 复制代码

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

name = "Qwen/Qwen2.5-0.5B-Instruct"
tok = AutoTokenizer.from_pretrained(name)
model = AutoModelForCausalLM.from_pretrained(name, torch_dtype="auto")
model.eval()

步骤 2：构造提示词

关键是让模型在回答的第一个 token 就输出选项字母：

python 复制代码

def build_prompt(question, options):
    lines = [f"{k}. {v}" for k, v in options.items()]
    content = question + "\n" + "\n".join(lines) + "\n只回答一个选项字母。"
    msgs = [{"role": "user", "content": content}]
    return tok.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False)

步骤 3：找到候选项对应的 token id

python 复制代码

letters = ["A", "B", "C", "D"]
cand_ids = [tok.encode(l, add_special_tokens=False)[0] for l in letters]
print(cand_ids)

注意：不同分词器对"A"和" A"（带空格）可能编码成不同 token。请打印出来确认每个字母都对应单个 token，并且与模型实际会生成的形式一致。

步骤 4：一次前向，读取概率

python 复制代码

@torch.no_grad()
def choose(question, options):
    prompt = build_prompt(question, options)
    inputs = tok(prompt, return_tensors="pt")
    logits = model(**inputs).logits[0, -1]      # 最后一个位置
    sub = logits[cand_ids]                      # 只保留候选项
    probs = torch.softmax(sub.float(), dim=-1)
    return {l: round(p.item(), 3) for l, p in zip(letters, probs)}

q = "用户评论："物流很快，但包装破损了。"这条评论的情感倾向是？"
opts = {"A": "正面", "B": "负面", "C": "中性或混合", "D": "无法判断"}
print(choose(q, opts))

输出是一个字典，例如每个选项对应一个概率值。取概率最高的选项作为答案即可。

步骤 5：批量处理

把待分类的文本放进列表循环调用即可。数据量大时，可以把多条提示词组成一个批次，注意设置 tok.padding_side = "left"，保证每条样本的"最后一个位置"都是真正的末尾。

五、关于"置信度"的重要提醒

拿到概率之后，很容易想当然地认为"概率 0.9 就代表 90% 会答对"。但正如原帖作者的实测所示，模型给出的概率往往并不校准：它可能表现得非常确定，实际准确率却低得多。

比较稳妥的做法是：

准备一个带标准答案的验证集；
把预测结果按概率分组，比如 0.9 以上、0.7～0.9、0.7 以下；
分别统计每组的实际准确率；
根据结果决定阈值，例如"只有概率高于某值的样本才自动处理"。

如果想更进一步，可以了解"温度缩放"（temperature scaling）等经典校准方法：在验证集上学习一个温度参数，让 logits 除以它之后再做 softmax，使概率更接近真实准确率。

六、常见问题 FAQ

Q1：选项超过 4 个怎么办？

同样适用，只要每个选项标签都能映射到单个 token。也可以用数字 1～9 作为标签。

Q2：选项标签被分词成了多个 token？

换用确定是单 token 的标签，或者改为计算整个选项序列的对数概率之和（需要对每个选项分别做前向计算）。

Q3：小模型准确率不够高？

可以在提示词里加入几条示例（少样本），或者用少量标注数据做微调。也可以把低置信度样本交给更大的模型复核。

Q4：这个方法能用在云端 API 上吗？

部分 API 提供返回候选 token 对数概率的参数，但能否限制候选范围、返回多少个候选，取决于具体服务，请查阅对应文档。

Q5：和让模型输出 JSON 相比如何选择？

需要多个字段、结构复杂时用结构化输出；只需要从固定选项中选一个时，本文的方法更快更稳。

七、实用技巧

选项顺序会影响结果：可以打乱顺序多测几次，观察是否存在位置偏好；
加一个"无法判断"选项：给模型留出拒答空间，减少硬猜；
记录概率而不只是答案：后续调整阈值时非常有用；
先人工标注一小批数据：没有验证集，就无法知道模型是否可靠；
保持提示词稳定：上线后修改提示词，需要重新做一次验证。

八、总结

把"生成答案"变成"比较候选项的概率"，是小模型落地分类任务的一个简单有效的技巧：构造提示词 → 取最后位置 logits → 只保留候选项做 softmax → 用验证集确定阈值。它快速、格式稳定，还能给出可用于分流的置信度，但别忘了先验证概率是否可信。

参考资料

Hugging Face Transformers 官方文档：模型输出与 logits 说明
论文：On Calibration of Modern Neural Networks（温度缩放方法）
CommonsenseQA 数据集论文
X 用户 @ataosky 关于小模型限定选项与置信度校准的分享