大语言模型控制生成的过程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

相关推荐
海森大数据1 分钟前
三步破局:一致性轨迹强化学习开启扩散语言模型“又快又好”推理新时代
人工智能·语言模型·自然语言处理
Tencent_TCB3 分钟前
云开发CloudBase AI+实战:快速搭建AI小程序全流程指南
人工智能·ai·小程序·ai编程·云开发
Sunhen_Qiletian5 分钟前
基于OpenCV与Python的身份证号码识别案例详解
人工智能·opencv·计算机视觉
AustinCyy11 分钟前
【论文笔记】Introduction to Explainable AI
论文阅读·人工智能
岁月宁静29 分钟前
在富文本编辑器中封装实用的 AI 写作助手功能
前端·vue.js·人工智能
末世灯光30 分钟前
时间序列入门第一问:它和普通数据有什么不一样?(附 3 类典型案例)
人工智能·python·机器学习·时序数据
Yann-企业信息化35 分钟前
AI 开发工具对比:Dify 与 Coze Studio(开源版)差异对比
人工智能·开源
2401_8369003339 分钟前
YOLOv4:集大成的目标检测王者
人工智能·yolov4
Xi xi xi43 分钟前
苏州唯理科技近期也正式发布了国内首款神经腕带产品
大数据·人工智能·经验分享·科技
www.0243 分钟前
微信克隆人,聊天记录训练专属AI(2.WeClone训练模型)
人工智能·python·微信·聊天克隆人·微信克隆人