使用自己的领域数据扩充baichuan模型词表（其他模型也一样）

前言

总的来说，扩充词表可以加快解码速度，对于对中文支持不太友好的模型（如llama），扩充词表还能提升模型在中文的表现。

环境

jsonlines==3.1.0
sentencepiece==0.1.99
transformers==4.28.1

项目结构

其中 tokenization_baichuan.py是直接从百川模型文件夹里复制过来的

一、使用步骤

pip install -r requirements.txt
cd src
python main.py

二、训练词表

训练的代码不多

def train_sp(train_file, domain_sp_model_name):
    max_sentence_length = 200000
    pad_id = 3
    model_type = "BPE"
    vocab_size = 8000


    spm.SentencePieceTrainer.train(
        input=train_file,
        model_prefix=domain_sp_model_name,
        shuffle_input_sentence=False,
        train_extremely_large_corpus=True,
        max_sentence_length=max_sentence_length,
        pad_id=pad_id,
        model_type=model_type,
        vocab_size=vocab_size,
        split_digits=True,
        split_by_unicode_script=True,
        byte_fallback=True,
        allow_whitespace_only_pieces=True,
        remove_extra_whitespaces=False,
        normalization_rule_name="nfkc",
        num_threads=16
    )

参数说明：

--input（以逗号分隔的输入句子列表） type: std::string default: ""
--input_format（输入格式。支持的格式为 `text` 或 `tsv`.) type: std::string default: ""
--model_prefix（输出模型前缀） type: std::string default: ""
--model_type（模型算法：unigram、bpe、word 或 char） type: std::string 默认值："unigram"
--vocab_size（词汇量大小） 类型：int32 默认值：8000
--accept_language（以逗号分隔的该模型可接受的语言列表 type: std::string 默认值: ""
--self_test_sample_size（自我测试样本的大小） type: int32 default: 0
--character_coverage（用于确定最小符号的字符覆盖率） 类型： double 默认值： 0.9995
--input_sentence_size（训练器加载句子的最大大小） 类型： std::uint64_t 默认值： 0
--shuffle_input_sentence（提前随机抽样输入句子。当--input_sentence_size > 0时有效） type: bool 默认值: true
--seed_sentencepiece_size（种子句子片段的大小） 类型： int32 默认值：1000000
--shrinking_factor（相对于损失，保留收缩因子最高的句子类型：double 默认值：0.75
--num_threads（训练的线程数） 类型： int32 默认值： 16
--num_sub_iterations（EM 子迭代次数） 类型： int32 默认值： 2
--max_sentencepiece_length（句子片段的最大长度） 类型： int32 默认值： 16
--max_sentence_length（以字节为单位的句子最大长度） 类型： int32 默认值： 4192
--split_by_unicode_script（使用 Unicode 脚本分割句子类型：bool 默认值：true
--split_by_number（按数字（0-9）分割词块） 类型：bool 默认值：true
--split_by_whitespace（使用空格分割句子类型：bool 默认值：true
--split_digits（将所有数字（0-9）分割成独立片段） 类型：bool 默认值：false
--treat_whitespace_as_suffix（将空白标记视为后缀而非前缀） type: bool 默认值: false
--allow_whitespace_only_piece（允许只包含（连续）空白标记的片段 type: bool 默认值: false
--control_symbols（以逗号分隔的控制符号列表 type: std::string 默认值: ""
--control_symbols_file（从文件加载控制符号） type: std::string default: ""
--user_defined_symbols（以逗号分隔的用户自定义符号列表） type: std::string default: ""
--user_defined_symbols_file（从文件加载用户定义的符号） type: std::string default: ""
--required_chars（无论 --character_coverage 为何，此标志中的 UTF8 字符始终用于字符集） type: std::string 默认值: ""
--required_chars_file（从文件加载 required_chars） type: std::string default: ""
--byte_fallback（将未知片段分解为 UTF-8 字节片段） type: bool default: false
--vocabulary_output_piece_score（在词汇文件中定义分数） type: bool default: true
--normalization_rule_name（规范化规则名称，可从 nfkc 或 identity 中选择 type: std::string default: "nmt_nfkc"
--normalization_rule_tsv（规范化规则 TSV 文件） type: std::string default: ""
--denormalization_rule_tsv（去规范化规则 TSV 文件） type: std::string default: ""
--add_dummy_prefix（在文本开头添加虚假空白） type: bool default: true
--remove_extra_whitespaces（删除前导、尾部和重复的内部空白） type: bool 默认值: true
--hard_vocab_limit（如果设置为 false，--vocab_size 将被视为软限制） type: bool 默认值: true
--use_all_vocab（如果设置为 true，则使用所有标记作为词汇表。 对单词/字符模型有效。
--unk_id（覆盖 UNK () id。）类型：int32 默认值：0
--bos_id（覆盖 BOS () id。设置 -1 可禁用 BOS。） 类型：int32 默认值：1
--eos_id （覆盖 EOS () id。设置 -1 可禁用 EOS。） 类型：int32 默认值：2
--pad_id （覆盖 PAD () id。设置 -1 可禁用 PAD。） 类型：int32 默认值：-1
--unk_piece（覆盖 UNK () piece。）类型：std::string 默认值：""
--bos_piece (覆盖 BOS () piece.) 类型: std::string 默认值: ""
--eos_piece（覆盖 EOS () piece。） 类型：std::string 默认值：""
--pad_piece (覆盖 PAD () 部件。) 类型：std::string 默认值：""
--unk_surface （的虚拟表面字符串。在解码时， 被解码为 `unk_surface`.) type: std::string default: " ⁇ "
--train_extremely_large_corpus（增加单字节标记化的位深度。） type: bool default: false
--random_seed （随机生成器的种子值） 类型： uint32 默认值： 4294967295
--enable_differential_privacy（是否在训练时添加 DP。目前仅 UNIGRAM 模型支持。） 类型：bool 默认值：false
--differential_privacy_noise_level（为 DP 增加的噪声量） type: float default: 0
--differential_privacy_clipping_threshold（为 DP 削除计数的阈值） type: std::uint64_t default: 0

