[NLP] LLM---扩充词表LLama2-构建中文tokenization

使用SentencePiece的除了从0开始训练大模型的土豪和大公司外，大部分应该都是使用其为当前开源的大模型扩充词表，比如为LLama扩充通用中文词表（通用中文词表，或者垂直领域词表）。那这部分工作有没有意义呢？或者说值不值得投入资源去做呢？先说自己的结论，有，以下两点的作用，第三点不确定：

1.提高模型的编解码的效率，在LLaMa原来的词表上，一个汉字平均1.45个token，扩充后的Chinese-LLaMa为0.65个token；那在垂直领域内呢？比如在LLaMa在继续扩充领域内词表，金融或者医疗等等，把“负债表”，“糖尿病”等领域词汇也加入词表里，那更加能提高其编解码的效率。

2.提高模型的上下文窗口长度，原LLaMa上下文长度是4096个token，不扩充词表前，按1.45来算就是最多只能输入2824个汉字，扩充后以0.65来算的话就是6301，垂直领域会更大。这点带来的好处是实打实的。

3.提高模型的效果？提高LLaMa在中文的表现？提高开源模型在垂直领域的表现？这一点上难以下结论，目前好像也没有确定的结论，自我感觉会有，但是不多，而且可能在垂直领域扩充词表后，垂直领域词太多过拟合影响通用领域效果，还有就是扩充完词表后还要经过一系列的后续处理和训练，可以控制变量的研究一下，但需要很多的资源哈哈。但是前两点的好处是实打实的，所以在有资源的情况下，扩充词表还是可以尝试的。

目前，大语言模型呈爆发式的增长，其中，基于llama家族的模型占据了半壁江山。而原始的llama模型对中文的支持不太友好，接下来本文将讲解如何去扩充vocab里面的词以对中文进行token化。

第一阶段的主要工作是通过扩展词表和 embedding 来提升 llama2 的中文能力。今天的工作是获得一个中文的bpe分词模型。

BPE模型对后期的 token 长度、token效果影响较大，而且我们希望训练的 llama2 模型具有通用价值，所以训练的数据应该尽可能多样。

liwu/MNBVC · Datasets at Hugging Face 是一个不错的数据集，多样性丰富。
https://github.com/aceimnorstuvwxz/toutiao-text-classfication-dataset 今日头条短文本
GitHub - shjwudp/shu: 中文书籍收录整理, Collection of Chinese Books 一些中文书籍
维基百科

可以加入更多的数据

数据预处理

对斗破苍穹语料进行预处理，每一行为一句或多句话。

with open("data/《斗破苍穹》.txt", "r", encoding="utf-8") as fp:
    data = fp.read().strip().split("\n")
sentences = []

for d in data:
    d = d.strip()
    if "===" in d or len(d) == 0 or d == "《斗破苍穹》来自:":
        continue
    sentences.append(d)

with open("data/corpus.txt", "w", encoding="utf-8") as fp:
    fp.write("\n".join(sentences))

最终得到corpus.txt。

训练一个sentencepiece模型

首先，我们需要去构建中文的词库。一般的，目前比较主流的是使用sentencepiece训练中文词库。安装指令也很简单：pip install sentencepiece。然后，我们准备好语料，这里我们使用的语料是斗破苍穹小说。

直接看代码：

import sentencepiece as spm
spm.SentencePieceTrainer.train(
    input='data/corpus.txt',
    model_prefix='tokenizer',
    vocab_size=30000,
    user_defined_symbols=['foo', 'bar'],
    character_coverage=1.0,
    model_type="bpe",
)

这里讲下每个参数的作用：

input：指定输入文本文件的路径或者是一个目录，可以指定多个输入文件或目录。其中每一行可以是一句话或者多句话。
tokenizer：保存的模型的名称前缀。
vocab_size：设置的词表大小。
user_defined_symbols：用于指定用户自定义的符号。这些符号将会被视为单独的 Token，不会被拆分成子词。这个参数的作用是将一些用户定义的特殊符号作为一个整体加入到生成的词表中，以便于后续的模型使用。这里我们简单进行了测试。
model_type: 指定模型的类型，有三种可选参数：unigram, bpe, char. word。
character_coverage指定覆盖字符的数量，可以理解为限制字符集的大小。默认值为 1.0，即覆盖全部字符。
unk_id: 指定未登录词的 ID 号，即在词表中为未登录词分配一个整数 ID。默认值为 0。
bos_id: 指定句子开头符号的 ID 号，即在词表中为句子开头符号分配一个整数 ID。默认值为 1。
eos_id: 指定句子结束符号的 ID 号，即在词表中为句子结束符号分配一个整数 ID。默认值为 2。
pad_id: 指定填充符号的 ID 号，即在词表中为填充符号分配一个整数 ID。默认值为 -1，即不使用填充符号。

