Blog of TheChef
0%

LLM之分词器(tokenizer)

发表于 分类于

LLM之分词器(tokenizer)

一、引言

大语言模型在自然语言处理领域取得了巨大成功，而分词器作为其关键组成部分，对模型的性能和效果有着显著影响。今天，我们将探讨四种常见的分词方法：BPE、WordPiece、SentencePiece 和 unigram，剖析它们的技术原理、实现细节及应用场景。

在此之前，我们需要了解一下什么是tokenizition。

任何一段文本，输入给模型，都是要转换成一串embedding。这个过程简单概括为：

  1. 分词，并把词转换为token（即词的ID）
  2. token转换成embedding

而tokenization就是在做这第一步。而对于第二步就是常见的Embedding查表操作，即根据token_id的值，去Embedding矩阵中查找第token_id行的数据作为embedding。

二、方法介绍

1. BPE

BPE 是一种基于频率的分词算法，最初用于机器翻译任务中的词汇扩展。其核心思想是通过不断合并文本中出现频率最高的字节对来构建词汇表。具体来说，首先将文本中的每个字符视为一个独立的 token，然后统计所有相邻字符对的频率，选择最频繁的字符对进行合并，并更新词汇表和文本表示，重复这一过程直到达到预设的词汇表大小。其实现方法如下

  1. 初始化词汇表：以字符为粒度，将文本中的每个字符作为初始词汇表的元素。
  2. 统计字节对频率：遍历文本，统计所有相邻字符对的出现次数。
  3. 合并高频字节对：选择频率最高的字节对进行合并，生成新的 token，并更新词汇表和文本表示。
  4. 迭代更新：重复统计频率和合并操作，直到词汇表大小达到设定值。

Byte-Pair Encoding(BPE)是最广泛采用的subword分词器。

  • 训练方法：从字符级的小词表出发，训练产生合并规则以及一个词表
  • 编码方法：将文本切分成字符，再应用训练阶段获得的合并规则
  • 经典模型：GPT, GPT-2, RoBERTa, BART, LLaMA, ChatGLM等

1.1. 训练阶段

在训练环节，目标是给定语料，通过训练算法，生成合并规则和词表。 BPE算法是从一个字符级别的词表为基础，合并pair并添加到词表中，逐步形成大词表。合并规则为选择相邻pair词频最大的进行合并。

假定训练的语料(已归一化处理)为4个句子。

1
2
3
4
5
6
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.",
]

首先进行预切分处理。这里采用gpt2的预切分逻辑。 具体会按照空格和标点进行切分，并且空格会保留成特殊的字符“Ġ”。

1
2
3
4
5
6
7
8
from transformers import AutoTokenizer

# init pre tokenize function
gpt2_tokenizer = AutoTokenizer.from_pretrained("gpt2")
pre_tokenize_function = gpt2_tokenizer.backend_tokenizer.pre_tokenizer.pre_tokenize_str

# pre tokenize
pre_tokenized_corpus = [pre_tokenize_str(text) for text in corpus]

获得的pre_tokenized_corpus如下，每个单元分别为[word, (start_index, end_index)]

1
2
3
4
5
6
[
    [('This', (0, 4)), ('Ġis', (4, 7)), ('Ġthe', (7, 11)), ('ĠHugging', (11, 19)), ('ĠFace', (19, 24)), ('ĠCourse', (24, 31)), ('.', (31, 32))], 
    [('This', (0, 4)), ('Ġchapter', (4, 12)), ('Ġis', (12, 15)), ('Ġabout', (15, 21)), ('Ġtokenization', (21, 34)), ('.', (34, 35))], 
    [('This', (0, 4)), ('Ġsection', (4, 12)), ('Ġshows', (12, 18)), ('Ġseveral', (18, 26)), ('Ġtokenizer', (26, 36)), ('Ġalgorithms', (36, 47)), ('.', (47, 48))], 
    [('Hopefully', (0, 9)), (',', (9, 10)), ('Ġyou', (10, 14)), ('Ġwill', (14, 19)), ('Ġbe', (19, 22)), ('Ġable', (22, 27)), ('Ġto', (27, 30)), ('Ġunderstand', (30, 41)), ('Ġhow', (41, 45)), ('Ġthey', (45, 50)), ('Ġare', (50, 54)), ('Ġtrained', (54, 62)), ('Ġand', (62, 66)), ('Ġgenerate', (66, 75)), ('Ġtokens', (75, 82)), ('.', (82, 83))]
]

