1. 算法原理介绍

        与 BPE 分词（参考《BPE原理及代码示例》）和 WordPiece 分词（参考《WordPiece原理及代码示例》）不同，Unigram 分词方法【1】是从一个包含足够多字符串或词元的初始集合开始，迭代地删除其中的词元，直到达到预期的词表大小。该方法假设通过删除某个词元能够增加训练语料的似然性，并以此作为选择标准。这个过程依赖于一个训练好的单一语言模型。

        为了估计单一语言模型，在每次迭代中，首先根据旧的语言模型找到当前最优的分词方式，然后重新估计单一概率以更新语言模型。通常，使用动态规划算法（如维特比算法，Viterbi Algorithm）来高效找到语言模型对词汇的最优分词方案。

        Unigram算法通常用于SentencePiece，这是AlBERT、T5、mBART、Big Bird和XLNet等模型使用的分词算法。如前述，Unigram训练算法与BPE和WordPiece相比，Unigram的工作方式相反：它从一个大词汇开始，逐步删除其中的标记，直到达到所需的词汇大小。可以使用多种方法构建基础词汇，例如提取预分词单词中最常见的子串，或在具有较大词汇量的初始语料库上应用BPE。

        在训练的每一步，Unigram算法根据当前词汇计算语料库的损失。然后，对于词汇中的每个符号，算法计算如果移除该符号整体损失将增加多少，并寻找对整体损失影响最小的符号。这些符号对语料库的整体损失影响较小，因此在某种意义上，它们是“需求较低”的，最适合被删除。可以预见，这个过程非常耗时。需要注意基本字符绝不删除，以确保任何单词都可以被分词。

2. 案例说明

        重用在BPE中的语料库：

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

        在这个例子中，将所有严格子串作为初始词汇：

         ["h", "u", "g", "hu", "ug", "p", "pu", "n", "un", "b", "bu", "s", "hug", "gs", "ugs"]

        分词算法 Unigram模型是一种语言模型，认为每个标记与之前的标记是独立的。这是最简单的语言模型，因为给定之前的上下文，标记X的概率仅为标记X的概率。因此，如果使用Unigram语言模型生成文本，总是会预测最常见的标记。给定标记的概率是其在原始语料库中的频率（出现次数），除以词汇中所有标记频率的总和（以确保概率总和为1）。例如，“ug”出现在“hug”、“pug”和“hugs”中，因此在当前语料库中它的频率为20。

        以下是词汇中所有可能子词的频率：

 ("h", 15) ("u", 36) ("g", 20) ("hu", 15) ("ug", 20) ("p", 17) ("pu", 17) ("n", 16) ("un", 16) ("b", 4) ("bu", 4) ("s", 5) ("hug", 15) ("gs", 5) ("ugs", 5)

        因此，所有频率的总和为210，而子词“ug”的概率为20/210。

        现在，为了对给定单词进行分词，查看所有可能的分词方式，并根据Unigram模型计算每种方式的概率。由于所有标记被视为独立，因此该概率只是每个标记概率的乘积。例如，单词“pug”的分词["p", "u", "g"]的概率为：

P ( [ "p" , "u" , "g" ] ) = P ( "p" ) × P ( "u" ) × P ( "g" ) = 5/210 × 36/210 × 20/210= 0.000389 

        相比之下，分词["pu", "g"]的概率为：

        P ( [ "pu" , "g" ] ) = P ( "pu" ) × P ( "g" ) = 5/210 × 20/210 = 0.0022676

        因此，这种分词的概率要高得多。一般来说，标记数量最少的分词将具有最高的概率（因为对于每个标记都有210的除法），这对应于直观上的目的：将单词分成尽可能少的标记。

        使用Unigram模型对单词进行分词的结果是概率最高的分词。在“pug”的例子中，为每个可能的分词计算的概率如下：

         ["p", "u", "g"] : 0.000389

         ["p", "ug"] : 0.0022676

         ["pu", "g"] : 0.0022676

        因此，“pug”的分词将为["p", "ug"]或["pu", "g"]，具体取决于首先遇到的分词。

        在这个情况下，很容易找到所有可能的分词并计算它们的概率，但通常这会更困难。采用维特比算法，可以构建一个图来检测给定单词的可能分词，通过说如果从字符a到字符b有一个子词，则从a到b有一个分支，并将子词的概率赋予该分支。

        为了找到在该图中具有最佳分数的路径，维特比算法确定每个单词位置的最佳分词，它结束于该位置。由于是从开始到结束，因此通过遍历所有以当前字符结束的子词来找到最佳分数，然后使用该子词起始位置的最佳分词分数。接下来，只需解开到达结尾的路径。

        看一个使用我们的词汇和单词“unhug”的例子。对于每个位置，最佳分数的子词如下：

        字符0 (u): "u" (分数0.171429)

        字符1 (n): "un" (分数0.076191)

        字符2 (h): "un" "h" (分数0.005442)

        字符3 (u): "un" "hu" (分数0.005442)

        字符4 (g): "un" "hug" (分数0.005442)

        因此，“unhug”的分词将为["un", "hug"]。

        基于上述例子，已经了解了分词是如何工作的，进一步探讨训练期间使用的损失函数。在任何给定阶段，这个损失是通过对语料库中的每个单词进行分词计算的，使用当前的词汇表和根据每个标记在语料库中的频率确定的单语模型。

        语料库中的每个单词都有一个分数，损失则是这些分数的负对数似然——即语料库中所有单词的 -log(P(单词)) 的总和。

        回到以下示例语料库：

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

        每个单词的分词及其相应的分数为：

