如何在transformers库中正确实现stopping_criteria参数配置

核心问题说明 #

你原有代码直接实例化StoppingCriteria是错误的，该类是抽象基类，必须继承后实现判定逻辑才能使用，且最终传入的参数必须是StoppingCriteriaList类型的实例。

完整实现代码 #

from transformers import pipeline, StoppingCriteria, StoppingCriteriaList
import torch

# 自定义停止条件判定类
class StopAtSpecificStr(StoppingCriteria):
    def __init__(self, stop_strs, tokenizer):
        self.stop_strs = stop_strs
        self.tokenizer = tokenizer
        # 提前将停止字符串转换为对应的token id列表，避免重复计算
        self.stop_token_list = [tokenizer.encode(s, add_special_tokens=False) for s in stop_strs]

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        # 遍历所有停止序列，检查当前生成内容末尾是否匹配
        for stop_tokens in self.stop_token_list:
            # 已生成内容长度小于停止序列长度时直接跳过
            if input_ids.shape[1] < len(stop_tokens):
                continue
            # 对比最后n个token和停止序列id，n为停止序列的token长度
            if torch.all(input_ids[0, -len(stop_tokens):] == torch.tensor(stop_tokens, device=input_ids.device)):
                # 匹配成功返回True，触发停止生成
                return True
        return False

# 初始化文本生成pipeline
generator = pipeline('text-generation', model='EleutherAI/gpt-neo-125M')

# 实例化自定义停止条件，传入停止字符串和pipeline对应的tokenizer
stop_condition = StopAtSpecificStr(stop_strs=["###"], tokenizer=generator.tokenizer)
# 将停止条件实例包装为StoppingCriteriaList，即使只有一个条件也需要包裹
stopping_criteria = StoppingCriteriaList([stop_condition])

# 调用生成接口传入参数
prompt = "请输入你的生成提示词"
res = generator(
    prompt,
    do_sample=True,
    stopping_criteria=stopping_criteria,
    max_new_tokens=200
)
print(res[0]["generated_text"])

注意事项 #

• __call__方法是必须实现的接口，返回True时就会终止生成
• 停止序列转换为token时不要添加特殊token（add_special_tokens=False），否则会出现匹配失败的问题
• 判定时要将对比的token张量放到和input_ids相同的设备上，避免GPU/CPU设备不匹配报错
• 所有停止条件必须放在StoppingCriteriaList列表中传入，不能直接传入自定义类的实例

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