进一步统计每个整词的词频

1
2
3
4
word2count = defaultdict(int)
for split_text in pre_tokenized_corpus:
    for word, _ in split_text:
        word2count[word] += 1

获得word2count如下

1
defaultdict(<class '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})

因为BPE是从字符级别的小词表，逐步合并成大词表，所以需要先获得字符级别的小词表。

1
2
3
4
vocab_set = set()
for word in word2count:
    vocab_set.update(list(word))
vocabs = list(vocab_set)

获得的初始小词表vocabs如下:

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

基于小词表就可以对每个整词进行切分

1
2
3
4
5
6
7
8
9
word2splits = {word: [c for c in word] for word in word2count}

'This': ['T', 'h', 'i', 's'], 
'Ġis': ['Ġ', 'i', 's'], 
'Ġthe': ['Ġ', 't', 'h', 'e'], 
...
'Ġand': ['Ġ', 'a', 'n', 'd'], 
'Ġgenerate': ['Ġ', 'g', 'e', 'n', 'e', 'r', 'a', 't', 'e'], 
'Ġtokens': ['Ġ', 't', 'o', 'k', 'e', 'n', 's']

基于word2splits统计vocabs中相邻两个pair的词频pair2count

1
2
3
4
5
6
7
8
9
10
def _compute_pair2score(word2splits, word2count):
    pair2count = defaultdict(int)
    for word, word_count in word2count.items():
        split = word2splits[word]
        if len(split) == 1:
            continue
        for i in range(len(split) - 1):
            pair = (split[i], split[i + 1])
            pair2count[pair] += word_count
    return pair2count

获得pair2count如下：

1
defaultdict(<class 'int'>, {('T', 'h'): 3, ('h', 'i'): 3, ('i', 's'): 5, ('Ġ', 'i'): 2, ('Ġ', 't'): 7, ('t', 'h'): 3, ..., ('n', 's'): 1})

统计当前频率最高的相邻pair

1
2
3
4
5
6
7
8
def _compute_most_score_pair(pair2count):
    best_pair = None
    max_freq = None
    for pair, freq in pair2count.items():
        if max_freq is None or max_freq < freq:
            best_pair = pair
            max_freq = freq
    return best_pair

经过统计，当前频率最高的pair为: (‘Ġ’, ‘t’)， 频率为7次。 将(‘Ġ’, ‘t’)合并成一个词并添加到词表中。同时在合并规则中添加(‘Ġ’, ‘t’)这条合并规则。

1
2
3
4
merge_rules = []
best_pair = self._compute_most_score_pair(pair2score)
vocabs.append(best_pair[0] + best_pair[1])
merge_rules.append(best_pair)

此时的vocab词表更新成:

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

根据更新后的vocab重新对word2count进行切分。具体实现上，可以直接在旧的word2split上应用新的合并规则(‘Ġ’, ‘t’)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
def _merge_pair(a, b, word2splits):
    new_word2splits = dict()
    for word, split in word2splits.items():
        if len(split) == 1:
            new_word2splits[word] = split
            continue
        i = 0
        while i < len(split) - 1:
            if split[i] == a and split[i + 1] == b:
                split = split[:i] + [a + b] + split[i + 2:]
            else:
                i += 1
        new_word2splits[word] = split
    return new_word2splits

从而获得新的word2split

1
2
3
4
5
6
{'This': ['T', 'h', 'i', 's'], 
'Ġis': ['Ġ', 'i', 's'], 
'Ġthe': ['Ġt', 'h', 'e'], 
'ĠHugging': ['Ġ', 'H', 'u', 'g', 'g', 'i', 'n', 'g'],
...
'Ġtokens': ['Ġt', 'o', 'k', 'e', 'n', 's']}

可以看到新的word2split中已经包含了新的词”Ġt”。

重复上述循环直到整个词表的大小达到预先设定的词表大小。

1
2
3
4
5
6
while len(vocabs) < vocab_size:
    pair2score = self._compute_pair2score(word2splits, word2count)
    best_pair = self._compute_most_score_pair(pair2score)
    vocabs.append(best_pair[0] + best_pair[1])
    merge_rules.append(best_pair)
    word2splits = self._merge_pair(best_pair[0], best_pair[1], word2splits)

假定最终词表的大小为50，经过上述迭代后我们获得的词表和合并规则如下：

1
2
3
vocabs = ['i', 't', 'p', 'o', 'r', 'm', 'e', ',', 'y', 'v', 'Ġ', 'F', 'a', 'C', 'H', '.', 'f', 'l', 'u', 'c', 'T', 'k', 'h', 'z', 'd', 'g', 'w', 'n', 's', 'b', 'Ġt', 'is', 'er', 'Ġa', 'Ġto', 'en', 'Th', 'This', 'ou', 'se', 'Ġtok', 'Ġtoken', 'nd', 'Ġis', 'Ġth', 'Ġthe', 'in', 'Ġab', 'Ġtokeni', 'Ġtokeniz']

merge_rules = [('Ġ', 't'), ('i', 's'), ('e', 'r'), ('Ġ', 'a'), ('Ġt', 'o'), ('e', 'n'), ('T', 'h'), ('Th', 'is'), ('o', 'u'), ('s', 'e'), ('Ġto', 'k'), ('Ġtok', 'en'), ('n', 'd'), ('Ġ', 'is'), ('Ġt', 'h'), ('Ġth', 'e'), ('i', 'n'), ('Ġa', 'b'), ('Ġtoken', 'i'), ('Ġtokeni', 'z')]

至此我们就根据给定的语料完成了BPE分词器的训练。

1.2. 推理阶段

在推理阶段，给定一个句子，我们需要将其切分成一个token的序列。 具体实现上需要先对句子进行预分词并切分成字符级别的序列，然后根据合并规则进行合并。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
def tokenize(self, text: str) -> List[str]:
    # pre tokenize
    words = [word for word, _ in self.pre_tokenize_str(text)]
    # split into char level
    splits = [[c for c in word] for word in words]
    # apply merge rules
    for merge_rule in self.merge_rules:
        for index, split in enumerate(splits):
            i = 0
            while i < len(split) - 1:
                if split[i] == merge_rule[0] and split[i + 1] == merge_rule[1]:
                    split = split[:i] + ["".join(merge_rule)] + split[i + 2:]
                else:
                    i += 1
            splits[index] = split
    return sum(splits, [])

例如

1
2
>>> tokenize("This is not a token.")
>>> ['This', 'Ġis', 'Ġ', 'n', 'o', 't', 'Ġa', 'Ġtoken', '.']

2. WordPiece

WordPiece 也是一种基于频率的分词方法，与 BPE 不同的是，它在选择合并单元时，不仅考虑字节对的出现频率，还引入了语言模型的似然估计。其目标是最小化语言模型的困惑度，即选择使语言模型概率最大的子词划分方式。[只是在训练阶段合并pair的策略不是pair的频率而是互信息。]

1
socre=log(p(ab))−(log(p(a))+log(p(b)))=log(p(ab)/p(a)p(b))

这里的动机是一个pair的频率很高，但是其中pair的一部分的频率更高，这时候不一定需要进行该pair的合并。 而如果一个pair的频率很高，并且这个pair的两个部分都是只出现在这个pair中，就说明这个pair很值得合并。实现步骤如下：

  1. 初始化词汇表：通常以字符为初始 token。
  2. 训练语言模型：使用初始词汇表对文本进行编码，并训练一个语言模型。
  3. 寻找最优合并：在每次迭代中，尝试所有可能的子词对合并，计算合并后的语言模型困惑度，选择使困惑度最小的合并对。
  4. 更新词汇表和语言模型：将新合并的子词加入词汇表，并重新训练语言模型。
  5. 重复迭代：直到达到预设的词汇表大小或困惑度不再显著降低。
  • 训练方法：从字符级的小词表出发，训练产生合并规则以及一个词表
  • 编码方法：将文本切分成词，对每个词在词表中进行最大前向匹配
  • 经典模型：BERT及其系列DistilBERT，MobileBERT等

2.1. 训练阶段

在训练环节，给定语料，通过训练算法，生成最终的词表。 WordPiece算法也是从一个字符级别的词表为基础，逐步扩充成大词表。合并规则为选择相邻pair互信息最大的进行合并。

下面进行具体手工实现。

假定训练的语料(已归一化处理)为

1
2
3
4
5
6
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.",
]

