Palind's Blog
0%

CS336 学习笔记（一）

发表于 更新于 分类于 ， Disqus：
本文字数： 6.8k 阅读时长 ≈ 6 分钟

Lecture 1

第一次课程包括实现 Byte-Pair Encoding (BPE) tokenizer。

tokenizer 不应该使用字节级别操作，因为会导致长 sequence。

如果 tokenizer 使用 word-based tokenization 会导致 vocabulary size 过大。

因此采用 BPE。核心逻辑是把经常一起出现的 byte 凑成一个新组合，并反复执行这个过程：

  1. 把每个字节（Byte）都看作一个独立的 Token。
  2. 看看哪两个 Token 挨在一起出现的次数最多。
  3. 把这对 Token 合并成一个新的 Token。
  4. 重复这个过程直到达到合适的词表大小。

基本代码实现如下。

  1. ABC 类 Tokenizer
1
2
3
4
5
6
7
8
9
10
11
from abc import ABC, abstractmethod

class Tokenizer(ABC):
    @abstractmethod
    def encode(self, string: str) -> list[int]:
        pass

    @abstractmethod
    def decode(self, indices: list[int]) -> str:
        pass

这是抽象基类。

  1. merge 函数
1
2
3
4
5
6
7
8
9
10
11
12
def merge(indices: list[int], pair: tuple[int, int], new_index: int) -> list[int]:
    """Return `indices`, but with all instances of `pair` replaced with `new_index`."""
    new_indices = []
    i = 0
    while i < len(indices):
        if i + 1 < len(indices) and indices[i] == pair[0] and indices[i + 1] == pair[1]:
            new_indices.append(new_index)
            i += 2
        else:
            new_indices.append(indices[i])
            i += 1
    return new_indices
  1. BPETokenizerParams 类
1
2
3
4
@dataclass(frozen=True)
class BPETokenizerParams:
    vocab: dict[int, bytes]
    merges: dict[tuple[int, int], int]

@dataclass 告知这个 BPETokenizerParams 类主要是用来存数据。(frozen=True) 使得创建参数对象后无法修改里面的内容。

成员 含义 例子
vocab 编号对应的原文 {256: b'th'} (256 代表 “th”)
merges 合并规则 {(116, 104): 256} (116(‘t’) 和 104(‘h’) 遇到就变 256)
  1. BPETokenizer 类
1
2
3
4
5
6
7
8
9
10
11
12
13
14
class BPETokenizer(Tokenizer):
    def __init__(self, params: BPETokenizerParams):
        self.params = params
    
    def encode(self, string: str) -> list[int]:
        indices = list(string.encode("utf-8"))  # 第一步：文字转字节
        for pair, new_index in self.params.merges.items(): # 第二步：按规则合并
            indices = merge(indices, pair, new_index)
        return indices

    def decode(self, indices: list[int]) -> str:
        bytes_list = list(map(self.params.vocab.get, indices))
        string = b"".join(bytes_list).decode("utf-8")
        return string

class BPETokenizer(Tokenizer) 表示 BPETokenizer 继承自 Tokenizer。def __init__ 初始化函数用 params 初始化 self.params。

encode 进行编码。string.encode("utf-8") 把字符串转换成基础的字节（0-255）。for 循环取出训练时存下的合并规则，按照训练时的顺序查看规则，并调用 merge 函数。比如第一条规则是 (116, 104) -> 256（即 t + h -> th）。就会调用 merge 函数，把当前列表里所有的 116, 104 组合全部换成 256。

map(self.params.vocab.get, indices) 对于 indices 里的每一个数字（比如 256），去 vocab 字典里查它对应什么字节（比如 b'th'）。

  1. train_bpe 函数
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
def train_bpe(string: str, num_merges: int) -> BPETokenizerParams:
    indices = list(map(int, string.encode("utf-8")))
    merges: dict[tuple[int, int], int] = {} 
    vocab: dict[int, bytes] = {x: bytes([x]) for x in range(256)}
    for i in range(num_merges):
        # Count the number of occurrences of each pair of tokens
        counts = defaultdict(int)
        for index1, index2 in zip(indices, indices[1:]):  # For each adjacent pair
            counts[(index1, index2)] += 1  # @inspect counts
        # Find the most common pair.
        pair = max(counts, key=counts.get)  # @inspect pair
        index1, index2 = pair
        # Merge that pair.
        new_index = 256 + i  # @inspect new_index
        merges[pair] = new_index  # @inspect merges
        vocab[new_index] = vocab[index1] + vocab[index2]  # @inspect vocab
        indices = merge(indices, pair, new_index)  # @inspect indices
    return BPETokenizerParams(vocab=vocab, merges=merges)

类型注解 string: str 告知 string 是字符串，返回类型为 BPETokenizerParams。

map 批量把刚才得到的每一个字节都明确转成 int（整数）类型。list(...) 把 map 的惰性结果变成列表。

