前言

开始学习cs336，目前已经做完了作业A1，但算力有限，所以最后超参数调优实验没做，目前来总结一下整个实验的知识和内容。

分词器

对于语言，人类可以自然的理解为连续的内容，而在计算机中，一般是按照离散的字节进行存储的。对于一句话“We are eating bread"，计算机会把这句话进行拆分，例如最简单的拆分方法便是根据空格进行拆分，即[“we”, “are”, “eating”, “bread”]，然后对应的得到token2id和id2token。即{0: “we”, 1:“are”, 2: “eating”, 3: “bread”}和{“we”: 0, “are”: 1, “eating”: 2, “bread”: 3}。这样在训练的时候，我们只需要获取对应token的id，在最后生成文本的时候，只需要根据id就可以找到对应token

BPE分词器

作业一开始便让实现一个分词器，其实现思路是：

把任意语言转换为通用处理的字符，例如unicode
对unicode字符进行编码，例如utf-8
进行训练

我们来拆开讲一讲。

unicode编码

unicode编码是表示字符的通用方法，具体来说，unicode会给每个字符分配一个数字，该数字就是unicode的编码后的结果，比如对于"牛"这个字来说，请看以下示例程序

ord: ord(character, /) Return the ordinal value of a character.
If the argument is a one-character string, return the Unicode code
point of that character.

If the argument is a bytes or bytearray object of length 1, return its
single byte value.
chr: chr(i, /) Return a Unicode string of one character with ordinal i; 0 <= i <= 0x10ffff.

1
2
3
4


>>>>> ord('牛')
29275
>>> chr(29275)
'牛'

这里，我们可以想，是否可以用unicode来作为词汇表呢？答案是不行，首先，unicode是一个超大的集合，而且其中有许多字符都很罕见。所以，一个合适的词汇表是内容数量适中的。我们正好可以利用unicode的三种编码方式：UTF-8，UTF-16和UTF-32，主流的编码方式是UTF-8。

UTF-8编码

UTF-8编码本质上是去把一个unicode编码后的数字转为一个字节或多个字节序列，每个字节取值范围都是0~255，可以很方便的使用encode和decode去进行编码和解码

encode示例

1
2
3
4


>>> "hello".encode("utf-8")
b'hello'
>>>> "牛子".encode("utf-8")
b'\xe7\x89\x9b\xe5\xad\x90'

decode示例

1
2
3
4


>>> b'hello'.decode()
'hello'
>>> b'\xe7\x89\x9b\xe5\xad\x90'.decode()
'牛子'

注意，对于中文示例，虽然整个字符串只有两个字，但是其utf-8编码长度却为6，所以我们不可以简单的遍历每个字节去进行解码，而是要去整体解码。

预分词

文档中写不直接在text上进行BPE合并的原因是因为，对于一些词来说，例如dog!和dog.，这两个词的含义不应该受到结尾标点的影响而被看作是两个不同的词，所以我们要进行预分词处理。在早期的预分词处理就是去根据空格把text切分，而文档里面给出了一个正则pat来进行切分：

1

PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""

文档也给出了使用例子：

1
2
3


>>> import regex as re
>>> re.findall(PAT, "some text that i'll pre-tokenize")
['some', ' text', ' that', ' i', "'ll", ' pre', '-', 'tokenize']

其中文档也说了，在处理大型数据的时候，最好使用迭代器的方式避免一次加载进内存中，可以使用re.finditer来替换

re.finditer返回一个迭代器matches，对于这个matches，我们可以进行遍历来取出里面的内容
1
2
3
4
5
matches = re.finditer(PAT, text)

for match in matches:
   print(match.group())
match.group() 或 match.group(0)	整个匹配内容

训练BPE

BPE是一个压缩合并算法，具体来说就是不断去合并字符对，这里有一个训练示例：

1
2
3


low low low low low
lower lower widest widest widest
newest newest newest newest newest newest

我们来看看怎么合并：

统计词频

1

{low: 5, lower: 2, widest: 3, newest: 6}

统计相邻字符个数

当前词汇是：

1
2
3
4
5
6


{
    "low": ('l', 'o', 'w'),
    "lower": ('l', 'o', 'w', 'e', 'r'),
    "widest": ('w', 'i', 'd', 'e', 's', 't'),
    "newest": ('n', 'e', 'w', 'e', 's', 't')
}

得到他们的字节对：

1
2
3
4
5
6


{
    low: (lo, ow), 5
    lower: (lo, ow, we, er), 2
    widest: (wi, id, de, es, st), 3
    newest: (ne, ew, we, es, st), 6
}

统计个数后我们得到

1

{lo: 7, ow: 7, we: 8, er: 2, wi: 3, id: 3, de: 3, es: 9, st: 9, ne: 6, ew: 6}

合并出现频率最高的

可以看到，在上面那个例子中，出现频率最高的是
1

('e', 's')和('s', 't')
这两个字节对，在合并时，如果遇到了出现频率一样，那么优先合并字典序大的那个元素，所以，现在的待合并词变为了
1 2 3 4 5 6

