大语言模型控制生成的过程Trick：自定义LogitsProcessor实践

前言

在大模型的生成过程中，部分原生的大语言模型未经过特殊的对齐训练，往往会“胡说八道”的生成一些敏感词语等用户不想生成的词语，最简单粗暴的方式就是在大模型生成的文本之后，添加敏感词库等规则手段进行敏感词过滤，但是在生成过程中，生成敏感词仍然耗费了时间和算力成本。

本文以chatglm2-6B为例，通过自定义LogitsProcessor，实践大模型在生成过程中控制一些词语的生成。

LogitsProcessor

从下面代码可以看到，LogitsProcessor的作用就是在生成过程中修改score，改变模型输出的概率分布的工具。

class LogitsProcessor:
    """Abstract base class for all logit processors that can be applied during generation."""

    @add_start_docstrings(LOGITS_PROCESSOR_INPUTS_DOCSTRING)
    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        raise NotImplementedError(
            f"{self.__class__} is an abstract class. Only classes inheriting this class can be called."
        )
        
class LogitsProcessorList(list):
    """
    This class can be used to create a list of [`LogitsProcessor`] or [`LogitsWarper`] to subsequently process a
    `scores` input tensor. This class inherits from list and adds a specific *__call__* method to apply each
    [`LogitsProcessor`] or [`LogitsWarper`] to the inputs.
    """

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> torch.FloatTensor:
        r"""
        Args:
            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
                Indices of input sequence tokens in the vocabulary. [What are input IDs?](../glossary#input-ids)
            scores (`torch.FloatTensor` of shape `(batch_size, config.vocab_size)`):
                Prediction scores of a language modeling head. These can be logits for each vocabulary when not using
                beam search or log softmax for each vocabulary token when using beam search
            kwargs (`Dict[str, Any]`, *optional*):
                Additional kwargs that are specific to a logits processor.

        Return:
            `torch.FloatTensor` of shape `(batch_size, config.vocab_size)`:
                The processed prediction scores.

        """
        for processor in self:
            function_args = inspect.signature(processor.__call__).parameters
            if len(function_args) > 2:
                if not all(arg in kwargs for arg in list(function_args.keys())[2:]):
                    raise ValueError(
                        f"Make sure that all the required parameters: {list(function_args.keys())} for "
                        f"{processor.__class__} are passed to the logits processor."
                    )
                scores = processor(input_ids, scores, **kwargs)
            else:
                scores = processor(input_ids, scores)
        return scores

自定义LogitsProcessor实践

回到正题，如何自定义LogitsProcessor控制大模型生成的过程呢？下面直接上实践代码：

class new_logits_processor(LogitsProcessor):
    def __init__(self, forbid_token_id_list: List[int] = None):
        self.forbid_token_id_list = forbid_token_id_list

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        for id_ in self.forbid_token_id_list:
            scores[:, id_] = -float('inf')
        return scores

forbid_token_id_list是不让模型生成词语的id映射列表，对于这些抑制生成的词语，在自定义logits_processor时将其概率推向负无穷大即可。

chatglm2-6B详细实践代码：

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, TextStreamer
from transformers.generation.logits_process import LogitsProcessor, LogitsProcessorList
from typing import List
import torch


class new_logits_processor(LogitsProcessor):
    def __init__(self, forbid_token_id_list: List[int] = None):
        self.forbid_token_id_list = forbid_token_id_list

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
        for id_ in self.forbid_token_id_list:
            scores[:, id_] = -float('inf')
        return scores


model_path = "THUDM/chatglm2-6b"
tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
model = AutoModelForSeq2SeqLM.from_pretrained(model_path, trust_remote_code=True).to('mps')


def add_forbid_words():
    '''
    添加需要抑制的词语，这里简单添加了数字和几个词语进行对比
    :return:list
    '''
    forbid_words = []
    for i in range(10):
        forbid_words.append(tokenizer.convert_tokens_to_ids(str(i)))
    forbid_words.append(tokenizer.convert_tokens_to_ids("首先"))
    forbid_words.append(tokenizer.convert_tokens_to_ids("积极"))
    forbid_words.append(tokenizer.convert_tokens_to_ids("回答"))
    forbid_words.append(tokenizer.convert_tokens_to_ids("勇敢"))
    forbid_words.append(tokenizer.convert_tokens_to_ids("勇气"))
    return forbid_words


logits_processor = LogitsProcessorList()
logits_processor.append(new_logits_processor(add_forbid_words()))

streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)

input = "列举出10个积极的词语："

outputs = model.generate(
    tokenizer(input, return_tensors='pt').input_ids.to("mps"),
    max_new_tokens=1024,
    logits_processor=logits_processor,  # 不开启注释即可
    streamer=streamer
)
decode_text = tokenizer.batch_decode(outputs, streamer=streamer)[0]
print(decode_text)

抑制前输出：

1. 勇敢
2. 快乐
3. 成功
4. 努力
5. 积极
6. 乐观
7. 自信
8. 开朗
9. 团结
10. 奋斗

抑制后输出：

- 积极主动
- 乐观向上
- 自信
- 自律
- 诚实守信
- 乐于助人
- 勇于尝试
- 坚韧不拔
- 乐观开朗
- 团结一心

小结

本文通过自定义LogitsProcessor，简单的实践了大语言模型在生成过程中屏蔽生成用户自定义词语的trick。在现实场景中，根据特定场景探索如何灵活的利用LogitsProcessor进行有针对性的控制生成模型的生成过程非常重要。

参考文献

【1】https://github.com/huggingface/transformers/blob/v4.31.0/src/transformers/generation/logits_process.py