三、合并词表

def merge_vocab(src_tokenizer_dir, domain_sp_model_file, output_sp_dir, output_hf_dir, domain_vocab_file=None):

    # load
    src_tokenizer = AutoTokenizer.from_pretrained(src_tokenizer_dir, trust_remote_code=True)
    domain_sp_model = spm.SentencePieceProcessor()
    domain_sp_model.Load(domain_sp_model_file)

    src_tokenizer_spm = sp_pb2_model.ModelProto()
    src_tokenizer_spm.ParseFromString(src_tokenizer.sp_model.serialized_model_proto())
    domain_spm = sp_pb2_model.ModelProto()
    domain_spm.ParseFromString(domain_sp_model.serialized_model_proto())


    print("原始词表大小： ", len(src_tokenizer))
    print("用自己的数据训练出来的词表大小", len(domain_sp_model))
    
    # 原始词表集合
    src_tokenizer_spm_tokens_set = set(p.piece for p in src_tokenizer_spm.pieces)
    print("原始词表 example: ", list(src_tokenizer_spm_tokens_set)[:10])


    # 拓展词表
    added_set = set()
    for p in domain_spm.pieces:
        piece = p.piece
        if piece not in src_tokenizer_spm_tokens_set:
            # print('新增词：', piece)
            new_p = sp_pb2_model.ModelProto().SentencePiece()
            new_p.piece = piece
            new_p.score = 0
            src_tokenizer_spm.pieces.append(new_p)
            added_set.add(piece)
    
    print(f"加了sp训练出来的词表后，新词表大小: {len(src_tokenizer_spm.pieces)}")


    # 加载自定义的词表
    if domain_vocab_file:
        word_freqs = load_user_vocab(domain_vocab_file)
        domain_vocab_set = set([i[0] for i in word_freqs if i])
        print('自定义的词表 example:', list(domain_vocab_set)[:10])
        print('自定义的词表大小:', len(domain_vocab_set))
        for p in domain_vocab_set:
            piece = p
            if piece not in src_tokenizer_spm_tokens_set and piece not in added_set:
                # print('新增词：', piece)
                new_p = sp_pb2_model.ModelProto().SentencePiece()
                new_p.piece = piece
                new_p.score = 0
                src_tokenizer_spm.pieces.append(new_p)
        print(f"加了自定义的词表后，新词表大小: {len(src_tokenizer_spm.pieces)}")

    os.makedirs(output_sp_dir, exist_ok=True)
    with open(output_sp_dir + '/new_tokenizer.model', 'wb') as f:
        f.write(src_tokenizer_spm.SerializeToString())
    tokenizer = BaichuanTokenizer(vocab_file=output_sp_dir + '/new_tokenizer.model')

    tokenizer.save_pretrained(output_hf_dir)
    print(f"new tokenizer has been saved to {output_hf_dir}")

四、效果

将原始词表大小64000扩充到129570，增加了一倍。

要扩充其他模型的词表也是一样的，将tokenization_baichuan.py换成其他模型，再修改下对应的代码即可。

全部代码已上传至Github： https://github.com/seanzhang-zhichen/expand-baichuan-tokenizer