merges = {}: 创建一个空字典，准备记录合并规则。

{x: bytes([x]) for x in range(256)} 字典推导式初始化一个词典，里面存好 0-255 每个数字对应的原始字节。

defaultdict(int) 使得如果去查一个没存进去的键，不会报错，而是自动给0。

indices[1:] 的切片操作取从第二个开始到最后。

zip(list1, list2) 把两个列表对齐。如 [A, B, C] 和 [B, C] ，变成 (A, B) 和 (B, C)。

max(counts, key=counts.get) 这是在字典里找“值最大”的那个键。也就是找出出现频率最高的那对搭档（比如 (104, 105)）。

new_index = 256 + i 新组合编号。

完整代码如下：

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
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
from abc import ABC, abstractmethod

class Tokenizer(ABC):
    @abstractmethod
    def encode(self, string: str) -> list[int]:
        pass
    
    @abstractmethod
    def decode(self, tokens: list[int]) -> str:
        pass
    
def merge(indices: list[int], pair: tuple[int, int], new_index: int) -> list[int]:
    """Return `indices` with all instances of `pair` replaced with new_index."""
    new_indices = []
    i = 0
    while i < len(indices):
        if i + 1 < len(indices) and indices[i] == pair[0] and indices[i + 1] == pair[1]:
            new_indices.append(new_index)
            i += 2
        else:
            new_indices.append(indices[i])
            i += 1
    return new_indices

from dataclasses import dataclass

@dataclass(frozen=True)
class BPETokenizerParams:
    vocab: dict[int, bytes]
    merges: dict[tuple[int, int], int]

class BPETokenizer(Tokenizer):
    def __init__(self, params: BPETokenizerParams):
        self.params = params
    
    def encode(self, string: str) -> list[int]:
        indices = list(string.encode("utf-8"))  # 第一步：文字转字节
        for pair, new_index in self.params.merges.items(): # 第二步：按规则合并
            indices = merge(indices, pair, new_index)
        return indices

    def decode(self, indices: list[int]) -> str:
        bytes_list = list(map(self.params.vocab.get, indices))
        string = b"".join(bytes_list).decode("utf-8")
        return string
    
from collections import defaultdict

def train_bpe(string: str, num_merges: int) -> BPETokenizerParams:
    indices = list(map(int, string.encode("utf-8")))
    merges: dict[tuple[int, int], int] = {} 
    vocab: dict[int, bytes] = {x: bytes([x]) for x in range(256)}
    for i in range(num_merges):
        # Count the number of occurrences of each pair of tokens
        counts = defaultdict(int)
        for index1, index2 in zip(indices, indices[1:]):  # For each adjacent pair
            counts[(index1, index2)] += 1  # @inspect counts
        # Find the most common pair.
        pair = max(counts, key=counts.get)  # @inspect pair
        index1, index2 = pair
        # Merge that pair.
        new_index = 256 + i  # @inspect new_index
        merges[pair] = new_index  # @inspect merges
        vocab[new_index] = vocab[index1] + vocab[index2]  # @inspect vocab
        indices = merge(indices, pair, new_index)  # @inspect indices
    return BPETokenizerParams(vocab=vocab, merges=merges)

if __name__ == "__main__":
    string = "the cat in the hat"  # @inspect string
    params = train_bpe(string, num_merges=3)
    tokenizer = BPETokenizer(params)
    string = "the quick brown fox"  # @inspect string
    indices = tokenizer.encode(string)  # @inspect indices
    reconstructed_string = tokenizer.decode(indices)  # @inspect reconstructed_string
    assert string == reconstructed_string
    print("Original string:", string)
    print("Encoded indices:", indices)
    print("Reconstructed string:", reconstructed_string)
    print("=============================")
    print("Vocabulary:")
    for index, byte_seq in params.vocab.items():
        print(f"  {index}: {byte_seq}")
    print("=============================")
    print("Merges:")
    for pair, new_index in params.merges.items():
        print(f"  {pair} -> {new_index}")

测试结果如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
lecture01  python .\lecture01.py
Original string: the quick brown fox
Encoded indices: [258, 113, 117, 105, 99, 107, 32, 98, 114, 111, 119, 110, 32, 102, 111, 120]
Reconstructed string: the quick brown fox
=============================
Vocabulary:
  0: b'\x00'
  1: b'\x01'
  2: b'\x02'
  [略]
  254: b'\xfe'
  255: b'\xff'
  256: b'th'
  257: b'the'
  258: b'the '
=============================
Merges:
  (116, 104) -> 256
  (256, 101) -> 257
  (257, 32) -> 258

可见分词器提取出了“the ”这个 token。

在 assignment 1 中提出了更高的要求，例如只循环有意义的合并，处理特殊 Token，使用预分词，性能优化等。准备之后完成。