LLM Tokenizer 介绍

https://xiaosheng.blog/2023/01/07/add-new-token

Tokenizer 初始化

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")

添加新 token

• 添加普通token

new_tokens = ["new_tok1", "my_new-tok2"]
new_tokens = set(new_tokens) - set(tokenizer.vocab.keys())
tokenizer.add_tokens(list(new_tokens))

num_added_toks = tokenizer.add_tokens(["new_tok1", "my_new-tok2"])

• 添加特殊token

special_tokens_dict = {"cls_token": "[MY_CLS]"}
  
num_added_toks = tokenizer.add_special_tokens(special_tokens_dict)


########################

num_added_toks = tokenizer.add_tokens(["[NEW_tok3]", "[NEW_tok4]"], special_tokens=True)

PS：特殊 token 的标准化 (normalization) 过程与普通 token 有一些不同，比如不会被小写。这里我们使用的是不区分大小写的 BERT 模型，因此分词后添加的普通 token [NEW_tok1] 和 [NEW_tok2] 都被处理为了小写，而特殊 token [NEW_tok3] 和 [NEW_tok4] 则维持大写，与 [CLS] 等自带特殊 token 保持一致。

一个例子

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())

调整 embedding 矩阵

注意！无论使用哪种方式向词表中添加新 token 后，都需要重置模型 token embedding 矩阵的大小，也就是向矩阵中添加新 token 对应的 embedding，这样模型才可以正常工作（将 token 映射到对应的 embedding）。

PS：不同的语言模型获取 embedding 的方式会有不同，需要根据模型结构进行检索。LLM一般可以通过 get_input_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 embedding 初始化

一种常见的做法是根据新 token 的语义，使用对应的描述文本来完成初始化，即将值初始化为描述文本中所有 token 的平均值。

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>)
'''