HuggingfaceNLP笔记6.6WordPiece tokenization

WordPiece是Google为预训练BERT开发的分词算法，它被应用于基于BERT的Transformer模型，如DistilBERT、MobileBERT、Funnel Transformers和MPNET等。它与BPE在训练方法上相似，但实际的分词过程有所不同。

💡 这部分深入讲解WordPiece，甚至会展示完整的实现。如果你只想了解分词算法的概要，可以跳到最后。

训练算法

⚠️ Google并未公开其WordPiece训练算法的实现，以下内容基于已发表文献的推测，可能不完全准确。

与BPE类似，WordPiece从包含模型特殊符号和初始字母表的小词汇表开始。由于它通过添加前缀（如BERT中的##）来识别子词，每个单词最初都是通过在单词内部添加该前缀来分割的。例如，"word"会被这样分割：

w ##o ##r ##d

因此，初始字母表包含单词开头的所有字符以及带有WordPiece前缀的单词内部字符。

然后，与BPE类似，WordPiece学习合并规则。主要区别在于合并对的选择方式。WordPiece不会选择最频繁的对，而是为每对计算分数，使用以下公式：
score=(freq_of_pair)/(freq_of_first_element×freq_of_second_element)。

通过将对的频率除以每个部分频率的乘积，算法优先合并那些部分在词汇表中频率较低的对。例如，即使("un", "##able")这对在词汇表中出现频率很高，WordPiece也不一定会将其合并，因为"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" "##gs", 5)

因此，初始词汇表将是["b", "h", "p", "##g", "##n", "##s", "##u"]（忽略特殊符号）。最频繁的对是("##u", "##g")（出现20次），但由于"##u"的频率很高，其分数不是最高的（为1/36）。所有包含"##u"的对实际上都有相同的分数（1/36），所以最好的分数给了没有"##u"的对("##g", "##s")，分数为1/20，第一个学习的合并是("##g", "##s") -> ("##gs")。

请注意，当我们合并时，会移除两个标记之间的##，所以我们添加"##gs"到词汇表中，并在语料库中的单词中应用合并：

Vocabulary: ["b", "h", "p", "##g", "##n", "##s", "##u", "##gs"]
Corpus: ("h" "##u" "##g", 10), ("p" "##u" "##g", 5), ("p" "##u" "##n", 12), ("b" "##u" "##n", 4), ("h" "##u" "##gs", 5)

到目前为止，所有可能的对"##u"都有相同的分数。假设在这种情况下，我们首先合并第一个对，即("h", "##u") -> "hu"。这样我们就到了：