{ "low": ('l', 'o', 'w'), "lower": ('l', 'o', 'w', 'e', 'r'), "widest": ('w', 'i', 'd', 'e', 'st'), "newest": ('n', 'e', 'w', 'e', 'st'), }
然后一直合并下去。

在合并的过程中，我们还维护一个merges代表先后合并顺序，即每次要合并的pair添加到这个merges中，方便编码时使用。

BPE编码与解码

编码

经过上面的训练后，我们就得到了一个完整的词汇表，即一个id和token相互对应的字典。我们来看看给定一句话，怎么把这句话转换为token_id

例如， “the cat ate"这句话经过预分词后变为["the", " cat", " ate"]，注意其中的空格是存在的。然后进行拆分并获取字节对。假设我们的词汇表是

1

{0: b' ', 1: b'a', 2:b'c', 3: b'e', 4: b'h', 5: b't', 6: b'th', 7: b' c', 8: b' a', 9: b'the', 10: b' at'}

合并的先后顺序是：

1

[(b't', b'h'), (b' ', b'c'), (b' ', b'a'), (b'th', b'e'), (b' a',b't')]

拆分

1
2
3
4
5


{
    "the": ['t', 'h', 'e'],
    " cat": [' c', 'a', 't'],
    " ate": [' a', 't', 'e'],
}

查看合并对中是否存在合并

对于第一个单词”the”，其字节对是[('t', 'h'), ('h', 'e')]，正好merges中第一个就是th，合并并替换掉，现在待合并变为
1 2 3 4 5

{ "the": ['th', 'e'], " cat": [' c', 'a', 't'], " ate": [' a', 't', 'e'], }
对于第一个单词”the"，其字节对是[('th'), ('e')]，正好在merges中，合并并替换掉，现在待合并变为
1 2 3 4 5

{ "the": ['the'], " cat": [' c', 'a', 't'], " ate": [' a', 't', 'e'], }
下面两个词的合并过程也是一样
查表得到token id

得到这句话的id集合[9, 7, 1, 5, 10, 3]

解码

可以看到，对于这一句话，我们想要解码的话，需要去词汇表中拿到对于utf-8编码后的字节。注意，我们不能直接去根据单个id去进行解码，因为一个词可能会被拆分为多个id，我们需要将这一一组id全部传进去，然后去使用decode函数。

例如这个实现：

1
2
3


def decode(self, ids: list[int]) -> str:
    ans = b"".join(self.vocab[idx] for idx in ids)
    return ans.decode('utf-8', errors= 'replace')

实验杂项和细节

把理论部分说完了，现在来看看实验中的一些实现细节。

文档读取

实验给了两个数据集：

Tinystory
openwebtext

前者数据规模小一些，文档中给了一个根据special_token进行切分的示例代码

 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


import os
from typing import BinaryIO


def find_chunk_boundaries(
    file: BinaryIO,
    desired_num_chunks: int,
    split_special_token: bytes,
) -> list[int]:
    """
    Chunk the file into parts that can be counted independently.
    May return fewer chunks if the boundaries end up overlapping.
    """
    assert isinstance(split_special_token, bytes), "Must represent special token as a bytestring"

    # Get total file size in bytes
    file.seek(0, os.SEEK_END)
    file_size = file.tell()
    file.seek(0)

    chunk_size = file_size // desired_num_chunks

    # Initial guesses for chunk boundary locations, uniformly spaced
    # Chunks start on previous index, don't include last index
    chunk_boundaries = [i * chunk_size for i in range(desired_num_chunks + 1)]
    chunk_boundaries[-1] = file_size

    mini_chunk_size = 4096  # Read ahead by 4k bytes at a time

    for bi in range(1, len(chunk_boundaries) - 1):
        initial_position = chunk_boundaries[bi]
        file.seek(initial_position)  # Start at boundary guess
        while True:
            mini_chunk = file.read(mini_chunk_size)  # Read a mini chunk

            # If EOF, this boundary should be at the end of the file
            if mini_chunk == b"":
                chunk_boundaries[bi] = file_size
                break

            # Find the special token in the mini chunk
            found_at = mini_chunk.find(split_special_token)
            if found_at != -1:
                chunk_boundaries[bi] = initial_position + found_at
                break
            initial_position += mini_chunk_size

    # Make sure all boundaries are unique, but might be fewer than desired_num_chunks
    return sorted(set(chunk_boundaries))


## Usage
with open(..., "rb") as f:
    num_processes = 4
    boundaries = find_chunk_boundaries(f, num_processes, b"<|endoftext|>")

    # The following is a serial implementation, but you can parallelize this
    # by sending each start/end pair to a set of processes.
    for start, end in zip(boundaries[:-1], boundaries[1:]):
        f.seek(start)
        chunk = f.read(end - start).decode("utf-8", errors="ignore")
        # Run pre-tokenization on your chunk and store the counts for each pre-token

我觉得比较值得学习的地方在于，我们可以利用多进程去读写训练数据。我们把数据集分为多个chunk，每个进程负责去维护一个词频表，最后把每个chunk的词频表进行合并。

 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


