LLM 课程文档

WordPiece 分词

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

WordPiece 分词

Ask a Question Open In Colab Open In Studio Lab

WordPiece 是 Google 为预训练 BERT 而开发的分词算法。此后,它被 BERT 衍生的许多 Transformer 模型重用,例如 DistilBERT、MobileBERT、Funnel Transformers 和 MPNET。它在训练方面与 BPE 非常相似,但实际的分词方式不同。

💡 本节深入探讨 WordPiece,甚至展示了完整的实现。如果您只想大致了解分词算法,可以跳到末尾。

训练算法

⚠️ Google 从未开源其 WordPiece 训练算法的实现,因此以下内容是我们根据已发布的文献做出的最佳猜测。它可能不是 100% 准确。

与 BPE 类似,WordPiece 从一个包含模型使用的特殊标记和初始字母表的小词汇表开始。由于它通过添加前缀(例如 BERT 的 ##)来识别子词,因此每个词最初都是通过向词内所有字符添加该前缀来拆分的。例如,"word" 会这样拆分

w ##o ##r ##d

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

然后,再次与 BPE 类似,WordPiece 学习合并规则。主要区别在于选择要合并的对的方式。WordPiece 不选择最频繁的对,而是使用以下公式计算每对的分数score=(freq_of_pair)/(freq_of_first_element×freq_of_second_element)\mathrm{score} = (\mathrm{freq\_of\_pair}) / (\mathrm{freq\_of\_first\_element} \times \mathrm{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" 添加到词汇表中,并在语料库的单词中应用合并

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,我们会按顺序应用学到的合并规则,并将其分词为 ["hu", "##gs"],因此编码是不同的。

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

当分词达到无法在词汇表中找到子词的阶段时,整个单词被标记为未知——例如,"mug" 将被标记为 ["[UNK]"]"bum" 也是(即使我们可以从 "b""##u" 开始,"##m" 不在词汇表中,最终的分词将只是 ["[UNK]"],而不是 ["b", "##u", "[UNK]"])。这是与 BPE 的另一个区别,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。

< > 在 GitHub 上更新