训练自定义BPE分词器时添加非特殊Token致ID重叠的解决方法

问题背景 #

从头训练BPE分词器时，希望词汇表包含部分可能不在训练数据中的Token（比如0-9数字），同时加入特殊Token（<UNK>、<BOS>、<EOS>）。但按原代码操作后，出现Token ID重叠的情况，比如<UNK>和\n共用ID 0，<BOS>和空格共用ID 1等。

问题原因 #

原代码中提前用tokenizer.add_special_tokens()和tokenizer.add_tokens()添加的Token，会被train_from_iterator()的训练流程覆盖——BPE训练会从头构建词汇表，之前手动添加的Token ID会被新生成的Token占用，导致冲突。

正确实现方式 #

要确保自定义Token和特殊Token被正确纳入词汇表且不冲突，需要通过训练器（Trainer）的参数来指定，而不是提前手动添加：

1. 特殊Token：通过BpeTrainer的special_tokens参数传入，训练时会自动将这些Token加入词汇表并分配唯一ID，不会被覆盖。
2. 强制保留的自定义Token（如数字）：通过BpeTrainer的initial_alphabet参数传入，确保这些Token被强制包含在初始词汇表中，训练过程中不会被合并或覆盖。

修改后的代码如下：

from datasets import load_dataset
from tokenizers import models, trainers, Tokenizer

# 加载数据集
ds = load_dataset('HuggingFaceFW/fineweb', streaming=True)['train']
texts = [sample['text'] for sample in ds.take(10_000)]

# 初始化BPE分词器
tokenizer = Tokenizer(models.BPE(unk_token="<UNK>", byte_fallback=True))

# 定义特殊Token和强制保留的自定义Token
special_tokens = ["<UNK>", "<BOS>", "<EOS>"]
forced_tokens = [str(num) for num in range(10)]  # 0-9数字

# 配置BPE训练器
trainer = trainers.BpeTrainer(
    vocab_size=32_000,
    min_frequency=2,
    max_token_length=32,
    show_progress=True,
    special_tokens=special_tokens,  # 传入特殊Token
    initial_alphabet=forced_tokens  # 传入需要强制保留的自定义Token
)

# 开始训练
tokenizer.train_from_iterator(texts, trainer=trainer)

关键说明 #

• initial_alphabet参数会将指定的Token加入初始字母表，训练时这些Token会被视为基础单元，不会被拆分或合并，确保它们始终存在于词汇表中。
• special_tokens参数会自动处理特殊Token的ID分配，避免与普通Token冲突，同时<UNK>会和初始化时指定的unk_token对应。

内容的提问来源于stack exchange，提问作者meliksahturker