如何构建最简轻量HuggingFace分词器用于pytest单元测试？

问题原因 #

你初始化WordPiece模型时指定了unk_token="[UNK]"，但未给分词器提供包含该token的词汇表，导致编码操作触发错误。

极简分词器实现方案 #

以下是轻量、无需预训练模型的实现，完全适配单元测试需求：

import tokenizers
from tokenizers.pre_tokenizers import Whitespace
from tokenizers.processors import TemplateProcessing

# 1. 初始化WordPiece模型，指定未知token
model = tokenizers.models.WordPiece(unk_token="[UNK]")

# 2. 构建极简词汇表，必须包含[UNK]，可额外添加测试用token
vocab = {
    "[UNK]": 0,
    "this": 1,
    "is": 2,
    "a": 3,
    "test": 4,
    "text": 5
}
model.vocab = vocab

# 3. 初始化分词器并配置核心组件
tokenizer = tokenizers.Tokenizer(model)

# 设置预分词器：按空格拆分文本，模拟常规分词逻辑
tokenizer.pre_tokenizer = Whitespace()

# 可选：设置后处理器，生成带[CLS][SEP]的格式，对齐BERT类分词器输出
tokenizer.post_processor = TemplateProcessing(
    single="[CLS] $A [SEP]",
    pair="[CLS] $A [SEP] $B:1 [SEP]:1",
    special_tokens=[
        ("[CLS]", 6),
        ("[SEP]", 7),
    ],
)

# 测试编码功能
output = tokenizer.encode("this is a test text")
print("Tokens:", output.tokens)
print("IDs:", output.ids)
print("Attention Mask:", output.attention_mask)

代码说明 #

• 手动构建极小词汇表，保证包含[UNK]及测试所需基础token
• 搭配Whitespace预分词器，实现最基础的按空格拆分逻辑
• 可选添加后处理器，生成和主流预训练分词器一致的特殊token格式
• 所有组件均本地构建，无需下载任何外部模型

测试效果 #

运行后会输出标准分词器格式的结果：

Tokens: ['[CLS]', 'this', 'is', 'a', 'test', 'text', '[SEP]']
IDs: [6, 1, 2, 3, 4, 5, 7]
Attention Mask: [1, 1, 1, 1, 1, 1, 1]

若输入词汇表外的词（如"unknown"），会自动替换为[UNK]：

output = tokenizer.encode("this is unknown")
print(output.tokens)  # ['[CLS]', 'this', 'is', '[UNK]', '[SEP]']

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