首先进行预切分处理。这里采用BERT的预切分逻辑。具体会按照空格和标点进行切分。

1
2
3
4
5
6
7
8
from transformers import AutoTokenizer

# init pre tokenize function
bert_tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
pre_tokenize_function = bert_tokenizer.backend_tokenizer.pre_tokenizer.pre_tokenize_str

# pre tokenize
pre_tokenized_corpus = [pre_tokenize_str(text) for text in corpus]

获得的pre_tokenized_corpus如下，每个单元分别为[word, (start_index, end_index)]

1
2
3
4
5
6
[
    [('This', (0, 4)), ('is', (5, 7)), ('the', (8, 11)), ('Hugging', (12, 19)), ('Face', (20, 24)), ('Course', (25, 31)), ('.', (31, 32))], 
    [('This', (0, 4)), ('chapter', (5, 12)), ('is', (13, 15)), ('about', (16, 21)), ('tokenization', (22, 34)), ('.', (34, 35))], 
    [('This', (0, 4)), ('section', (5, 12)), ('shows', (13, 18)), ('several', (19, 26)), ('tokenizer', (27, 36)), ('algorithms', (37, 47)), ('.', (47, 48))], 
    [('Hopefully', (0, 9)), (',', (9, 10)), ('you', (11, 14)), ('will', (15, 19)), ('be', (20, 22)), ('able', (23, 27)), ('to', (28, 30)), ('understand', (31, 41)), ('how', (42, 45)), ('they', (46, 50)), ('are', (51, 54)), ('trained', (55, 62)), ('and', (63, 66)), ('generate', (67, 75)), ('tokens', (76, 82)), ('.', (82, 83))]
]

