如何让HuggingFace Transformer翻译返回N条排序后的推理结果？

一、修改T5模型生成代码 #

要返回N条按置信度排序的命令结果，核心是调整generate方法的参数，获取每个候选的置信度得分，再排序后输出。以下是修改后的代码段：

predictions = []
# 配置生成多候选的参数，合并到generate_kwargs中
generate_kwargs.update({
    "num_return_sequences": N,  # N为你需要的候选数量
    "num_beams": N,             # beam数量与候选数一致，保证生成不同结果
    "return_scores": True,      # 返回每个候选的生成得分
    "early_stopping": True,
    "return_dict_in_generate": True  # 以字典格式返回结果，方便提取得分
})

for batch in dataset:
    input_tensor = self.to_tensor(batch["inputs"])
    # 生成候选序列及对应得分
    generate_output = self._model.generate(input_ids=input_tensor, **generate_kwargs)
    predicted_sequences = generate_output.sequences
    scores = generate_output.scores
    
    batch_predictions = []
    for seq_idx in range(N):
        # 解码单条候选序列为命令文本
        seq = predicted_sequences[seq_idx::N]  # 按输入分组提取候选
        pred_text = vocabs["targets"].decode(seq.cpu().numpy().tolist()[0])
        # 提取对应得分（T5得分是负对数概率，越接近0置信度越高）
        confidence = scores[seq_idx].item()
        batch_predictions.append((pred_text, confidence))
    
    # 按置信度降序排序
    batch_predictions.sort(key=lambda x: x[1], reverse=True)
    # 提取排序后的文本，也可保留(文本, 置信度)对用于后续处理
    sorted_texts = [item[0] for item in batch_predictions]
    predictions.extend(sorted_texts)

# 日志输出每个输入的所有候选结果
for inp_idx, inp in enumerate(inputs):
    start = inp_idx * N
    end = start + N
    input_preds = predictions[start:end]
    logging.info("%s\n  -> %s", inp, "\n  -> ".join(input_preds))

if output_file is not None:
    # 按输入分组写入，每组候选用制表符分隔
    grouped_preds = [predictions[i*N:(i+1)*N] for i in range(len(inputs))]
    utils.write_lines_to_file(["\t".join(preds) for preds in grouped_preds], output_file)

关键参数说明： #

• num_return_sequences: 指定返回的候选数量
• num_beams: 使用beam search生成多样化候选，数量与候选数一致可避免重复结果
• return_scores + return_dict_in_generate=True: 用于获取每个候选的生成得分，得分越接近0，模型对该候选的置信度越高
• 排序逻辑：直接按得分降序排列，符合置信度从高到低的需求

二、自然语言转命令的模型建议 #

针对你40种固定无嵌套命令的场景，除了原版T5，还可以考虑以下模型：

• FLAN-T5: 基于T5的指令微调版本，对文本到指令类任务适配性更强，微调后生成的命令更贴合需求
• DistilBERT: 将任务转为文本分类，把40种命令作为分类标签，直接输出概率排序的标签结果，推理速度快，适合固定类别的场景
• GPT-2/GPT-Neo: 小体量生成式模型，微调后可生成多候选命令，再过滤到你的命令集合中，适合需要一定灵活性的场景
• DistilT5: 轻量化T5模型，推理速度比原版T5快60%，精度损失小，适合部署到资源有限的环境中

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