在自然语言处理（NLP）领域，分词器（Tokenizer）是连接原始文本数据与深度学习模型之间的关键桥梁。它负责将人类可读的文本转换成模型能够理解的数值序列（token IDs），并添加必要的辅助信息，如注意力掩码（attention mask）和 token 类型 ID（token type IDs）。本项目的目标是探索超越传统 Transformer 模型的高效序列建模方法，因此，一个灵活且高效的分词器对于支持不同模型架构的输入需求至关重要。
本项目采用 tokenizers 库来构建自定义的 Byte-Pair Encoding (BPE) 分词器。选择 BPE 算法的原因是其在处理 OOV (Out-Of-Vocabulary) 词汇方面的鲁棒性，以及能够通过子词（subword）表示来平衡词汇表大小和表示能力。
为了训练出能够良好覆盖领域词汇的分词器，我们使用了 GLUE (General Language Understanding Evaluation) Benchmark 数据集作为训练语料。GLUE 包含多种自然语言理解任务，其多样性有助于分词器学习到更广泛的语言模式。
语料提取过程：
- 数据源: 从
../dataset/glue_data_std目录读取 GLUE 数据。 - 文件遍历: 遍历 GLUE 数据根目录下的所有子目录，寻找
train.tsv和dev.tsv文件。 - 文本抽取: 对于每个找到的
.tsv文件，使用csv.DictReader解析，并从中提取sentence、sentence1或sentence2列的文本内容。 - 去重与过滤: 提取的句子会被添加到
set中，以确保语料的唯一性。同时，过滤掉长度小于等于 3 个词的短句子，以提高语料质量。 - 语料保存: 最终，所有去重并过滤后的句子将保存到一个名为
glue_corpus.txt的文本文件中，路径为../dataset/glue_corpus.txt。
# ... existing code ...
glue_root = "../dataset/glue_data_std"
# 遍历所有子目录中的 train.tsv 文件
all_sentences = set()
for task_name in os.listdir(glue_root):
task_path = os.path.join(glue_root, task_name)
for task in ["train.tsv", "dev.tsv"]:
train_file = os.path.join(task_path, task)
if not os.path.isfile(train_file):
continue
with open(train_file, "r", encoding="utf-8") as f:
tsv_reader = csv.DictReader(f, delimiter=" ")
for row in tsv_reader:
for key in ["sentence", "sentence1", "sentence2"]:
if key in row:
text = row[key].strip()
if len(text.split()) > 3:
all_sentences.add(text)
# 保存为语料文本文件
corpus_path = "../dataset/glue_corpus.txt"
with open(corpus_path, "w", encoding="utf-8") as f:
for line in all_sentences:
f.write(line + "\n")
print("提取完成，共计句子：", len(all_sentences))- 文件:
tokenizer/extract_glue_corpus.py
我们训练了一个自定义的 BPE 分词器，以满足特定模型的需求。
训练过程详解：
- 初始化: 创建一个基于 BPE 模型的
Tokenizer实例。 - 预分词器设置: 使用
Whitespace()预分词器，在 BPE 算法应用之前，首先通过空格将文本拆分成单词。 - 特殊 Token 定义: 定义了一组在序列建模中常用的特殊 token，包括：
[PAD]: 填充 token，用于将序列填充到固定长度。[UNK]: 未知 token，用于表示在词汇表中不存在的词。[CLS]: 分类 token，通常位于序列的开头，用于表示整个序列的语义。[SEP]: 分隔 token，用于分隔序列中的不同部分（如句子对）。[MASK]: 掩码 token，在某些预训练任务（如 Masked Language Modeling）中使用。
- 训练器配置: 配置
BpeTrainer，指定词汇表大小 (vocab_size=8000) 和定义的特殊 token。 - 分词器训练: 调用
tokenizer.train()方法，使用准备好的语料文件 (glue_corpus.txt) 进行训练。 - 解码器设置: 设置
BPEDecoder()，使其能够将 token ID 序列解码回原始文本。 - 后处理设置 (Transformer): 针对 Transformer 模型，设置
TemplateProcessing后处理。这会自动在单句输入前后添加[CLS]和[SEP]，在句对输入中添加[CLS]、[SEP]和[SEP]，确保输入格式符合 Transformer 的惯例。这些特殊 token 的 ID 会从分词器中获取。 - 保存: 将训练好的分词器保存为 JSON 文件 (
custom_bpe_tokenizer.json)。
# ... existing code ...
def train_custom_bpe_tokenizer(corpus_file, vocab_size=8000, save_path="custom_bpe_tokenizer.json"):
# 初始化BPE tokenizer
tokenizer = Tokenizer(models.BPE())
tokenizer.pre_tokenizer = pre_tokenizers.Whitespace()
# 定义特殊token集合
special_tokens = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
# 配置训练器
trainer = trainers.BpeTrainer(
vocab_size=vocab_size,
special_tokens=special_tokens,
show_progress=True
)
# 训练分词器
tokenizer.train([corpus_file], trainer=trainer)
# 设置解码器
tokenizer.decoder = decoders.BPEDecoder()
# 设置针对Transformer模型的后处理（自动添加CLS/SEP）
tokenizer.post_processor = processors.TemplateProcessing(
single="[CLS] $A [SEP]",
pair="[CLS] $A [SEP] $B [SEP]",
special_tokens=[
("[CLS]", tokenizer.token_to_id("[CLS]")),
("[SEP]", tokenizer.token_to_id("[SEP]")),
],
)
# 保存分词器
tokenizer.save(save_path)
print(f"Tokenizer saved to {save_path}")
return tokenizer
# ... existing code ...- 文件:
tokenizer/train_tokenizer.py
为了适配不同的模型架构，特别是 Transformer、Mamba 和 Fusion 模型，我们设计了灵活的编码函数。这些函数负责将原始文本转换为模型所需的 input_ids、attention_mask 和 token_type_ids。
tokenizer/tokenizerEncode.py 文件中提供了一个统一的 encode 函数，根据 model_name 参数选择不同的编码策略。
# ... existing code ...
def encode(tokenizer: Tokenizer,
text_a: str,
text_b: str = None,
max_length: int = 128,
model_name: str = 'transformer'):
"""
根据模型类型选择不同的编码方式
"""
if model_name == 'transformer':
return encode_transformer(tokenizer, text_a, text_b, max_length)
elif model_name == 'mamba':
return encode_mamba(tokenizer, text_a, text_b, max_length)
elif model_name == 'fusion':
if text_b is None:
raise ValueError("Fusion model requires both text_a and text_b")
return encode_fusion(tokenizer, text_a, text_b, max_length)
else:
raise ValueError(f"Unsupported model type: {model_name}")- 文件:
tokenizer/tokenizerEncode.py
此函数用于为标准的 Transformer 模型生成输入。它利用了 tokenizers 库中预设的 post_processor 功能，该功能在训练时已配置为自动添加 [CLS] 和 [SEP]。
处理流程：
- 分词: 调用
tokenizer.encode(text_a, text_b)进行分词。如果提供了text_b，则处理为句对；否则，处理为单句。 - 截断: 将
input_ids和type_ids截断到max_length。 - 注意力掩码: 初始化
attention_mask，所有有效 token 对应的值为 1。 - 填充: 如果序列长度小于
max_length，则使用[PAD]token ID 填充input_ids，并用 0 填充type_ids和attention_mask。 - 输出: 返回包含
input_ids、attention_mask和token_type_ids的字典，并转换为 PyTorch 张量，形状为(1, max_length)。
# ... existing code ...
def encode_transformer(tokenizer: Tokenizer, text_a: str, text_b: Optional[str] = None, max_length: int = 128):
"""
Transformer编码，自动添加 [CLS],[SEP]，输出token ids、attention mask、token type ids
"""
if text_b is not None:
encoding = tokenizer.encode(text_a, text_b)
else:
encoding = tokenizer.encode(text_a)
input_ids = encoding.ids[:max_length]
type_ids = encoding.type_ids[:max_length]
attention_mask = [1] * len(input_ids)
pad_len = max_length - len(input_ids)
if pad_len > 0:
pad_token_id = tokenizer.token_to_id("[PAD]")
input_ids += [pad_token_id] * pad_len
type_ids += [0] * pad_len
attention_mask += [0] * pad_len
return {
"input_ids": torch.tensor(input_ids, dtype=torch.long).unsqueeze(0), # shape (1, max_length)
"attention_mask": torch.tensor(attention_mask, dtype=torch.long).unsqueeze(0),
"token_type_ids": torch.tensor(type_ids, dtype=torch.long).unsqueeze(0)
}- 文件:
tokenizer/tokenizerEncode.py
Mamba 模型可能不需要 tokenizers 库提供的复杂 post_processor，因此此函数会暂时禁用它，然后手动构建输入序列。
处理流程：
- 禁用
post-processor: 暂时将tokenizer.post_processor设置为None，以避免自动添加特殊 token。 - 获取特殊 token ID: 获取
[CLS]、[SEP]和[PAD]的 token ID。 - 编码第一个句子: 分词
text_a。 - 构建
input_ids和token_type_ids(text_a):input_ids以[CLS]开始，后接text_a的 token ID。token_type_ids全为 0，表示第一个句子及其[CLS]。
- 处理第二个句子 (text_b):
- 如果
text_b存在，则在input_ids后追加[SEP]、text_b的 token ID 和另一个[SEP]。 token_type_ids相应地追加 0 ( for[SEP]) 和 1 ( fortext_b和末尾的[SEP])，表示第二个句子。- 如果
text_b不存在，则只在input_ids后追加一个[SEP]，token_type_ids相应追加 0。
- 如果
- 截断: 将
input_ids和token_type_ids截断到max_length。 - 注意力掩码: 所有有效 token 对应的值为 1。
- 填充: 如果序列长度小于
max_length，则使用[PAD]token ID 填充input_ids，并用 0 填充token_type_ids和attention_mask。 - 恢复
post-processor: 将tokenizer.post_processor恢复到原始值。 - 输出: 返回包含
input_ids、attention_mask和token_type_ids的字典，并转换为 PyTorch 张量。
# ... existing code ...
def encode_mamba(tokenizer: Tokenizer, text_a: str, text_b: Optional[str] = None, max_length: int = 128):
"""
Mamba编码，支持句对任务，自动添加 [CLS],[SEP]，输出token ids、attention mask、token type ids
"""
# 禁用 post-processor
old_post_processor = tokenizer.post_processor
tokenizer.post_processor = None
# 获取特殊token的id
cls_token_id = tokenizer.token_to_id("[CLS]")
sep_token_id = tokenizer.token_to_id("[SEP]")
pad_token_id = tokenizer.token_to_id("[PAD]")
# 编码第一个句子
encoding_a = tokenizer.encode(text_a)
input_ids = [cls_token_id] + encoding_a.ids
token_type_ids = [0] * (len(encoding_a.ids) + 1) # CLS + 第一句
# 如果有第二个句子，添加SEP和第二个句子
if text_b is not None:
encoding_b = tokenizer.encode(text_b)
input_ids += [sep_token_id] + encoding_b.ids + [sep_token_id]
token_type_ids += [0] + [1] * (len(encoding_b.ids) + 1) # SEP + 第二句 + SEP
else:
input_ids += [sep_token_id]
token_type_ids += [0] # 只有SEP
# 截断到最大长度
input_ids = input_ids[:max_length]
token_type_ids = token_type_ids[:max_length]
attention_mask = [1] * len(input_ids)
# 填充到最大长度
pad_len = max_length - len(input_ids)
if pad_len > 0:
input_ids += [pad_token_id] * pad_len
token_type_ids += [0] * pad_len # padding的token type id设为0
attention_mask += [0] * pad_len
# 恢复post-processor
tokenizer.post_processor = old_post_processor
return {
"input_ids": torch.tensor(input_ids, dtype=torch.long).unsqueeze(0),
"attention_mask": torch.tensor(attention_mask, dtype=torch.long).unsqueeze(0),
"token_type_ids": torch.tensor(token_type_ids, dtype=torch.long).unsqueeze(0)
}- 文件:
tokenizer/tokenizerEncode.py
Fusion 模型可能需要两个独立的序列作为输入，因此此函数会分别编码 text_a 和 text_b。
处理流程：
- 禁用
post-processor: 暂时将tokenizer.post_processor设置为None。 - 获取特殊 token ID: 获取
[CLS]、[SEP]和[PAD]的 token ID。 - 编码第一个序列 (text_a):
- 分词
text_a。 - 构建
input_ids1：以[CLS]开始，后接截断到max_length-2的text_atoken ID，最后以[SEP]结束。 token_type_ids1全为 0。attention_mask1全为 1。
- 分词
- 编码第二个序列 (text_b):
- 分词
text_b。 - 构建
input_ids2：以[CLS]开始，后接截断到max_length-2的text_btoken ID，最后以[SEP]结束。 token_type_ids2全为 0。attention_mask2全为 1。
- 分词
- 填充: 分别对
input_ids1/token_type_ids1/attention_mask1和input_ids2/token_type_ids2/attention_mask2进行填充，使其达到max_length。 - 恢复
post-processor: 将tokenizer.post_processor恢复到原始值。 - 输出: 返回包含两个序列编码结果的字典，键分别为
input_ids1、attention_mask1、token_type_ids1、input_ids2、attention_mask2、token_type_ids2，并转换为 PyTorch 张量。
# ... existing code ...
def encode_fusion(tokenizer: Tokenizer, text_a: str, text_b: str, max_length: int = 128):
"""
Fusion模型编码，支持文本对任务，自动添加特殊token，输出两个序列的编码结果
"""
# 禁用 post-processor
old_post_processor = tokenizer.post_processor
tokenizer.post_processor = None
# 获取特殊token的id
cls_token_id = tokenizer.token_to_id("[CLS]")
sep_token_id = tokenizer.token_to_id("[SEP]")
pad_token_id = tokenizer.token_to_id("[PAD]")
# 编码第一个序列
encoding_a = tokenizer.encode(text_a)
input_ids1 = [cls_token_id] + encoding_a.ids[:max_length-2] + [sep_token_id]
token_type_ids1 = [0] * len(input_ids1)
attention_mask1 = [1] * len(input_ids1)
# 编码第二个序列
encoding_b = tokenizer.encode(text_b)
input_ids2 = [cls_token_id] + encoding_b.ids[:max_length-2] + [sep_token_id]
token_type_ids2 = [0] * len(input_ids2)
attention_mask2 = [1] * len(input_ids2)
# 填充到最大长度
pad_len1 = max_length - len(input_ids1)
if pad_len1 > 0:
input_ids1 += [pad_token_id] * pad_len1
token_type_ids1 += [0] * pad_len1
attention_mask1 += [0] * pad_len1
pad_len2 = max_length - len(input_ids2)
if pad_len2 > 0:
input_ids2 += [pad_token_id] * pad_len2
token_type_ids2 += [0] * pad_len2
attention_mask2 += [0] * pad_len2
# 恢复post-processor
tokenizer.post_processor = old_post_processor
return {
"input_ids1": torch.tensor(input_ids1, dtype=torch.long).unsqueeze(0),
"attention_mask1": torch.tensor(attention_mask1, dtype=torch.long).unsqueeze(0),
"token_type_ids1": torch.tensor(token_type_ids1, dtype=torch.long).unsqueeze(0),
"input_ids2": torch.tensor(input_ids2, dtype=torch.long).unsqueeze(0),
"attention_mask2": torch.tensor(attention_mask2, dtype=torch.long).unsqueeze(0),
"token_type_ids2": torch.tensor(token_type_ids2, dtype=torch.long).unsqueeze(0)
}- 文件:
tokenizer/tokenizerEncode.py
本项目的分词器模块设计严谨，能够有效地支持不同高效序列建模架构的需求。通过使用 GLUE 数据集进行 BPE 分词器的训练，我们获得了对通用语言模式的良好覆盖。特别地，通过 tokenizerEncode.py 中为 Transformer、Mamba 和 Fusion 模型定制的编码函数，确保了模型能够接收正确格式的输入，这对于后续模型训练和性能评估至关重要。特殊 token 的合理使用和序列的截断与填充策略，使得分词器在保持灵活性的同时，也保证了输入序列的统一性和规范性。