"hug": ["hug"] (分数 0.071428)
"pug": ["pu", "g"] (分数 0.007710)
"pun": ["pu", "n"] (分数 0.006168)
"bun": ["bu", "n"] (分数 0.001451)
"hugs": ["hug", "s"] (分数 0.001701)

        因此，损失为：

10 * (-log(0.071428)) + 5 * (-log(0.007710)) + 12 * (-log(0.006168)) + 4 * (-log(0.001451)) + 5 * (-log(0.001701)) = 169.8

        现在需要计算每个标记的移除如何影响损失。为了简化，只为两个标记进行计算，并将整个过程留到有代码帮助时再进行。在这个特殊的情况下，有两个等价的标记化：例如，"pug" 可以标记为 ["p", "ug"]，并具有相同的分数。因此，从词汇表中移除 "pu" 将产生相同的损失。

        另一方面，移除 "hug" 将导致损失增加，因为 "hug" 和 "hugs" 的标记化将变为：

"hug": ["hu", "g"] (分数 0.006802)
"hugs": ["hu", "gs"] (分数 0.001701)

        这些变化将导致损失上升：

10 * (-log(0.071428)) + 10 * (-log(0.006802)) = 23.5

        因此，标记 "pu" 可能会从词汇表中移除，但 "hug" 则不会。

3. 代码实现

        依然沿用在BPE中的语料，在jupyter环境进行测试。本来我们验证的预训练模型为xlnet-base-cased，但model scope没有对应的，所以我们更换成T5，效果是一样的，因为使用的分词器是一样的。

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

同样，下载和加载T5预训练模型：

import torch
from modelscope import snapshot_download, AutoModel, AutoTokenizer
import os
from transformers import AutoTokenizer
 
model_dir = snapshot_download('AI-ModelScope/t5-base', cache_dir='/root/autodl-tmp', revision='master')
mode_name_or_path = '/root/autodl-tmp/AI-ModelScope/t5-base'
tokenizer = AutoTokenizer.from_pretrained(mode_name_or_path, trust_remote_code=True)

统计预分词的词频：

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

        接下来，需要将词汇初始化为一个比我们最终希望的词汇量更大的值。必须包含所有基本字符（否则将无法标记每个单词），但对于较大的子字符串，将仅保留最常见的，因此按频率对它们进行排序：        

char_freqs = defaultdict(int)
subwords_freqs = defaultdict(int)
for word, freq in word_freqs.items():
    for i in range(len(word)):
        char_freqs[word[i]] += freq
        # 循环遍历至少长度为2的子单词
        for j in range(i + 2, len(word) + 1):
            subwords_freqs[word[i:j]] += freq

# 按频率对子单词进行排序
sorted_subwords = sorted(subwords_freqs.items(), key=lambda x: x[1], reverse=True)
sorted_subwords[:10]

        将字符与最佳子单词分组，以达到初始词汇大小 300：

token_freqs = list(char_freqs.items()) + sorted_subwords[: 300 - len(char_freqs)]
token_freqs = {token: freq for token, freq in token_freqs}

        接下来，计算所有频率的总和，将频率转换为概率。对于我们的模型，将存储概率的对数，因为将对数相加比乘以小数字更数值稳定，并且这将简化模型损失的计算：