进一步统计词频

1
2
3
4
word2count = defaultdict(int)
for split_text in pre_tokenized_corpus:
    for word, _ in split_text:
        word2count[word] += 1

获得word2count如下

1
defaultdict(<class '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})

因为WordPiece同样是从字符级别的小词表，逐步合并成大词表，所以先获得字符级别的小词表。注意这里如果字符不是不一个词的开始，需要添加上特殊字符”##”。

1
2
3
4
5
vocab_set = set()
for word in word2count:
    vocab_set.add(word[0])
    vocab_set.update(['##' + c for c in word[1:]])
vocabs = list(vocab_set)

获得的初始小词表vocabs如下:

1
['##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']

基于小词表对每个词进行切分

1
2
3
4
5
6
7
8
9
word2splits = {word: [word[0]] + ['##' + c for c in word[1:]] for word in word2count}

{'This': ['T', '##h', '##i', '##s'], 
'is': ['i', '##s'], 
'the': ['t', '##h', '##e'], 
'Hugging': ['H', '##u', '##g', '##g', '##i', '##n', '##g'], 
...
'generate': ['g', '##e', '##n', '##e', '##r', '##a', '##t', '##e'], 
'tokens': ['t', '##o', '##k', '##e', '##n', '##s']}

进一步统计vocabs中相邻两个pair的互信息

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
def _compute_pair2score(word2splits, word2count):
    """
    计算每个pair的分数
    score=(freq_of_pair)/(freq_of_first_element×freq_of_second_element)
    :return:
    """
    vocab2count = defaultdict(int)
    pair2count = defaultdict(int)
    for word, word_count in word2count.items():
        splits = word2splits[word]
        if len(splits) == 1:
            vocab2count[splits[0]] += word_count
            continue
        for i in range(len(splits) - 1):
            pair = (splits[i], splits[i + 1])
            vocab2count[splits[i]] += word_count
            pair2count[pair] += word_count
        vocab2count[splits[-1]] += word_count
    scores = {
        pair: freq / (vocab2count[pair[0]] * vocab2count[pair[1]])
        for pair, freq in pair2count.items()
    }
    return scores

获得每个pair的互信息如下：

1
2
3
4
5
6
{('T', '##h'): 0.125, 
('##h', '##i'): 0.03409090909090909, 
('##i', '##s'): 0.02727272727272727, 
('a', '##b'): 0.2,
...
('##n', '##s'): 0.00909090909090909}

统计出互信息最高的相邻pair

1
2
3
4
5
6
7
8
def _compute_most_score_pair(pair2score):
    best_pair = None
    max_score = None
    for pair, score in pair2score.items():
        if max_score is None or max_score < score:
            best_pair = pair
            max_score = score
    return best_pair

此时互信息最高的pair为: (‘a’, ‘##b’) 将(‘a’, ‘##b’)合并成一个词’ab’并添加到词表中

1
2
best_pair = self._compute_most_score_pair(pair2score)
vocabs.append(best_pair[0] + best_pair[1])

这样vocab词表更新成:

1
2
['##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']

根据更新的vocab重新对word2count进行切分。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
def _merge_pair(a, b, word2splits):
    new_word2splits = dict()
    for word, split in word2splits.items():
        if len(split) == 1:
            new_word2splits[word] = split
            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
        new_word2splits[word] = split
    return new_word2splits

获得新的word2split

1
2
3
4
5
{'This': ['T', '##h', '##i', '##s'], 
'is': ['i', '##s'], 'the': ['t', '##h', '##e'], 
'Hugging': ['H', '##u', '##g', '##g', '##i', '##n', '##g'], 
'about': ['ab', '##o', '##u', '##t'], 
'tokens': ['t', '##o', '##k', '##e', '##n', '##s']}

可以看到新的word2split中已经包含了新的词”ab”。

重复上述步骤，直到整个词表的大小达到预先设定的词表大小。

1
2
3
4
5
6
while len(vocabs) < vocab_size:
    pair2score = self._compute_pair2score(word2splits, word2count)
    best_pair = self._compute_most_score_pair(pair2score)
    word2splits = self._merge_pair(best_pair[0], best_pair[1], word2splits)
    new_token = best_pair[0] + best_pair[1][2:] if best_pair[1].startswith('##') else best_pair[1]
    vocabs.append(new_token)

假定最终词表的大小为70，经过上述迭代后我们获得的词表如下：

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

注意词表中添加了特殊的token：[CLS], [MASK], [PAD], [SEP], [UNK] 至此我们就根据给定的语料完成了WordPiece分词器的训练。

2.2. 推理阶段

在推理阶段，给定一个句子，需要将其切分成一个token的序列。 具体实现上需要先对句子进行预分词，然后对每个词进行在词表中进行最大前向的匹配。如果词表中不存在则为UNK。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
def _encode_word(self, word):
    tokens = []
    while len(word) > 0:
        i = len(word)
        while i > 0 and word[:i] not in self.vocabs:
            i -= 1
        if i == 0:
            return ["[UNK]"]
        tokens.append(word[:i])
        word = word[i:]
        if len(word) > 0:
            word = f"##{word}"
    return tokens

def tokenize(self, text):
    words = [word for word, _ in self.pre_tokenize_str(text)]
    encoded_words = [self._encode_word(word) for word in words]
    return sum(encoded_words, [])

例如

1
2
>>> 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]']

3. Unigram

Unigram分词与BPE和WordPiece不同，是基于一个大词表逐步裁剪成一个小词表。 通过Unigram语言模型计算删除不同subword造成的损失来衡量subword的重要性，保留重要性较高的子词。

  • 训练方法：从包含字符和全部子词的大词表出发，逐步裁剪出一个小词表，并且每个词都有自己的分数。
  • 编码方法：将文本切分成词，对每个词基于Viterbi算法求解出最佳解码路径。
  • 经典模型：AlBERT, T5, mBART, Big Bird, XLNet

3.1. 训练阶段

在训练环节，目标是给定语料，通过训练算法，生成最终的词表，并且每个词有自己的概率值。 Unigram算法是从大词表为基础，逐步裁剪成小词表。裁剪规则是根据Unigram语言模型的打分依次裁剪重要度相对较低的词。

下面进行具体手工实现。

假定训练的语料(已归一化处理)为

1
2
3
4
5
6
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.",
]

