以Llama-2为例，在生成模型中使用自定义LogitsProcessor

1. 前言

在上一篇文章 以Llama-2为例，在生成模型中使用自定义StoppingCriteria中，介绍了怎样在生成的过程中，使用stopping criteria来控制生成过程的结束，本文将继续这一话题，结合具体的场景，介绍如何实现自定义的logits processor，并以此来控制生成的过程。

2. 场景介绍

场景延续上篇介绍stopping criteria的文章，假如我们希望使用Llama-2模型，来生成一篇新闻的概要，希望它能够生成一句简短的话，来描述这篇新闻中主要发生了什么。

在上一篇文章中，我们成功的使用stopping criteria解决了模型废话太多的问题，然而，在某些情况下，模型输出的结果并不是我们想要的，它没有用一句话概括，反而是一条一条列举了其中的主要信息，类似：

1. ......
2. ......
3. ......

针对这种情况，我们可以强制要求生成的第一个token，不可以是数字，这样的话，就只能从字母中选择合适的单词生成，也就达到我们的目的了。为了实现这一策略，就需要用到logits processor。

3. 解决方法

logits processor是在生成的过程中，每一个step的score计算完成之后，对score进行进一步的加工，改变模型输出的概率分布，从而影响后续生成的结果。

transformers模块中提供了若干内置的processor可以直接调用，具体的整理和简介可以参考之前的文章以beam search为例，详解transformers中generate方法（上）。

现在我们需要设计这样一个processor，判断如果是第一个生成的第一个token，则禁止它生成数字，也就是把所有数字对应的得分强制设置为负无穷。

首先，引入需要用到的类，与stopping criteria类似的，也是有要给基础类，和一个容器类：

from transformers.generation.logits_process import LogitsProcessor, LogitsProcessorList

然后继承基础类，实现我们所需的processor：

class SuppressSpecificBOSTokenLogitsProcessor(LogitsProcessor):
    """
    防止生成的第一个token是某些特定的token
    ---------------
    ver: 2023-08-02
    by: changhongyu
    """
    def __init__(self, bad_bos_token_id_list: List[int] = None):
        """
        :param bad_bos_token_id_list: 不可以作为第一个token的token的id列表
        """
        self.bad_bos_token_id_list = bad_bos_token_id_list
    
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        new_token_len = input_ids.shape[-1] - current_token_len
        if new_token_len == 0:
            for id_ in self.bad_bos_token_id_list:
                scores[:, id_] = -float('inf')
        return scores

logits processor的使用方法与stopping criteria是一样的，我们设计好自己的processor类之后，实例化一个容器，再将实例化的processor放到这个容器中就好了：

NUMBER_ID_LIST = []
for i in range(10):
    NUMBER_ID_LIST.append(tokenizer.convert_tokens_to_ids(str(i)))
logits_processor = LogitsProcessorList()
logits_processor.append(SuppressSpecificBOSTokenLogitsProcessor(NUMBER_ID_LIST))

如果有多个processor的话，可能需要注意一下放入容器的顺序。

最后在生成的时候，将它作为参数传给generate方法就好了。

例如，原本生成的代码是：

outputs = model.generate(**inputs)

使用processor的话，可以写作：

outputs = model.generate(logits_processor=logits_processor, **inputs)

注意在实现的时候有一个小细节，由于是对话模型，输入的除了当前的query之外，还包括历史的对话记录，二者拼接在一起才是完整的prompt（prompt构建参考这一篇），所以我们并不能仅仅根据当前输入input_ids的长度，来判断当前step是不是这一轮生成的第一个token，这就是为什么上面的代码中有一个为声明定义的变量current_token_len。

对于这个current_token_len，只需要在model.generate执行之前，对他global一下就可以了。

例如像这个样子，每次生成之前先计算一下截至生成之前的长度：

global current_token_len
current_token_len = inputs['input_ids'].shape[1]

outputs = model.generate(logits_processor=logits_processor, **inputs)

4. 结语

作为用户控制生成过程的主要手段，如何巧妙地利用好logits processor对使用生成式模型来说非常重要。在实际情况中，需要针对场景，发现其中地规律，然后又针对性地去设计一个processor。它主要解决的问题，是一些有规律可循的场景，从一定意义上理解，可以认为是对生成模型的解空间进行了限制和变换。在解决问题的风格上给人的感觉，有点像抽取式模型所做的风格了，比如对于一个关键词生成任务，如果我们不希望模型生成文章中没有出现过的token，那完全可以利用本文中类似的方法，把生成结果限定为文中出现过的token。

以上就是本文的全部内容，如果对你有所帮助或启发，记得留下一个免费的赞，我们下期再见。