运行后会得到tokenizer.model和tokenizer.vocab两个文件。

我们来看看tokenizer.vocab里面是什么：

除了一些特殊符号外，还有我们自定义的foo和bar，其余的一些词是BPE训练得到，具体什么是BPE算法这里不作展开了。

合并LLama2词表和中文词表

import os

os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python"
from transformers import LlamaTokenizer
from sentencepiece import sentencepiece_model_pb2 as sp_pb2_model
import sentencepiece as spm

llama2_tokenizer_dir = "llama2_tokenizer/tokenizer.model"
llama2_tokenizer = LlamaTokenizer.from_pretrained(llama2_tokenizer_dir)

chinese_sp_model = spm.SentencePieceProcessor()
chinese_sp_model_file = "tokenizer.model"
chinese_sp_model.Load(chinese_sp_model_file)

llama2_spm = sp_pb2_model.ModelProto()
llama2_spm.ParseFromString(llama2_tokenizer.sp_model.serialized_model_proto())

chinese_spm = sp_pb2_model.ModelProto()
chinese_spm.ParseFromString(chinese_sp_model.serialized_model_proto())


# print number of tokens
print(len(llama2_tokenizer), len(chinese_sp_model))
print(llama2_tokenizer.all_special_tokens)
print(llama2_tokenizer.all_special_ids)
print(llama2_tokenizer.special_tokens_map)

## Add Chinese tokens to LLaMA2 tokenizer
llama_spm_tokens_set = set(p.piece for p in llama2_spm.pieces)
print(len(llama_spm_tokens_set))
print(f"Before:{len(llama_spm_tokens_set)}")

for p in chinese_spm.pieces:
    piece = p.piece
    if piece not in llama_spm_tokens_set:
        new_p = sp_pb2_model.ModelProto().SentencePiece()
        new_p.piece = piece
        new_p.score = 0
        llama2_spm.pieces.append(new_p)
print(f"New model pieces: {len(llama2_spm.pieces)}")

## Save
output_sp_dir = 'llama2_chinese'
os.makedirs(output_sp_dir, exist_ok=True)
with open(output_sp_dir + '/chinese_llama2.model', 'wb') as f:
    f.write(llama2_spm.SerializeToString())
tokenizer = LlamaTokenizer(vocab_file=output_sp_dir + '/chinese_llama2.model')

output_hf_dir = 'llama2_chinese'  #
os.makedirs(output_hf_dir, exist_ok=True)
tokenizer.save_pretrained(output_hf_dir)
print(f"Chinese-LLaMA tokenizer has been saved to {output_hf_dir}")

# Test
llama_tokenizer = LlamaTokenizer.from_pretrained(llama2_tokenizer_dir)
chinese_llama_tokenizer = LlamaTokenizer.from_pretrained(output_hf_dir)
print(tokenizer.all_special_tokens)
print(tokenizer.all_special_ids)
print(tokenizer.special_tokens_map)
text = '''白日依山尽，黄河入海流。欲穷千里目，更上一层楼。
The primary use of LLaMA is research on large language models, including'''
print("Test text:\n", text)
print(f"Tokenized by LLaMA tokenizer:{llama_tokenizer.tokenize(text)}")
print(f"Tokenized by ChatGLM tokenizer:{chinese_llama_tokenizer.tokenize(text)}")

加入了我们定义的词表后确实能够正确的对中文进行分词了

核心部分是这一块：

for p in chinese_spm.pieces:
    piece = p.piece
    if piece not in llama_spm_tokens_set:
        new_p = sp_pb2_model.ModelProto().SentencePiece()
        new_p.piece = piece
        new_p.score = 0
        llama_spm.pieces.append(new_p)

也就是将原始词表中没有的新加入进去。

怎么使用修改后的词表

如果我们重新从头开始训练，那么其实使用起来很简单：

config = AutoConfig.from_pretrained(...)
tokenizer = LlamaTokenizer.from_pretrained(...)
model = LlamaForCausalLM.from_pretrained(..., config=config)
model_vocab_size = model.get_output_embeddings().weight.size(0)
model.resize_token_embeddings(len(tokenizer))

但是如果我们想要保留原始模型embedding的参数，那么我们可以这么做：

1、找到新词表和旧词表id之间的映射关系。
2、将模型里面新词表里面包含的旧词表用原始模型的embedding替换。
3、如果新词在旧词表里面没有出现就进行相应的初始化再进行赋值。比如transformers库中的llama是这么进行初始化的：

 def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()

具体怎么做可以参考一下这个：https://github.com/yangjianxin1/LLMPruner