首先进行预切分处理。这里采用xlnet的预切分逻辑。具体会按照空格进行切分，标点不会切分。并且空格会保留成特殊字符”▁”，句子开头也会添加特殊字符”▁”。

1
2
3
4
5
6
7
8
from transformers import AutoTokenizer

# init pre tokenize function
xlnet_tokenizer = AutoTokenizer.from_pretrained("xlnet-base-cased")
pre_tokenize_function = xlnet_tokenizer.backend_tokenizer.pre_tokenizer.pre_tokenize_str

# pre tokenize
pre_tokenized_corpus = [pre_tokenize_str(text) for text in corpus]

获得的pre_tokenized_corpus如下，每个单元分别为[word, (start_index, end_index)]

1
2
3
4
5
6
[
    [('▁This', (0, 4)), ('▁is', (5, 7)), ('▁the', (8, 11)), ('▁Hugging', (12, 19)), ('▁Face', (20, 24)), ('▁Course.', (25, 32))], 
    [('▁This', (0, 4)), ('▁chapter', (5, 12)), ('▁is', (13, 15)), ('▁about', (16, 21)), ('▁tokenization.', (22, 35))], 
    [('▁This', (0, 4)), ('▁section', (5, 12)), ('▁shows', (13, 18)), ('▁several', (19, 26)), ('▁tokenizer', (27, 36)), ('▁algorithms.', (37, 48))], 
    [('▁Hopefully,', (0, 10)), ('▁you', (11, 14)), ('▁will', (15, 19)), ('▁be', (20, 22)), ('▁able', (23, 27)), ('▁to', (28, 30)), ('▁understand', (31, 41)), ('▁how', (42, 45)), ('▁they', (46, 50)), ('▁are', (51, 54)), ('▁trained', (55, 62)), ('▁and', (63, 66)), ('▁generate', (67, 75)), ('▁tokens.', (76, 83))]
]