# 使用多进程去进行统计
    with open(input_path, 'rb') as f:
        num_processes = 8
        boundaries = find_chunk_boundaries(f, num_processes, b"<|endoftext|>")
        ...
        # 对每个boundary都去使用一个进程去统计
        # 创建一个线程池
    jobs = [
        (
            input_path,
            start, 
            end, 
            special_tokens,
            
        ) for start, end in zip(boundaries[:-1], boundaries[1:])
    ]    
    with mp.Pool(processes= num_processes) as pool:
        part_cnts = pool.starmap(count_chunk, jobs)
        
    
    # 得到词频统计
    cnt = Counter()
    for part_cnt in part_cnts:
        cnt.update(part_cnt)
    
    # 初始化split_token
    split_token = {
        pre_token: [bytes([value]) for value in pre_token]
        for pre_token in cnt.keys()
    }

预处理

在读取完文档后，我们需要去得到一个词汇表。也就是我们之前实现的train_bpe部分

 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
87


def train_bpe(input_path: str, vocab_size: int, special_tokens: list[str])-> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
    """训练BPE分词器，得到最终词汇表和合并过程
    """
    
    # 边界处理
    assert Path(input_path).exists(), "No such file"
    assert vocab_size >= 256, 'Not correct vocab_size'
    
    # 初始化词汇表
    vocab = {
        i: bytes([i]) for i in range(256)
    }
    
    # 添加特殊词汇
    for special_token in special_tokens:
        # 首先对特殊词汇进行 encoder
        special_token_encode = special_token.encode('utf-8')
        if special_token_encode not in vocab.values():
            # 执行添加操作
            vocab[len(vocab)] = special_token_encode
    
    assert len(vocab) <= vocab_size, '词汇表大小超出规定'
    
    # 使用多进程去进行统计
    with open(input_path, 'rb') as f:
        num_processes = 8
        boundaries = find_chunk_boundaries(f, num_processes, b"<|endoftext|>")
        ...
        # 对每个boundary都去使用一个进程去统计
        # 创建一个线程池
    jobs = [
        (
            input_path,
            start, 
            end, 
            special_tokens,
            
        ) for start, end in zip(boundaries[:-1], boundaries[1:])
    ]    
    with mp.Pool(processes= num_processes) as pool:
        part_cnts = pool.starmap(count_chunk, jobs)
        
    
    # 得到词频统计
    cnt = Counter()
    for part_cnt in part_cnts:
        cnt.update(part_cnt)
    
    # 初始化split_token
    split_token = {
        pre_token: [bytes([value]) for value in pre_token]
        for pre_token in cnt.keys()
    }
    
    # 构建索引
    token_index = build_index(split_token)
    # 1. 对相邻pair进行计数
    cnt_pair = count_pairs(
        split_token,
        cnt
    )
    # 合并结果
    merges = []
    
    while (len(vocab) < vocab_size):

        
        # 2. 找到出现频率最高的pair
        best_pair = max(
            cnt_pair,
            key= lambda pair: (
                cnt_pair[pair],
                pair
            )
        )
        
        # 3. 创建并记录新的pair
        left, right = best_pair
        new_token = left + right
        vocab[len(vocab)] = new_token
        merges.append(best_pair)
        
        
        # 5. 对受影响的token进行合并操作
        merge_and_update(split_token, cnt, cnt_pair, token_index, best_pair)
        
    return vocab, merges

为了后面利用，我们可以把vocab，merges和special_token进行序列化，存储为一个.pkl文件

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15


import picket
tokenizer_state = {
    "vocab": vocab,
    "merges": merges,
    "special_tokens": special_tokens,
}
# 写入数据
with open("data/tokenizer.pkl", "wb") as file:
    pickle.dump(tokenizer_state, file)

# 读取数据
import pickle

with open('file.pkl', 'rb') as f: # 注意是'rb' model
    data = pickle.load(f)

保存数据

为了方便训练，文档中给了一个使用np.memmap的方式去读取避免内存占用过多

 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


def encode_file(
    tokenizer: Tokenizer,
    input_path: str,
    output_path: str,
    chunk_size: int = 1_000_000,
):
    max_token_id = max(tokenizer.vocab)

    if max_token_id <= np.iinfo(np.uint16).max:
        dtype = np.uint16
    else:
        dtype = np.uint32

    Path(output_path).parent.mkdir(
        parents=True,
        exist_ok=True,
    )

    with (
        open(input_path, "r", encoding="utf-8") as source,
        open(output_path, "wb") as destination,
    ):
        token_iterator = tokenizer.encode_iterable(source)

        while True:
            chunk = np.fromiter(
                islice(token_iterator, chunk_size),
                dtype=dtype,
            )

            if chunk.size == 0:
                break

            chunk.tofile(destination)

    print(
        f"Saved {output_path}, "
        f"tokens={Path(output_path).stat().st_size // np.dtype(dtype).itemsize:,}, "
        f"dtype={dtype.__name__}"
    )

我们可以把训练数据保存为.bin二进制数据，从而直接使用去进行读写。

到此，我们就有了：

训练数据
词汇表
编码器

下面就可以愉快的进行模型结构和训练部分了！