Vocabulary: ["b", "h", "p", "##g", "##n", "##s", "##u", "##gs", "hu"]
Corpus: ("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），所以我们首先合并分数最高的对：

Vocabulary: ["b", "h", "p", "##g", "##n", "##s", "##u", "##gs", "hu", "hug"]
Corpus: ("hug", 10), ("p" "##u" "##g", 5), ("p" "##u" "##n", 12), ("b" "##u" "##n", 4), ("hu" "##gs", 5)

然后我们继续这个过程，直到达到所需的词汇表大小。

📝 现在轮到你了！ 下一步的合并规则会是什么？

分词算法

WordPiece和BPE在分词方面的不同在于，WordPiece只保存最终的词汇表，而不保存学习到的合并规则。从要分词的单词开始，WordPiece会找到词汇表中以该单词开头的最长子词，然后在该位置进行分割。例如，如果我们使用上述示例中的词汇表，对于单词"hugs"，从开头开始的最长在词汇表中的子词是"hug"，所以我们在此处分割，得到["hug", "##s"]。然后我们继续处理"##s"，它也在词汇表中，所以"hugs"的分词结果是["hug", "##s"]。

BPE则会应用学到的合并规则，将"hugs"分词为["hu", "##gs"]，因此编码方式不同。

再举个例子，我们来看看单词"bugs"的分词过程。从单词开头开始，最长在词汇表中的子词是"b"，所以我们在此处分割，得到["b", "##ugs"]。然后在"##ugs"中，最长的以开头的子词是"##u"，所以我们在此分割，得到["b", "##u", "##gs"]。最后，"##gs"在词汇表中，所以这个列表就是"bugs"的分词结果。

当分词到达无法在词汇表中找到子词的阶段时，整个单词将被标记为未知——例如，单词"mug"会被分词为["[UNK]"]，就像"bum"（尽管我们可以从"b"和"##u"开始，但"##m"不在词汇表中，最终分词结果将是["[UNK]"]，而不是["b", "##u", "[UNK]"]）。这与BPE不同，它只会将单个不在词汇表中的字符标记为未知。

📝 现在轮到你了！ "pugs"这个单词会被如何分词？

实现WordPiece

现在让我们来看看WordPiece算法的实现。就像BPE一样，这只是一个教学示例，你无法在大型语料库上使用它。

我们将使用与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分词器进行预处理：

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")

然后，我们像预处理一样计算语料库中每个单词的频率：

from collections import defaultdict

word_freqs = defaultdict(int)
for text in corpus:
    words_with_offsets = tokenizer.backend_tokenizer.pre_tokenizer.pre_tokenize_str(text)
    new_words = [word for word, offset in words_with_offsets]
    for word in new_words:
        word_freqs[word] += 1

word_freqs

defaultdict(
    int, {'This': 3, 'is': 2, 'the': 1, 'Hugging': 1, 'Face': 1, 'Course': 1, '.': 4, 'chapter': 1, 'about': 1,
    'tokenization': 1, 'section': 1, 'shows': 1, 'several': 1, 'tokenizer': 1, 'algorithms': 1, 'Hopefully': 1,
    ',': 1, 'you': 1, 'will': 1, 'be': 1, 'able': 1, 'to': 1, 'understand': 1, 'how': 1, 'they': 1, 'are': 1,
    'trained': 1, 'and': 1, 'generate': 1, 'tokens': 1})

正如我们之前所见，字母表是由所有单词的首字母以及以##开头的其他字母组成的唯一集合：

alphabet = []
for word in word_freqs.keys():
    if word[0] not in alphabet:
        alphabet.append(word[0])
    for letter in word[1:]:
        if f"##{letter}" not in alphabet:
            alphabet.append(f"##{letter}")

alphabet.sort()
alphabet

print(alphabet)

['##a', '##b', '##c', '##d', '##e', '##f', '##g', '##h', '##i', '##k', '##l', '##m', '##n', '##o', '##p', '##r', '##s',
 '##t', '##u', '##v', '##w', '##y', '##z', ',', '.', 'C', 'F', 'H', 'T', 'a', 'b', 'c', 'g', 'h', 'i', 's', 't', 'u',
 'w', 'y']

我们还需要在词汇表的开头添加模型使用的特殊标记。对于BERT，它是一个列表["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]：

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

接下来，我们需要将每个单词拆分成单个字母，所有非首字母前都加上##：

splits = {
    word: [c if i == 0 else f"##{c}" for i, c in enumerate(word)]
    for word in word_freqs.keys()
}

现在我们已经准备好训练了，让我们编写一个函数来计算每对字母的分数。在训练的每个步骤中，我们都需要使用这个函数：

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)
for i, key in enumerate(pair_scores.keys()):
    print(f"{key}: {pair_scores[key]}")
    if i >= 5:
        break

('T', '##h'): 0.125
('##h', '##i'): 0.03409090909090909
('##i', '##s'): 0.02727272727272727
('i', '##s'): 0.1
('t', '##h'): 0.03571428571428571
('##h', '##e'): 0.011904761904761904

现在找到最高分的对只需要快速遍历：

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') 0.2

所以第一个要学习的合并是('a', '##b') -> 'ab'，然后我们将’ab’添加到词汇表中：

vocab.append("ab")

接下来，我们需要在splits字典中应用这个合并。为此，我们可以定义一个函数：

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"]

['ab', '##o', '##u', '##t']

现在我们有了继续学习所有合并所需的工具。目标是达到70个词汇大小：

vocab_size = 70
while len(vocab) < vocab_size:
    scores = compute_pair_scores(splits)
    best_pair, max_score = "", None
    for pair, score in scores.items():
        if max_score is None or max_score < score:
            best_pair = pair
            max_score = score
    splits = merge_pair(*best_pair, splits)
    new_token = (
        best_pair[0] + best_pair[1][2:]
        if best_pair[1].startswith("##")
        else best_pair[0] + best_pair[1]
    )
    vocab.append(new_token)

然后我们可以查看生成的词汇表：

print(vocab)

['[PAD]', '[UNK]', '[CLS]', '[SEP]', '[MASK]', '##a', '##b', '##c', '##d', '##e', '##f', '##g', '##h', '##i', '##k',
 '##l', '##m', '##n', '##o', '##p', '##r', '##s', '##t', '##u', '##v', '##w', '##y', '##z', ',', '.', 'C', 'F', 'H',
 'T', 'a', 'b', 'c', 'g', 'h', 'i', 's', 't', 'u', 'w', 'y', 'ab', '##fu', 'Fa', 'Fac', '##ct', '##ful', '##full', '##fully',
 'Th', 'ch', '##hm', 'cha', 'chap', 'chapt', '##thm', 'Hu', 'Hug', 'Hugg', 'sh', 'th', 'is', '##thms', '##za', '##zat',
 '##ut']

如你所见，与BPE相比，这个分词器学习单词的部分更快。

💡 如果在相同的语料库上使用train_new_from_iterator()，不会得到完全相同的词汇表。这是因为🤗 Tokenizers库并未实现WordPiece的训练（因为不清楚其内部细节），而是使用了BPE替代。

要对新文本进行分词，我们首先预处理它，然后将其拆分成单词，接着在每个单词上应用分词算法。具体来说，我们从第一个单词的开头开始查找最大的子词并将其拆分，然后对第二个部分重复此过程，直到处理完该单词及其后的文本：

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("HOgging"))

['Hugg', '##i', '##n', '##g']
['[UNK]']

现在，我们编写一个函数来对文本进行分词：

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, [])

我们可以用任何文本来测试它：

tokenize("This is the Hugging Face course!")

['Th', '##i', '##s', 'is', 'th', '##e', 'Hugg', '##i', '##n', '##g', 'Fac', '##e', 'c', '##o', '##u', '##r', '##s', '##e', '[UNK]']

这就是WordPiece算法的全部内容！接下来，我们来看看Unigram。