大语言模型训练分词器（附源代码）

一、数据分割和填充

将数据集进行预处理，进行分割，不足长度的进行填充

def split_and_pad(text_list, max_token=10, pad_token="[PAD]"):
    """
    对输入的text_list按照max_token分割，并将最后不足max_token长度的列表填充至max_token长度。

    参数:
    - text_list (list): 输入的文本token列表。
    - max_token (int): 每个分段的最大token数。
    - pad_token (str): 用于填充的特殊token。

    返回:
    - result (list): 分割后的文本列表，每个子列表长度等于max_token。
    """
    # 初始化结果列表
    result = []

    # 按照max_token分割文本
    for i in range(0, len(text_list), max_token):
        chunk = text_list[i:i + max_token]

        # 如果最后一段长度不足max_token，填充pad_token
        if len(chunk) < max_token:
            chunk.extend([pad_token] * (max_token - len(chunk)))

        result.append(chunk)

    return result


# 示例用法
text_list = ["人工智能", "(AI)", "是一种", "正在", "迅速", "发展", "的", "技术", "。", "越来越多", "的", "领域",
             "正在", "使用", "AI", "技术", "，", "从", "医疗", "到", "金融", "，", "各行各业", "都在", "受益于",
             "它", "的", "强大", "能力", "。"]

max_token = 8  # 每段最多10个token
result = split_and_pad(text_list, max_token=max_token)

# 打印结果
for i, chunk in enumerate(result):
    print(len(chunk))
    print(f"Chunk {i + 1}: {chunk}")

二、自己训练一个分词器

分词器技术包括分词和编码，分词技术主流包含BEP、WordPiece、Unigram。分词之后会生成一个词汇表。接下来介绍参考BELLE的方法进行分词。

-- coding: utf-8 --
# 1、自定义领域语料的文本并保存到txt文件中
text = r"由于ChatGPT、GPT-4近期火爆整个互联网，掀起了人工智能相关的二次开发应用的热潮，博主同时也应广大网友私信请求(太多了，无法一一回复，感谢理解！！)。在本文章中，博主将会及时、持续地更新人工智能领域最新的PTMs算法模型，以及LLMs大模型的部署实战案例。同时，博主也会特持续收集很多基于AI的产品合集，以方便广大网友测试和测试，并同时反馈产品效果，博主会及时更新产品排序。如果大家有新的AI工具，也可留言，博主会将留言的提到的AI小工具，加入到本文章内容。"
sentences = text.split("。")
# 打开一个txt文件，将句子写入文件
with open("sentences.txt", "w", encoding="utf-8") as file:
for sentence in sentences:
file.write(sentence + "。\n")
print("文本已拆分并保存到sentences.txt文件中。")
# 2、训练分词模型
# 2.1、使用 SentencePiece 训练分词模型
model_name = 'tokenizer'
import sentencepiece as spm
spm.SentencePieceTrainer.train(
input='sentences.txt',
model_prefix=model_name,
vocab_size=450,
user_defined_symbols=['foo', 'bar'],  # 自定义的符号
character_coverage=1.0,
model_type="bpe",
)
# 运行后会得到tokenizer.model和tokenizer.vocab两个文件
# 2.2、查看生成的词汇表
vocab_path = "%s.vocab"% model_name
with open(vocab_path, "r", encoding="utf-8") as file:
vocab_contents = file.read()
print(vocab_contents)
# 3、测试分词模型：加载训练好的 SentencePiece 模型对新文本进行分词测试
import sentencepiece as spm
sp = spm.SentencePieceProcessor()
sp.Load("tokenizer.model")
text = "由于ChatGPT、GPT-4近期火爆整个互联网。"
tokens = sp.EncodeAsPieces(text)
print(tokens)

三、将训练出来的分词器进行合并

使用SentencePiece库来训练一个名为belle的BPE分词器→加载两个现有的分词器模型→合并词汇表→保存合并后的模型作为新的分词器模型

''' Train tokenizer'''
import sentencepiece as spm
# 1、使用SentencePiece库来训练一个名为belle的BPE分词器：词表大小为25000,覆盖率为99.95%。
'''
input: 指定了训练模型所需的输入文本文件的路径。
model_prefix: 定义了模型文件的前缀名称，将生成两个文件，一个是belle.model用于保存模型配置，另一个是belle.vocab用于保存词汇表。
model_type: 这里设置为'bpe'，表示使用BPE (Byte Pair Encoding) 分词算法。
vocab_size: 指定了词汇表的大小，这里设置为25000。
character_coverage: 指定了覆盖字符的比例，这里设置为0.9995，表示要尽量覆盖99.95%的字符。
'''
spm.SentencePieceTrainer.train(input='/path/to/input_text',
model_prefix='belle',
model_type='bpe',
vocab_size=25000,
character_coverage=0.9995)
# 2、加载两个现有的分词器模型文件orig_model和belle_model
''' Merge tokenizer '''
from sentencepiece import sentencepiece_model_pb2 as model
orig_model_path = '/path/to/llama/tokenizer.model'
belle_model_path = '/path/to/belle/belle.model'
orig_m = model.ModelProto()
belle_m = model.ModelProto()
# 使用ParseFromString方法从模型文件中读取模型配置信息，将信息解析到这两个对象中
orig_m.ParseFromString(open(orig_model_path, "rb").read())
belle_m.ParseFromString(open(belle_model_path, "rb").read())
print(len(orig_m.pieces), len(belle_m.pieces))
# 3、合并词汇表：将belle_model中的词汇表融合到orig_model中
'''
首先，遍历了两个模型的词汇表，将orig_model的词汇表片段存入orig_pieces列表
然后，遍历belle_model的词汇表片段，并检查belle_m中的词汇片段是否已经存在于orig_pieces中，
如果不存在，就将这个词汇片段添加到orig_m的词汇表中，同时也将它添加到orig_pieces列表中。
'''
orig_pieces = []
for piece in orig_m.pieces:
orig_pieces.append(piece.piece)
for piece in belle_m.pieces:
if piece.piece not in orig_pieces:
orig_m.pieces.append(piece)
orig_pieces.append(piece.piece)
print(len(orig_m.pieces))
# 4、保存合并后的模型：将融合后的orig_model序列化并保存到文件中,作为新的分词器模型
save_vocab_path = '/path/to/merge_tokenizer/tokenizer.model'
with open(save_vocab_path, 'wb') as f:
f.write(orig_m.SerializeToString())