谈谈 Tokenization 如何注入攻击防御，防止原语料中存在 special tokens 被直接编码

前言

在自然语言处理（NLP）和机器学习领域，特别是在使用预训练的大语言模型（如BERT、GPT、Llama等）前，往往需要用特定的 tokenizer ，将原始语料文本分解成一个个 tokens ，以让模型理解，这一过程被称为 tokenization，tokenization 是文本预处理的关键步骤，它影响着模型的性能，"special tokens" 是指一些具有特殊意义的 tokens，它们不对应于实际的单词或短语，但在模型的架构和处理流程中扮演重要角色。例如，下面是一些常见的特殊标记及其用途：

作为术语的“tokenization”、tokens (也译为标记)、tokenizer (也译为分词器) 在中文中尚无共识的概念对应，本文采用英文表达以利说明。

1. 开始（Start）和结束（End）:

   • [CLS]: 在BERT等模型中，用于表示序列的开始。通常用于分类任务中的输入标记。
   • [SEP]: 用于分隔两个句子或段落。在处理两个连续文本片段时，如在问答任务中分隔问题和上下文。
   • [EOS] 或 [PAD]: 用于表示序列的结束或进行填充以达到固定长度。[EOS]代表结束符，而[PAD]用于填充序列以确保输入长度一致。

2. 未知（Unknown）标记:

   • [UNK]: 用于表示模型词汇表中不存在的单词。这有助于模型处理未知或罕见词汇。

3. 掩码（Mask）标记:

   • [MASK]: 在训练过程中用于遮盖（masking）某些标记，迫使模型学习预测这些被遮盖的标记。这在BERT等模型的预训练任务中非常常见。

4. 特殊分类（Special Category）标记:

   • [CLS], [SEP], [PAD], [UNK], [MASK] 等也可以被视为特殊分类的标记，因为它们属于模型预定义的特殊类别，而不是普通的词汇。

这些特殊标记对于模型的训练和推理至关重要，因为它们提供了关于输入数据结构和任务类型的重要信息。例如，在BERT模型中，[CLS]标记的输出向量通常用于分类任务，而[SEP]标记有助于模型理解两个文本片段的分隔。通过这些特殊标记，模型能够更好地理解和处理各种NLP任务。

问题引入

最近在复现论文 Data Engineering for Scaling Language Models to 128K Context ，这篇论文对 slimpajama 数据集进行一个长文本的上采样，并按照各领域固定比例的方式组合文本长度，组件了一个新的数据集 yaofu/slimpajama-per-source-length-upsample 。

原论文中测试的模型是 Llama 2，因此，保存的数据集是已经经过 Llama 2 的 tokenizer 进行 tokenization 处理的 tokens 数据集，但我需要测试的模型是 Qwen-7B，这个模型的 tokenizer 并不是 Llama 2 的 tokenizer，显然，在用这个数据集之前，需要对这个数据集进行一个 decode。

另外一个重要的细节是，注意原数据集中的 special tokens，在 Llama2 的 tokenizer 模型中，设定的 special tokens 包括：开始标志 <s> 结束标志 </s> 和未知标志 <unk> 。这与 Qwen-base 中仅有的 special tokens 结束标志 <|endoftext|> 并不相同，当我以为只需要将 decode 后的数据进行一个 replace 处理即可时，我突然发现在原始语料文本中就存在这样的 special tokens，例如：

I'm making it to challenge my ambition and all that, so I won't be making just boring, samey stuff.
I want to focus on my site, my comics, paintings, just... things with more put into them.
700 Watcher Contest/Raffle!!!!!!! (50+ PRIZES) <s>500</s> 700 Watcher Contest!!!!
Official Group Rules! (Updated 11/26/2017)

这里的 <s> 与 </s> 我并不理解意义在哪里，但显然，这不是文本的开始标记和结束标记。

<p>但是在<s>JDK 6u132, JDK 7u122, JDK 8u113</s> JDK 6u141, JDK 7u131, JDK 8u121 中Java提升了JNDI 限制了Naming/Directory服务中JNDI Reference远程加载Object Factory类的特性。系统属性 com.sun.jndi.rmi.object.trustURLCodebase、com.sun.jndi.cosnaming.object.trustURLCodebase 的默认值变为false,即默认不允许从远程的Codebase加载Reference工厂类。如果需要开启 RMI Registry 或者 COS Naming Service Provider的远程类加载功能,需要将前面说的两个属性值设置为true。</p>

这一段语料来自一段 html 代码，我本以为 <s> 与 </s> 并不是 html 的特殊格式，结果发现：在HTML中，<s> 和 </s> 标签是一对用来表示删除线的标签，它们会使得被它们包围的文本中间显示一条横线，视觉上表示文本被删除或不再有效。这个标签通常用于编辑文档时，需要标记某些内容已被删除的情况。然而，需要注意的是，<s> 标签在HTML5中是不推荐的（deprecated），这意味着它可能在未来的HTML标准中被移除。取而代之的是 del 标签。

理论上，输入文本中不包含特殊token，它们仅在 tokenization 后由开发者手动加入。 但有些输入文本中含有特殊token的字面表达，这种情况很常见，但是显然这里不应该不做处理直接 encode，否则将会导致模型在训练过程中读入数据时，会将数据从异常的地方切割，这将产生很严重的问题。那么该如何解决这个问题呢？

基于 tiktoken 的 tokenizer 处理

参考 Tokenization

