如何让Hugging Face Pipeline调用facebook/opt模型仅输出单个单词？

要让模型直接输出单个预期下一词，关键是用**max_new_tokens**参数替代max_length，同时配合return_full_text=False来直接获取生成的内容，具体修改如下：

修改后的代码 #

from transformers import pipeline
generator = pipeline('text-generation', model="facebook/opt-1.3b")

generate_input = "This is a"
# 仅生成1个新token（对应单个下一词），且只返回新增内容
generate_output = generator(generate_input, max_new_tokens=1, return_full_text=False)
print(generate_output)
# 示例输出：[{'generated_text': 'test'}]

参数说明 #

• max_new_tokens=1：这个参数直接控制模型生成的新增token数量，设置为1就只会生成一个新token（绝大多数英文单词对应单个token，少数复杂词可能是多个，但这是最精准控制生成长度的方式）。相比你之前计算max_length的方式，它不需要考虑输入的长度，完全聚焦在新增内容上，还能避免长输入的警告。
• return_full_text=False：默认情况下pipeline会返回包含输入文本的完整内容，开启这个参数后，只会返回模型生成的新内容，省去你后续截取字符串的步骤。

为什么原来的方式不行？ #

你之前用max_length计算总长度的方式，会让模型生成远超过一个单词的内容，因为max_length是输入文本长度加上生成内容长度的总和，这样模型会持续生成直到达到总长度，既浪费计算资源，还需要额外处理结果。而max_new_tokens直接限制新增token数，从根源上减少了预测次数，降低计算耗时，完全符合你的需求。

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