前言
在大模型的生成過程中，部分原生的大語言模型未經過特殊的對齊訓練，往往會“胡說八道”的生成一些敏感詞語等用戶不想生成的詞語，最簡單粗暴的方式就是在大模型生成的文本之后，添加敏感詞庫等規則手段進行敏感詞過濾，但是在生成過程中，生成敏感詞仍然耗費了時間和算力成本。
本文以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 usingbeam search or log softmax for each vocabulary token when using beam searchkwargs (`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__).parametersif 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_listdef __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 torchclass new_logits_processor(LogitsProcessor):def __init__(self, forbid_token_id_list: List[int] = None):self.forbid_token_id_list = forbid_token_id_listdef __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:for id_ in self.forbid_token_id_list:scores[:, id_] = -float('inf')return scoresmodel_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_wordslogits_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
)