from math import log

total_sum = sum([freq for token, freq in token_freqs.items()])
model = {token: -log(freq / total_sum) for token, freq in token_freqs.items()}

        现在，主要函数是使用维特比算法进行单词标记化。如前所述，该算法计算每个子字符串的最佳分段，将其存储在名为 best_segmentations 的变量中。为单词中的每个位置（从 0 到单词的总长度）存储一个字典，包含两个键：最佳分段中最后一个标记的起始索引，以及最佳分段的分数。通过最后一个标记的起始索引，将能够在列表完全填充后检索完整的分段。

        填充列表只需两个循环：主循环遍历每个起始位置，第二个循环尝试从该起始位置开始的所有子字符串。如果子字符串在词汇中，我们就得到了到该结束位置的单词的新分段，将其与 best_segmentations 中的内容进行比较。一旦主循环完成，只需从末尾开始，跳过一个起始位置，记录标记，直到到达单词的起始位置：        

def encode_word(word, model):
    best_segmentations = [{"start": 0, "score": 1}] + [
        {"start": None, "score": None} for _ in range(len(word))
    ]
    for start_idx in range(len(word)):
        # This should be properly filled by the previous steps of the loop
        best_score_at_start = best_segmentations[start_idx]["score"]
        for end_idx in range(start_idx + 1, len(word) + 1):
            token = word[start_idx:end_idx]
            if token in model and best_score_at_start is not None:
                score = model[token] + best_score_at_start
                # If we have found a better segmentation ending at end_idx, we update
                if (
                    best_segmentations[end_idx]["score"] is None
                    or best_segmentations[end_idx]["score"] > score
                ):
                    best_segmentations[end_idx] = {"start": start_idx, "score": score}

    segmentation = best_segmentations[-1]
    if segmentation["score"] is None:
        # We did not find a tokenization of the word -> unknown
        return ["<unk>"], None

    score = segmentation["score"]
    start = segmentation["start"]
    end = len(word)
    tokens = []
    while start != 0:
        tokens.insert(0, word[start:end])
        next_start = best_segmentations[start]["start"]
        end = start
        start = next_start
    tokens.insert(0, word[start:end])
    return tokens, score

在一些单词上尝试下初始模型：

        

计算模型在语料库上的损失：

def compute_loss(model):
    loss = 0
    for word, freq in word_freqs.items():
        _, word_loss = encode_word(word, model)
        loss += freq * word_loss
    return loss

计算删除每个标记后模型的损失：

import copy

def compute_scores(model):
    scores = {}
    model_loss = compute_loss(model)
    for token, score in model.items():
        # 始终保留长度为1的标记
        if len(token) == 1:
            continue
        model_without_token = copy.deepcopy(model)
        _ = model_without_token.pop(token)
        scores[token] = compute_loss(model_without_token) - model_loss
    return scores

        由于 "ll" 用于 "Hopefully" 的标记化，移除它可能会导致使用标记 "l" 两次，因此预期它会产生正损失。

        这种方法效率非常低，因此 SentencePiece 使用了一个近似的模型损失计算：它并不是从头开始，而是用剩余词汇中标记 X 的分割替换标记 X。这样，所有分数可以在计算模型损失的同时一次性计算出来。在此基础上，需要做的最后一件事是将模型使用的特殊标记添加到词汇表中，然后循环直到我们从词汇表中修剪出足够的标记以达到所需大小：

percent_to_remove = 0.1
while len(model) > 100:
    scores = compute_scores(model)
    sorted_scores = sorted(scores.items(), key=lambda x: x[1])
    # 移除具有最低分数的 percent_to_remove 标记。
    for i in range(int(len(model) * percent_to_remove)):
        _ = token_freqs.pop(sorted_scores[i][0])

    total_sum = sum([freq for token, freq in token_freqs.items()])
    model = {token: -log(freq / total_sum) for token, freq in token_freqs.items()}

        然后，为了对一些文本进行标记化，只需应用预标记化，然后使用 encode_word() 函数：

def tokenize(text, model):
    words_with_offsets = tokenizer.backend_tokenizer.pre_tokenizer.pre_tokenize_str(text)
    pre_tokenized_text = [word for word, offset in words_with_offsets]
    encoded_words = [encode_word(word, model)[0] for word in pre_tokenized_text]
    return sum(encoded_words, [])


tokenize("This is the amazing course.", model)

4. 参考材料

【1】Unigram tokenization