进一步统计词频

1
2
3
4
word2count = defaultdict(int)
for split_text in pre_tokenized_corpus:
    for word, _ in split_text:
        word2count[word] += 1

获得word2count如下

1
defaultdict(<class 'int'>, {'▁This': 3, '▁is': 2, '▁the': 1, '▁Hugging': 1, '▁Face': 1, '▁Course.': 1, '▁chapter': 1, '▁about': 1, '▁tokenization.': 1, '▁section': 1, '▁shows': 1, '▁several': 1, '▁tokenizer': 1, '▁algorithms.': 1, '▁Hopefully,': 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})

统计词表的全部子词和词频，取前300个词，构成最初的大词表。为了避免OOV，char级别的词均需要保留。

1
2
3
4
5
6
7
8
9
10
char2count = defaultdict(int)
sub_word2count = defaultdict(int)
for word, count in word2count.items():
    for i in range(len(word)):
        char2count[word[i]] += count
        for j in range(i + 2, len(word) + 1):
            sub_word2count[word[i:j]] += count
sorted_sub_words = sorted(sub_word2count.items(), key=lambda x: x[1], reverse=True)
# init a large vocab with 300
tokens = list(char2count.items()) + sorted_sub_words[: 300 - len(char2count)]

获得的初始小词表vocabs如下:

1
[('▁', 31), ('T', 3), ('h', 9), ('i', 13), ('s', 13), ...,  ('several', 1)]

进一步统计每个子词的概率，并转换成Unigram里的loss贡献

1
2
3
4
5
6
7
8
9
10
11
12
13
token2count = {token: count for token, count in tokens}
total_count = sum([count for token, count in token2count.items()])
model = {token: -log(count / total_count) for token, count in token2count.items()}

model = {
    '▁': 2.952892114877499, 
    'T': 5.288267030694535, 
    'h': 4.189654742026425, 
    ..., 
    'sever': 6.386879319362645, 
    'severa': 6.386879319362645, 
    'several': 6.386879319362645
}

基于每个子词的loss以及Viterbi算法就可以求解出，输入的一个词的最佳分词路径。即整体语言模型的loss最小。词的长度为N，解码的时间复杂度为O(N^2)。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
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 (lower score) ending at end_idx
                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

例如：

1
2
3
4
>>> tokenize("This")
>>> (['This'], 6.288267030694535)
>>> tokenize("this")
>>>(['t', 'his'], 10.03608902044192)

基于上述的函数，可以获得任一个词的分词路径，以及loss。这样就可以计算整个语料上的loss。

1
2
3
4
5
6
def _compute_loss(self, model, word2count):
    loss = 0
    for word, freq in word2count.items():
        _, word_loss = self._encode_word(word, model)
        loss += freq * word_loss
    return loss

尝试移除model中的一个子词，并计算移除后新的model在全部语料上的loss，从而获得这个子词的score，即删除这个子词使得loss新增的量。

1
2
3
4
5
6
7
8
9
10
11
12
13
def _compute_scores(self, model, word2count):
    scores = {}
    model_loss = self._compute_loss(model, word2count)
    for token, score in model.items():
        # We always keep tokens of length 1
        if len(token) == 1:
            continue
        model_without_token = copy.deepcopy(model)
        _ = model_without_token.pop(token)
        scores[token] = self._compute_loss(model_without_token, word2count) - model_loss
    return scores

scores = self._compute_scores(model, word2count)

为了提升迭代效率，批量删除前10%的结果，即让整体loss增量最小的前10%的词。(删除这些词对整体loss的影响不大。)

