在使用预训练模型时，我们有时需要使用一些自定义 token 来增强输入，例如使用 [ENT_START] 和 [ENT_END] 在文本中标记出实体。由于自定义 token 并不在预训练模型原来的词表中，因此直接运用分词器 (Tokenizer) 处理输入就会出现问题。
例如直接使用 BERT 分词器处理下面的句子：
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
sentence = 'Two [ENT_START] cars [ENT_END] collided in a [ENT_START] tunnel [ENT_END] this morning.'
print(tokenizer(sentence).tokens())
['[CLS]', 'Two', '[', 'E', '##NT', '_', 'ST', '##AR', '##T', ']', 'cars', '[', 'E', '##NT', '_', 'E', '##ND', ']', 'collided', 'in', 'a', '[', 'E', '##NT', '_', 'ST', '##AR', '##T', ']', 'tunnel', '[', 'E', '##NT', '_', 'E', '##ND', ']', 'this', 'morning', '.', '[SEP]']
由于分词器无法识别 [ENT_START] 和 [ENT_END] ，将它们都当作未知字符处理，例如 [ENT_END] 被切分成了 '['、'E'、'##NT'、'_'、'E'、'##ND'、']' 七个 token，很明显不符合我们的预期。
此外，有时我们还会遇到一些领域相关词汇，例如医学领域的文本通常会包含大量的医学术语，它们可能并不在模型的词表中（例如一些术语是使用多个词语的缩写拼接而成），这时也会出现上面的问题。
此时我们就需要将这些新 token 添加到模型的词表中，让分词器与模型可以识别并处理这些 token。
添加新 token
添加方法
Huggingface 的 Transformers 库提供了两种方式来添加新 token，分别是：
add_tokens() 添加普通 token：添加新 token 列表，如果 token 不在词表中，就会被添加到词表的最后。tokenizer = AutoTokenizer.from_pretrained("bert-base-cased") num_added_toks = tokenizer.add_tokens(["new_tok1", "my_new-tok2"]) print("We have added", num_added_toks, "tokens") We have added 2 tokens 为了防止 token 已经包含在词表中，我们还可以预先对新 token 列表进行过滤：new_tokens = ["new_tok1", "my_new-tok2"] new_tokens = set(new_tokens) - set(tokenizer.vocab.keys()) tokenizer.add_tokens(list(new_tokens))add_special_tokens() 添加特殊 token：添加包含特殊 token 的字典，键值从 bos_token, eos_token, unk_token, sep_token, pad_token, cls_token, mask_token, additional_special_tokens 中选择。与 add_tokens() 类似，如果 token 不在词表中，就会被添加到词表的最后。添加后，还可以通过特殊属性来访问这些 token，例如 tokenizer.cls_token 就指向 cls token。tokenizer = AutoTokenizer.from_pretrained("bert-base-cased") special_tokens_dict = {"cls_token": "[MY_CLS]"} num_added_toks = tokenizer.add_special_tokens(special_tokens_dict) print("We have added", num_added_toks, "tokens") assert tokenizer.cls_token == "[MY_CLS]" We have added 1 tokens 我们也可以使用 add_tokens() 添加特殊 token，只需要额外设置参数 special_tokens=True：tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") num_added_toks = tokenizer.add_tokens(["[NEW_tok1]", "[NEW_tok2]"]) num_added_toks = tokenizer.add_tokens(["[NEW_tok3]", "[NEW_tok4]"], special_tokens=True) print("We have added", num_added_toks, "tokens") print(tokenizer('[NEW_tok1] Hello [NEW_tok2] [NEW_tok3] World [NEW_tok4]!').tokens()) We have added 2 tokens ['[CLS]', '[new_tok1]', 'hello', '[new_tok2]', '[NEW_tok3]', 'world', '[NEW_tok4]', '!', '[SEP]'] 特殊 token 的标准化 (normalization) 过程与普通 token 有一些不同，比如不会被小写。这里我们使用的是不区分大小写的 BERT 模型，因此分词后添加的普通 token [NEW_tok1] 和 [NEW_tok2] 都被处理为了小写，而特殊 token [NEW_tok3] 和 [NEW_tok4] 则维持大写，与 [CLS] 等自带特殊 token 保持一致。
对于之前的例子，很明显实体标记符 [ENT_START] 和 [ENT_END] 属于特殊 token，因此按添加特殊 token 的方式进行。如果使用 add_tokens() 则需要额外设置 special_tokens=True，或者也可以直接使用 add_special_tokens()。
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
num_added_toks = tokenizer.add_tokens(['[ENT_START]', '[ENT_END]'], special_tokens=True)
# num_added_toks = tokenizer.add_special_tokens({'additional_special_tokens': ['[ENT_START]', '[ENT_END]']})
print("We have added", num_added_toks, "tokens")
sentence = 'Two [ENT_START] cars [ENT_END] collided in a [ENT_START] tunnel [ENT_END] this morning.'
print(tokenizer(sentence).tokens())
We have added 2 tokens
['[CLS]', 'two', '[ENT_START]', 'cars', '[ENT_END]', 'collided', 'in', 'a', '[ENT_START]', 'tunnel', '[ENT_END]', 'this', 'morning', '.', '[SEP]']
可以看到，分词器成功地将 [ENT_START] 和 [ENT_END] 识别为 token，并且依旧保持大写。
调整 embedding 矩阵
注意！无论使用哪种方式向词表中添加新 token 后，都需要重置模型 token embedding 矩阵的大小，也就是向矩阵中添加新 token 对应的 embedding，这样模型才可以正常工作（将 token 映射到对应的 embedding）。
该操作通过调用预训练模型的 resize_token_embeddings() 函数来实现，例如对于上面的例子：
from transformers import AutoTokenizer, AutoModel
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModel.from_pretrained("bert-base-uncased")
print(len(tokenizer))
num_added_toks = tokenizer.add_tokens(['[ENT_START]', '[ENT_END]'], special_tokens=True)
print("We have added", num_added_toks, "tokens")
print(len(tokenizer))
model.resize_token_embeddings(len(tokenizer))
print(model.embeddings.word_embeddings.weight.size())
# Randomly generated matrix
print(model.embeddings.word_embeddings.weight[-2:, :])
30522
We have added 2 tokens
30524
torch.Size([30524, 768])
tensor([[-0.0325, -0.0224, 0.0044, ..., -0.0088, -0.0078, -0.0110],
[-0.0005, -0.0167, -0.0009, ..., 0.0110, -0.0282, -0.0013]],
grad_fn=<SliceBackward0>)
可以看到，在添加了特殊 token [ENT_START] 和 [ENT_END] 之后，分词器的词表大小从 30522 增加到了 30524，并且模型的 token embedding 矩阵大小也成功调整为了 30524×76830524×768。
我们还尝试打印出新添加 token 对应的 embedding。因为新 token 会添加在词表的末尾，因此只需打印出矩阵最后两行。如果你重复运行一下上面的代码，就会发现每次打印出的 [ENT_START] 和 [ENT_END] 的 embedding 是不同的。这是因为在默认情况下，这些新 token 的 embedding 是随机初始化的。
token embedding 初始化
如果有充分的训练语料对模型进行微调或者继续预训练，那么将新添加 token 初始化为随机向量没什么问题。但是如果训练语料较少，甚至是只有很少语料的 few-shot learning 场景下，这种做法就可能存在问题。研究表明，在训练数据不够多的情况下，这些新添加 token 的 embedding 只会在初始值附近小幅波动。换句话说，即使经过训练，它们的值事实上还是随机的。
因此，在很多情况下，我们需要手工初始化这些新 token 的 embedding。对于 Transformers 库来说，可以通过直接对 embedding 矩阵赋值来实现。例如对于上面的例子，我们将这两个新 token 的 embedding 都初始化为全零向量：
with torch.no_grad():
model.embeddings.word_embeddings.weight[-2:, :] = torch.zeros([2, model.config.hidden_size], requires_grad=True)
print(model.embeddings.word_embeddings.weight[-2:, :])
tensor([[0., 0., 0., ..., 0., 0., 0.],
[0., 0., 0., ..., 0., 0., 0.]], grad_fn=<SliceBackward0>)
注意，初始化 embedding 的过程并不可导，因此这里通过 torch.no_grad() 暂停梯度的计算。
初始化为已有 token 的值
比较常见的操作是根据新添加 token 的语义，将其值初始化为训练好 token 的 embedding。例如对于上面的例子，我们可以将 [ENT_START] 和 [ENT_END] 的值都初始化为“entity”对应的 embedding。因为 token id 就是 token 在矩阵中的索引，因此我们可以直接通过 weight[token_id] 取出“entity”对应的 embedding。
token_id = tokenizer.convert_tokens_to_ids('entity')
token_embedding = model.embeddings.word_embeddings.weight[token_id]
print(token_id)
with torch.no_grad():
for i in range(1, num_added_toks+1):
model.embeddings.word_embeddings.weight[-i:, :] = token_embedding.clone().detach().requires_grad_(True)
print(model.embeddings.word_embeddings.weight[-2:, :])
9178
tensor([[-0.0039, -0.0131, -0.0946, ..., -0.0223, 0.0107, -0.0419],
[-0.0039, -0.0131, -0.0946, ..., -0.0223, 0.0107, -0.0419]],
grad_fn=<SliceBackward0>)
可以看到最终结果符合我们的预期，[ENT_START] 和 [ENT_END] 被初始化为相同的 embedding。
另一种常见的做法是根据新 token 的语义，使用对应的描述文本来完成初始化。例如将值初始化为描述文本中所有 token 的平均值，假设新 token titi 的描述文本为 wi,1,wi,2,…,wi,nwi,1,wi,2,…,wi,n，那么 titi 的初始化 embedding 为：
这里 EE 表示预训练模型的 token embedding 矩阵。对于上面的例子，我们可以分别为 [ENT_START] 和 [ENT_END] 编写对应的描述，然后再对它们的值进行初始化：
descriptions = ['start of entity', 'end of entity']
with torch.no_grad():
for i, token in enumerate(reversed(descriptions), start=1):
tokenized = tokenizer.tokenize(token)
print(tokenized)
tokenized_ids = tokenizer.convert_tokens_to_ids(tokenized)
new_embedding = model.embeddings.word_embeddings.weight[tokenized_ids].mean(axis=0)
model.embeddings.word_embeddings.weight[-i, :] = new_embedding.clone().detach().requires_grad_(True)
print(model.embeddings.word_embeddings.weight[-2:, :])
['end', 'of', 'entity']
['start', 'of', 'entity']
tensor([[-0.0340, -0.0144, -0.0441, ..., -0.0016, 0.0318, -0.0151],
[-0.0060, -0.0202, -0.0312, ..., -0.0084, 0.0193, -0.0296]],
grad_fn=<SliceBackward0>)
可以看到，这里成功地将 [ENT_START] 初始化为“start”、“of”、“entity”三个 token embedding 的平均值，将 [ENT_END] 初始化为“end”、“of”、“entity” embedding 的平均值。
保存新的token词汇表和模型权重
# 保存tokenizer
tokenizer.save_pretrained(save_directory)
# 保存model
model.save_pretrained(save_directory)
