【大模型实战篇】大模型分词算法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 os
 
model_dir = snapshot_download('AI-ModelScope/bert-base-cased', cache_dir='/root/autodl-tmp', revision='master')

加载模型：

from transformers import AutoTokenizer

mode_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]] += freq
            continue
        for i in range(len(split) - 1):
            pair = (split[i], split[i + 1])
            letter_freqs[split[i]] += freq
            pair_freqs[pair] += freq
        letter_freqs[split[-1]] += freq

    scores = {
        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 = pair
        max_score = score

print(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:
            continue
        i = 0
        while i < len(split) - 1:
            if split[i] == a and split[i + 1] == b:
                merge = a + b[2:] if b.startswith("##") else a + b
                split = split[:i] + [merge] + split[i + 2 :]
            else:
                i += 1
        splits[word] = split
    return 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 -= 1
        if 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, [])

测试：