在这里，我们基于 Qwen 的 tokenizer 进行实验，Qwen-7B 采用 UTF-8 字节级别的 BPE tokenization 方式，并依赖tiktoken 这一高效的软件包执行分词。对于特殊 token，Qwen-7B中有 <|endoftext|> 由于特殊token和普通token概念上的差异，如果输入文本中含有特殊token的字面表达该如何处理？ 以下面文本为例：

print("<|endoftext|>")

其正确的 tokenization 为

ids:[1350, 9639, 91, 8691, 723, 427, 91, 82598]
tokens: [b'print', b'("<', b'|', b'endo', b'ft', b'ext', b'|', b'>")']

不是

ids: [1350, 445, 151643, 899]
tokens: [b'print', b'("', '<|endoftext|>', b'")']

注：<|endoftext|> 在词表中对应的 id 是 151643，第一种正确的原因在于将这个词拆开做 tokenization

如需启用注入攻击防御，以避免将原始语料中的 special tokens 直接做 tokenization，可以传入参数allowed_special=set()：

>>> tokenizer('print("<|endoftext|>")', allowed_special=set())
{'input_ids': [1350, 9639, 91, 8691, 723, 427, 91, 82598], 'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1]}

>>> tokenizer.encode('print("<|endoftext|>")', allowed_special=set())
[1350, 9639, 91, 8691, 723, 427, 91, 82598]

这里的参数 allowed_special 代表包含允许做 tokenization 的 special tokens 的集合，当你不允许任何输入文本中的 special tokens 做 tokenization 时，这里可以直接传入一个空的集合，当你允许部分 special tokens 做 tokenization 时，可以通过下面的方式：

>>> tokenizer('print("<|extra_0|>")<|endoftext|>', allowed_special={'<|endoftext|>'})
{'input_ids': [1350, 9639, 91, 15460, 62, 15, 91, 82598, 151643], 'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0, 0], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1]}

>>> tokenizer.encode('print("<|extra_0|>")<|endoftext|>', allowed_special={'<|endoftext|>'})
[1350, 9639, 91, 15460, 62, 15, 91, 82598, 151643]

如果希望输入中遇到特殊 token 的字面表达时，获得更直接的提醒，通过配置 disallowed_special 可以让tokenizer直接触发异常。

>>> tokenizer('print("<|extra_0|>")<|endoftext|>', allowed_special={'<|endoftext|>'}, disallowed_special=('<|extra_0|>', ))
...
ValueError: Encountered text corresponding to disallowed special token '<|extra_0|>'.
If you want this text to be encoded as a special token, pass it to `allowed_special`, e.g. `allowed_special={'<|extra_0|>', ...}`.
If you want this text to be encoded as normal text, disable the check for this token by passing `disallowed_special=(enc.special_tokens_set - {'<|extra_0|>'})`.
To disable this check for all special tokens, pass `disallowed_special=()`.

更多关于allowed_special和disallowed_special的信息, 请参阅tiktoken代码.

对于默认的 tokenization 操作，将允许任何 special tokens 被 tokenization，这篇论文的原作者在做 tokenization 时，采取了这种默认的行为，但这并不合适：

>>> tokenizer('print("<|endoftext|>")', allowed_special="all", disallowed_special=())
{'input_ids': [1350, 445, 151643, 899], 'token_type_ids': [0, 0, 0, 0], 'attention_mask': [1, 1, 1, 1]}

基于 transformers 的 tokenizer 处理

如果在尝试了上面的方法后，触发了异常提醒：

TypeError: PreTrainedTokenizerFast._batch_encode_plus() got an unexpected keyword argument 'allowed_special'

那就说明这个 tokenizer 的底层实现不是基于 tiktoken 的，因为我在实验时用的是 Qwen1.5，而我没有注意到上面的这个文档是 Qwen 1.0 的，Qwen 团队在切换版本时也切换了 tokenizer 的底层实现，作者在 issue #219 中提到：

Qwen2Tokenizer follows the practice of transformers and uses split_special_tokens to control the behaviour of the tokenizer regrading to special/control tokens.

那么我们可以根据 PreTrainedTokenizerFast 直接在 from_pretrained 时，指定参数 split_special_tokens=True :

>>> from transformers import AutoTokenizer

>>> tokenizer = AutoTokenizer.from_pretrained(path_to_model, split_special_tokens=True, use_fast=False)

>>> tokenizer("<|im_start|>This is a test.<|im_end|><|endoftext|>") # safe tokenization
{'input_ids': [27, 91, 318, 4906, 91, 29, 1986, 374, 264, 1273, 15757, 91, 318, 6213, 91, 1784, 91, 8691, 723, 427, 91, 29], 'attention_mask': [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]}

>>> tokenizer.convert_tokens_to_ids("<|im_start|>") # get special token ids manually
151644

>>> tokenizer.convert_tokens_to_ids("<|im_end|>") # get special token ids manually
151645

下面是官方文档中对这个参数的描述，解释的很清楚了：

split_special_tokens (bool, optional, defaults to False) — Whether or not the special tokens should be split during the tokenization process. The default behavior is to not split special tokens. This means that if <s> is the bos_token, then tokenizer.tokenize("<s>") = ['<s>]. Otherwise, if split_special_tokens=True, then tokenizer.tokenize("<s>") will be give ['<', 's', '>']. This argument is only supported for slow tokenizers for the moment.

参考内容

[1] https://github.com/QwenLM/Qwen/blob/main/tokenization_note_zh.md

[2] tiktoken

[3] Data Engineering for Scaling Language Models to 128K Context

[4] https://github.com/QwenLM/Qwen1.5/issues/14

[5] tokenizer#transformers.PreTrainedTokenizerFast