为 BERT 模型训练分词器

pip install tokenizers

# 示例文本数据
training_corpus = [
    "The quick brown fox jumps over the lazy dog.",
    "A lazy dog is not good for any one.",
    "The fox is quick and the dog is lazy.",
    "Quick brown fox, quick.",
    "A lazy dog.",
    "The dog is quick.",
]

# 将数据保存到文件
corpus_file = "training_data.txt"
with open(corpus_file, "w", encoding="utf-8") as f:
    for sentence in training_corpus:
        f.write(sentence + "\n")

print(f"Training data saved to {corpus_file}")

from tokenizers import Tokenizer
from tokenizers.models import WordPiece
from tokenizers.trainers import WordPieceTrainer
from tokenizers.pre_tokenizers import Whitespace

# 1. 初始化 WordPiece 模型
# 我们将使用一个基本的 WordPiece 模型
tokenizer = Tokenizer(WordPiece(unk_token="[UNK]"))

# 2. 设置预分词器 (Pre-tokenizer)
# BERT通常使用基于空格的分词
tokenizer.pre_tokenizer = Whitespace()

# 3. 设置训练器
# 使用 BERT 常见的词汇表大小和特殊标记
special_tokens = ['[UNK]', '[CLS]', '[SEP]', '[PAD]', '[MASK]']
vocab_size = 30000
min_frequency = 2

trainer = WordPieceTrainer(
    vocab_size=vocab_size,
    min_frequency=min_frequency,
    special_tokens=special_tokens
)

# 4. 训练分词器
files = [corpus_file]
print("Starting tokenizer training...")
tokenizer.train(files, trainer=trainer)
print("Tokenizer training complete.")

# 检查词汇表大小
print(f"Vocabulary size: {tokenizer.get_vocab_size()}")

# 打印词汇表中的前20个标记
vocab = tokenizer.get_vocab()
print("Top 20 tokens:")
# 按频率排序（如果tokenizers库没有内置频率排序，这里只是展示部分）
# 我们通过词典的键来展示
token_list = sorted(vocab, key=vocab.get, reverse=True)
print(token_list[:20])

# 1. 编码示例文本
text_to_encode = "The quick brown fox jumps over the lazy dog."
encoding = tokenizer.encode(text_to_encode)

print(f"Original Text: {text_to_encode}")
print(f"Tokens: {encoding.tokens}")
print(f"IDs: {encoding.ids}")

# 2. 解码令牌序列
decoded_text = tokenizer.decode(encoding.ids)
print(f"Decoded Text: {decoded_text}")

# 3. 检查特殊标记
print(f"[UNK] token ID: {tokenizer.token_to_id('[UNK]')}")

# 保存分词器
save_path = "bert_wordpiece_tokenizer.json"
tokenizer.save(save_path)
print(f"Tokenizer saved to {save_path}")

# 加载分词器 (在新的会话或脚本中)
from tokenizers import Tokenizer

loaded_tokenizer = Tokenizer.from_file(save_path)

# 测试加载的分词器
loaded_encoding = loaded_tokenizer.encode("A lazy dog is quick.")
print(f"Loaded Tokenizer Tokens: {loaded_encoding.tokens}")

from tokenizers import BertWordPieceTokenizer

# 初始化新的 BERT 分词器
tokenizer = BertWordPieceTokenizer(
    clean_text=True,
    handle_chinese_chars=False,
    strip_accents=True,
    lowercase=True,
)

# 训练器参数（与之前类似）
special_tokens = ['[UNK]', '[CLS]', '[SEP]', '[PAD]', '[MASK]']

# 训练
tokenizer.train(
    files=[corpus_file],
    vocab_size=vocab_size,
    min_frequency=min_frequency,
    special_tokens=special_tokens
)

# 编码示例（它会自动添加 [CLS] 和 [SEP]）
bert_encoding = tokenizer.encode("The quick brown fox.")

print("\n--- Using BertWordPieceTokenizer ---")
print(f"Tokens: {bert_encoding.tokens}")
print(f"IDs: {bert_encoding.ids}")

# 保存
bert_tokenizer_path = "bert_tokenizer.json"
tokenizer.save_model("bert_tokenizer_model", "bert_tokenizer")
print(f"BertTokenizer saved to bert_tokenizer_model directory.")