【大模型实战篇】大模型分词算法WordPiece分词及代码示例

继《大模型数据词元化处理BPE(Byte-Pair Encoding tokenization)》之后，我们针对大模型原始数据的分词处理，继续分享WordPiece分词技术【1】。

1. 原理分析

WordPiece 是 Google 开发的分词算法，用于预训练 BERT。此后，它被多个基于 BERT 的 Transformer 模型重用，如 DistilBERT、MobileBERT、Funnel Transformers 和 MPNET。与 BPE 的训练过程非常相似，但实际的分词方式有所不同。与 BPE 类似，WordPiece 从一个包含模型使用的特殊标记和初始字母表的小词汇表开始。由于它通过添加前缀（如 BERT 中的 ##）来识别子词，因此每个单词最初通过在单词内的所有字符前添加该前缀来分割。例如，“word” 被分割为：

w ##o ##r ##d

因此，初始字母表包含单词开头的所有字符以及以 WordPiece 前缀开头的单词内字符。然后，像 BPE 一样，WordPiece 学习合并规则。主要区别在于选择合并对的方式。WordPiece 不是选择最频繁的对，而是对每一对计算一个分数，使用以下公式：

$\text{score} = \frac{\text{freq\_of\_pair}}{\text{freq\_of\_first\_element} \times \text{freq\_of\_second\_element}}$

通过将对的频率除以其组成部分的频率的乘积，算法优先合并在词汇表中频率较低的组成部分。例如，它不会合并（"un", "##able"），即使该对在词汇表中非常频繁，因为“un”和“##able”这两个部分可能会在许多其他单词中出现并具有高频率。相比之下，像（"hu", "##gging"）这样的对可能会更快被合并（假设单词“hugging”在词汇表中频繁出现），因为“hu”和“##gging”各自的频率较低。

来看一个与 BPE 训练示例相同的词汇表：

("hug", 10), ("pug", 5), ("pun", 12), ("bun", 4), ("hugs", 5)

此处的分割将是：

("h", "##u", "##g", 10), ("p", "##u", "##g", 5), ("p", "##u", "##n", 12), ("b", "##u", "##n", 4), ("h", "##u", "##g", "##s", 5)

所以初始词汇表将是 ["b", "h", "p", "##g", "##n", "##s", "##u"]。最频繁的对是（"##u", "##g"）（出现 20 次），但“##u”的单独频率非常高，因此其分数不是最高（为 1 / 36）。所有带有“##u”的对实际上都有相同的分数（1 / 36），因此最佳分数属于对（"##g", "##s"）——唯一没有“##u”的对——为 1 / 20，第一次学习的合并是（"##g", "##s"）->（"##gs"）。

当合并时，会移除两个标记之间的 ##，因此将“##gs”添加到词汇表，并在语料库中的单词上应用合并：

词汇表: ["b", "h", "p", "##g", "##n", "##s", "##u", "##gs"]
语料库: ("h", "##u", "##g", 10), ("p", "##u", "##g", 5), ("p", "##u", "##n", 12), ("b", "##u", "##n", 4), ("h", "##u", "##gs", 5)

此时，“##u”出现在所有可能的对中，因此它们都最终得到了相同的分数。假设在这种情况下，第一个对被合并，所以（"h", "##u"）->“hu”。因此可以得到：

词汇表: ["b", "h", "p", "##g", "##n", "##s", "##u", "##gs", "hu"]
语料库: ("hu", "##g", 10), ("p", "##u", "##g", 5), ("p", "##u", "##n", 12), ("b", "##u", "##n", 4), ("hu", "##gs", 5)

接下来，最好的分数由（"hu", "##g"）和（"hu", "##gs"）共享（分别为 1/15，与其他所有对的 1/21 相比），因此具有最大分数的第一个对被合并：

词汇表: ["b", "h", "p", "##g", "##n", "##s", "##u", "##gs", "hu", "hug"]
语料库: ("hug", 10), ("p", "##u", "##g", 5), ("p", "##u", "##n", 12), ("b", "##u", "##n", 4), ("hu", "##gs", 5)

从上述例子可以看到，WordPiece 和 BPE 的分词方式不同，WordPiece 只保存最终的词汇表，而不保存学习到的合并规则（也就是我们在BPE中提到的merges对象）。从要分词的单词开始，WordPiece 查找词汇表中最长的子词，然后在其上进行分割。例如，如果使用上述示例中学习到的词汇表，对于单词“hugs”，从开头开始最长的子词是“hug”，所以我们在这里进行分割，得到 ["hug", "##s"]。接着继续处理“##s”，它在词汇表中，因此“hugs”的分词是 ["hug", "##s"]。

使用 BPE，则会按顺序应用学习到的合并，将其分词为 ["hu", "##gs"]，因此编码是不同的。

