不论是Tensorflow版本或者PyTorch版本的NLP预训练模型，我们都会在模型文件中看到vocab.txt文件，这个文件就是该预训练模型的词汇表。通常，模型本身都会自带词汇表文件，这是在模型预训练的时候训练得到的词汇表，具有代表性，一般不可随意更改。同时vocab.txt文件中也保留了一定数量的未使用（unuserd）词汇，用于添加新词。
  本文将介绍如何在BERT模型中添加自己的词汇，其它预训练模型原理相同。
  我们将通过三个常见的模块来介绍，分别是keras-bert，transformers,tokenizer。其中keras-bert是Keras框架实现的模块，transformers主要是PyTorch实现的模块，也可用于TensorFlow2.0版本以上，tokenizer是一个专门用于切分词（tokenize）的模块。
  通常，往预训练模型中添加新词有两种实现方式，如下：

• 直接在词汇表vocab.txt中替换[unused]
• 通过重构词汇矩阵来增加新词

keras-bert

  在keras-bert模块中，首先观察不添加新词时的切分词结果。我们以特殊标识jjj为例，代码如下：

# -*- coding: utf-8 -*-
from keras_bert import Tokenizer

# 加载词典
dict_path = './chinese_L-12_H-768_A-12/vocab.txt'
token_dict = {}
with open(dict_path, 'r', encoding='utf-8') as reader:
    for line in reader:
        token = line.strip()
        token_dict[token] = len(token_dict)

tokenizer = Tokenizer(token_dict)
text = 'jjj今天天气很好。'
tokens = tokenizer.tokenize(text)
print(tokens)

输出结果如下：

['[CLS]', 'jj', '##j', '今', '天', '天', '气', '很', '好', '。', '[SEP]']

  可以看到，如果直接按照原有模型词汇表，则不会将特殊标识jjj作为整体切分，而是按照现有切分逻辑进行切分。
  我们将模型词汇表文件中的[unused1]替换成jjj，则切分结果如下：

['[CLS]', 'jjj', '今', '天', '天', '气', '很', '好', '。', '[SEP]']

  或者不修改vocab.txt，在上述代码中将token_dict中将key[unused1]替换成jjj，比如：token_dict['jjj'] = token_dict.pop('[unused1]')。
  bert4keras模块添加新词同理。

transformers

  transformers模块添加新词也是上述两种方式，在词汇表vocab.txt中替换[unused]这种方式不再赘述，介绍如何通过重构词汇矩阵来增加新词，代码如下：

# -*- coding: utf-8 -*-
from transformers import BertTokenizer

tokenizer = BertTokenizer("./bert-base-chinese/vocab.txt")
text = 'jjj今天天气很好。'
tokens = tokenizer.tokenize(text)
print('未添加新词前:', tokens)
tokenizer.add_tokens('jjj')
tokens = tokenizer.tokenize(text)
print('添加新词后:', tokens)

输出结果结果如下：

未添加新词前: ['jj', '##j', '今', '天', '天', '气', '很', '好', '。']
添加新词后: ['jjj', '今', '天', '天', '气', '很', '好', '。']

需要注意的是，加载的模型需要略作调整，如下：

model.resize_token_embeddings(len(tokenizer))

tokenizer

  tokenizer模块添加新词也是上述两种方式，在词汇表vocab.txt中替换[unused]这种方式不再赘述，介绍如何通过重构词汇矩阵来增加新词，代码如下：

# -*- coding: utf-8 -*-
from tokenizers import BertWordPieceTokenizer
tokenizer = BertWordPieceTokenizer("./bert-base-chinese/vocab.txt", lowercase=True)

context = '今天jjj天气很好。'
tokenized_context = tokenizer.encode(context)
print(tokenized_context.ids)
print(len(tokenized_context.ids))
print("未添加新词前:", [tokenizer.id_to_token(_) for _ in tokenized_context.ids])
print("词汇表大小:", tokenizer.get_vocab_size())
tokenizer.add_special_tokens(['jjj'])
tokenized_context = tokenizer.encode(context)
print(tokenized_context.ids)
print(len(tokenized_context.ids))
print("添加新词后:", [tokenizer.id_to_token(_) for _ in tokenized_context.ids])
print("词汇表大小:", tokenizer.get_vocab_size())

输出结果如下：

[101, 791, 1921, 11095, 8334, 1921, 3698, 2523, 1962, 511, 102]
11
未添加新词前: ['[CLS]', '今', '天', 'jj', '##j', '天', '气', '很', '好', '。', '[SEP]']
词汇表大小: 21128
[101, 791, 1921, 21128, 1921, 3698, 2523, 1962, 511, 102]
10
添加新词后: ['[CLS]', '今', '天', 'jjj', '天', '气', '很', '好', '。', '[SEP]']
词汇表大小: 21129

问题探讨

  上述方式对于一般的新词，均可起效。但对于另一类特殊的新词，比如<e>,</e>等，需要另加分析，我们以tokenizer模块进行分析，如下：

# -*- coding: utf-8 -*-
from tokenizers import BertWordPieceTokenizer
tokenizer = BertWordPieceTokenizer("./bert-base-chinese/vocab.txt", lowercase=True)
# tokenizer.add_special_tokens(['<e>', '</e>', '</ec>'])

context = '<e>苹果</e>树尽早疏蕾，能节省营养，利于坐大果，促果高桩。'
tokenized_context = tokenizer.encode(context)
print(tokenized_context.ids)
print(len(tokenized_context.ids))
print([tokenizer.id_to_token(_) for _ in tokenized_context.ids])
print(tokenizer.get_vocab_size())

我们在词汇表vocab.txt中替换[unused]，但不会起效，输出结果如下：

[101, 133, 147, 135, 5741, 3362, 133, 120, 147, 135, 3409, 2226, 3193, 4541, 5945, 8024, 5543, 5688, 4689, 5852, 1075, 8024, 1164, 754, 1777, 1920, 3362, 8024, 914, 3362, 7770, 3445, 511, 102]
34
['[CLS]', '<', 'e', '>', '苹', '果', '<', '/', 'e', '>', '树', '尽', '早', '疏', '蕾', '，', '能', '节', '省', '营', '养', '，', '利', '于', '坐', '大', '果', '，', '促', '果', '高', '桩', '。', '[SEP]']
21128

但add_special_tokens会起效，原因为<,e,>和<e>均存在于vocab.txt，但前三者的优先级高于<e>，而add_special_tokens会起效，却会使得词汇表大小增大，从而需另外调整模型size。
  但是，如果同时在词汇表vocab.txt中替换[unused]，同时add_special_tokens，则新增词会起效，同时词汇表大小不变。

总结

  本文介绍如何在BERT模型中添加自己的词汇，其它预训练模型原理相同。同时，tokenizer也是一个不错的切分词的模块，建议读者有空可以尝试~