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

前言

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

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

LogitsProcessor

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

python 复制代码
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控制大模型生成的过程呢?下面直接上实践代码:

python 复制代码
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详细实践代码:

python 复制代码
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

相关推荐
飞睿科技16 分钟前
乐鑫代理商飞睿科技,2025年AI智能语音助手市场发展趋势与乐鑫芯片解决方案分析
人工智能
许泽宇的技术分享18 分钟前
从新闻到知识图谱:用大模型和知识工程“八步成诗”打造科技并购大脑
人工智能·科技·知识图谱
亚马逊云开发者28 分钟前
全景解读亚马逊云科技的 GenBI 解决方案:三大路径助力企业智能决策升级
sql·llm
坤坤爱学习2.032 分钟前
求医十年,病因不明,ChatGPT:你看起来有基因突变
人工智能·ai·chatgpt·程序员·大模型·ai编程·大模型学
蹦蹦跳跳真可爱5891 小时前
Python----循环神经网络(Transformer ----注意力机制)
人工智能·深度学习·nlp·transformer·循环神经网络
空中湖3 小时前
tensorflow武林志第二卷第九章:玄功九转
人工智能·python·tensorflow
lishaoan773 小时前
使用tensorflow的线性回归的例子(七)
人工智能·tensorflow·线性回归
千宇宙航6 小时前
闲庭信步使用SV搭建图像测试平台:第三十一课——基于神经网络的手写数字识别
图像处理·人工智能·深度学习·神经网络·计算机视觉·fpga开发
onceco7 小时前
领域LLM九讲——第5讲 为什么选择OpenManus而不是QwenAgent(附LLM免费api邀请码)
人工智能·python·深度学习·语言模型·自然语言处理·自动化
天水幼麟7 小时前
动手学深度学习-学习笔记(总)
笔记·深度学习·学习