1
2
3
4
sorted_scores = sorted(scores.items(), key=lambda x: x[1])
# Remove percent_to_remove tokens with the lowest scores.
for i in range(int(len(model) * 0.1)):
    _ = token2count.pop(sorted_scores[i][0])

获得新的词表后，重新计算每个词的概率，获得新的模型。并重复以上步骤，直到裁剪到词表大小符合要求。

1
2
3
4
5
6
7
8
while len(model) > vocab_size:
    scores = self._compute_scores(model, word2count)
    sorted_scores = sorted(scores.items(), key=lambda x: x[1])
    # Remove percent_to_remove tokens with the lowest scores.
    for i in range(int(len(model) * percent_to_remove)):
        _ = token2count.pop(sorted_scores[i][0])
    total_count = sum([freq for token, freq in token2count.items()])
    model = {token: -log(count / total_count) for token, count in token2count.items()}

假定预设的词表的大小为100，经过上述迭代后我们获得词表如下:

1
2
3
4
5
6
7
8
9
10
11
model = {
    '▁': 2.318585434340487, 
    'T': 4.653960350157523, 
    'h': 3.5553480614894135, 
    'i': 3.1876232813640963, 
    ...
    'seve': 5.752572638825633, 
    'sever': 5.752572638825633, 
    'severa': 5.752572638825633, 
    'several': 5.752572638825633
}

3.2. 推理阶段

在推理阶段，给定一个句子，需要将其切分成一个token的序列。 具体实现上先对句子进行预分词，然后对每个词基于Viterbi算法进行解码。

1
2
3
4
def tokenize(self, text):
    words = [word for word, _ in self.pre_tokenize_str(text)]
    encoded_words = [self._encode_word(word, self.model)[0] for word in words]
    return sum(encoded_words, [])

例如

1
2
>>> tokenize("This is the Hugging Face course!")
>>> ['▁This', '▁is', '▁the', '▁Hugging', '▁Face', '▁', 'c', 'ou', 'r', 's', 'e', '.']

基于Viterbi的切分获得的是最佳切分，基于unigram可以实现一个句子的多种切分方式，并且可以获得每种切分路径的打分。

4. SentencePiece

SentencePiece是Google出的一个分词工具，是一种基于 BPE 的分词工具，但它与 BPE 有所不同。它直接对原始文本进行处理，不需要预先进行空格分隔等预处理，并且可以生成子词单位。SentencePiece 将文本转换为unicode码点序列，然后对码点序列应用 BPE 算法，还可以对罕见码点进行 utf-8 编码转换:

  1. 文本预处理：将文本转换为unicode码点序列。
  2. BPE 训练：使用 BPE 算法对码点序列进行分词训练，生成子词单元。
  3. 罕见码点处理：对于低频码点，可以选择保留或进行 utf-8 编码转换。
  4. 词汇表生成：根据训练结果生成包含子词单元的词汇表。
  • 内置BPE，Unigram，char和word的分词方法
  • 无需预分词，以unicode方式直接编码整个句子，空格会被特殊编码为▁
  • 相比传统实现进行优化，分词速度速度更快

当前主流的大模型都是基于sentencepiece实现，例如ChatGLM的tokenizer。

1
2
3
4
5
6
7
8
9
10
11
12
13
...
class TextTokenizer:
    def __init__(self, model_path):
        self.sp = spm.SentencePieceProcessor()
        self.sp.Load(model_path)
        self.num_tokens = self.sp.vocab_size()

    def encode(self, text):
        return self.sp.EncodeAsIds(text)

    def decode(self, ids: List[int]):
        return self.sp.DecodeIds(ids)
...

三、对比分析

分词方法 特点 优势 应用场景
BPE 基于字节对频率合并，简单高效 平衡词汇表大小和文本粒度，处理罕见词效果好 GPT-2 等模型，多种自然语言处理任务
WordPiece 基于语言模型似然估计合并，考虑语义信息 更好地处理长尾词汇，提升模型泛化能力 BERT、多语言模型
SentencePiece 基于 BPE，直接处理原始文本，支持多种语言 处理无空格语言能力强，适用于多语言任务 中文、日语等语言处理，跨语言迁移任务
unigram 基于 unigram 语言模型概率，动态调整词汇表 提高语言模型准确性，适应文本统计特性 语言模型训练、语音识别等

引用

大模型基础组件 - Tokenizer