作为另一个例子，单词“bugs”将如何被分词。“b”是从单词开头开始的最长子词，因此我们在此分割，得到 ["b", "##ugs"]。然后，“##u”是“##ugs”开头的最长子词，所以在此分割，得到 ["b", "##u", "##gs"]。最后，“##gs”在词汇表中，因此这个列表就是“bugs”的分词。

当分词到达一个阶段，无法在词汇表中找到子词时，整个单词将被标记为未知——例如，“mug”将被分词为 ["[UNK]"]，而“bum”也是如此（即使可以从“b”和“##u”开始，“##m”不在词汇表中，最终的分词将只是 ["[UNK]"]，而不是 ["b", "##u", "[UNK]"]）。这是与 BPE 的另一个区别，后者只将不在词汇表中的个别字符标记为未知。

2. 代码实现示例

使用与 BPE 示例相同的语料库：

corpus = ["This is the Hugging Face Course.","This chapter is about tokenization.","This section shows several tokenizer algorithms.","Hopefully, you will be able to understand how they are trained and generate tokens.",
]

首先，需要将语料库预分词为单词。由于需要复现一个 WordPiece 分词器（如 BERT），因此将使用 bert-base-cased 分词器进行预分词，同样的，我们从model scope上下载bert-base-cased预训练模型。下载速度很快。

import torch
from modelscope import snapshot_download, AutoModel, AutoTokenizer
import osmodel_dir = snapshot_download('AI-ModelScope/bert-base-cased', cache_dir='/root/autodl-tmp', revision='master')

加载模型：

from transformers import AutoTokenizermode_name_or_path = '/root/autodl-tmp/AI-ModelScope/bert-base-cased'
tokenizer = AutoTokenizer.from_pretrained(mode_name_or_path, trust_remote_code=True)

在预分词的过程中计算语料库中每个单词的频率。字母表是由所有单词的首字母和带前缀 "##" 的所有其他字母组成的唯一集合。

将模型使用的特殊词元添加到词汇表的开头。在 BERT 的情况下，这个列表为 ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]：

vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"] + alphabet.copy()

将每个单词拆分，除了第一个字母外的所有字母前面加上 "##"：

接下来计算词对的得分：

def compute_pair_scores(splits):letter_freqs = defaultdict(int)pair_freqs = defaultdict(int)for word, freq in word_freqs.items():split = splits[word]if len(split) == 1:letter_freqs[split[0]] += freqcontinuefor i in range(len(split) - 1):pair = (split[i], split[i + 1])letter_freqs[split[i]] += freqpair_freqs[pair] += freqletter_freqs[split[-1]] += freqscores = {pair: freq / (letter_freqs[pair[0]] * letter_freqs[pair[1]])for pair, freq in pair_freqs.items()}return scores

查看初始拆分后的字典部分，找到得分最高的词对。

pair_scores = compute_pair_scores(splits)
best_pair = ""
max_score = None
for pair, score in pair_scores.items():if max_score is None or max_score < score:best_pair = pairmax_score = scoreprint(best_pair, max_score)

基于最大得分，第一个学习的合并是 ('a', '##b') -> 'ab'，并将 'ab' 添加到词汇表：

vocab.append("ab")

在拆分字典中应用该合并：

def merge_pair(a, b, splits):for word in word_freqs:split = splits[word]if len(split) == 1:continuei = 0while i < len(split) - 1:if split[i] == a and split[i + 1] == b:merge = a + b[2:] if b.startswith("##") else a + bsplit = split[:i] + [merge] + split[i + 2 :]else:i += 1splits[word] = splitreturn splits

splits = merge_pair("a", "##b", splits)
splits["about"]

接下来，将目标设定为词汇大小为 70，进行循环找到合并词对，生成词汇表：

对新文本进行分词，先进行预分词，拆分，然后在每个单词上应用分词算法。也就是说，从第一个单词的开始寻找最大的子词并进行拆分，然后对剩下的部分重复这个过程：

def encode_word(word):tokens = []while len(word) > 0:i = len(word)while i > 0 and word[:i] not in vocab:i -= 1if i == 0:return ["[UNK]"]tokens.append(word[:i])word = word[i:]if len(word) > 0:word = f"##{word}"return tokens

测试一下：

print(encode_word("Hugging"))
print(encode_word("yuanquan"))

设置对文本进行分词处理的函数（先进行预分词，然后再应用WordPiece分词算法）：

def tokenize(text):pre_tokenize_result = tokenizer._tokenizer.pre_tokenizer.pre_tokenize_str(text)pre_tokenized_text = [word for word, offset in pre_tokenize_result]encoded_words = [encode_word(word) for word in pre_tokenized_text]return sum(encoded_words, [])

测试：