|
| 1 | +"""BPE tokenizer training (最简写法版本)。 |
| 2 | +
|
| 3 | +本文件刻意只用最基础的 Python 语法: 普通 for 循环、if 判断、list、dict。 |
| 4 | +不使用推导式、lambda、zip、enumerate、Counter 等简写工具, 方便逐行理解。 |
| 5 | +
|
| 6 | +核心数据结构: |
| 7 | +- words: list[list[bytes]] 每个元素是一个 pre-token(预分词单元), 内容是字节 token 的列表 |
| 8 | + 例如 "low" -> [b'l', b'o', b'w'] |
| 9 | +- counts: list[int] counts[i] 是 words[i] 在语料里出现的次数 |
| 10 | + words 和 counts 下标一一对应 |
| 11 | +- vocab: dict[int, bytes] 词表, token id -> 字节串 |
| 12 | +- merges: list[tuple] 按创建顺序记录的合并规则 |
| 13 | +""" |
| 14 | + |
| 15 | +import os |
| 16 | + |
| 17 | +import regex as re |
| 18 | + |
| 19 | +# GPT-2 的预分词正则(来自 assignment 文档)。 |
| 20 | +# 它把文本切成: 常见缩写('s 'll 've 're 't 'd 'm)、带前导空格的字母串、 |
| 21 | +# 带前导空格的数字串、带前导空格的符号串、以及各种空白。 |
| 22 | +PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""" |
| 23 | + |
| 24 | + |
| 25 | +def _pretokenize(text, special_tokens): |
| 26 | + """把一段文本预分词, 返回 dict: pre-token(字节 tuple) -> 出现次数。 |
| 27 | +
|
| 28 | + 特殊 token 会先把文本切开, 这样后续的合并不会跨越特殊 token 边界, |
| 29 | + 并且特殊 token 本身不参与统计(它们已经作为整体加入词表)。 |
| 30 | + """ |
| 31 | + counts = {} |
| 32 | + |
| 33 | + # ---- 用特殊 token 把文本切成若干段 ---- |
| 34 | + if special_tokens: |
| 35 | + # 按长度从大到小排序, 保证像 "<|eot|><|eot|>" 这种更长的 token 优先匹配。 |
| 36 | + sorted_specials = sorted(special_tokens, key=len, reverse=True) |
| 37 | + |
| 38 | + # 把每个特殊 token 转义后收进一个列表, 再用 | 连成 "a|b|c" 形式的正则。 |
| 39 | + escaped_specials = [] |
| 40 | + for tok in sorted_specials: |
| 41 | + escaped_specials.append(re.escape(tok)) |
| 42 | + split_pat = "|".join(escaped_specials) |
| 43 | + |
| 44 | + segments = re.split(split_pat, text) |
| 45 | + else: |
| 46 | + segments = [text] |
| 47 | + |
| 48 | + # ---- 对每一段(不含特殊 token)做 GPT-2 正则预分词 ---- |
| 49 | + for segment in segments: |
| 50 | + for match in re.finditer(PAT, segment): |
| 51 | + piece = match.group() |
| 52 | + |
| 53 | + # 把 pre-token 编码成 UTF-8 字节, 再拆成"单字节 token"的列表。 |
| 54 | + byte_list = [] |
| 55 | + for b in piece.encode("utf-8"): |
| 56 | + byte_list.append(bytes([b])) |
| 57 | + |
| 58 | + # 转成 tuple 才能当字典的键(list 不能当键, 因为它是可变的)。 |
| 59 | + token = tuple(byte_list) |
| 60 | + |
| 61 | + # 查字典累加: 没见过就从 0 开始。 |
| 62 | + if token in counts: |
| 63 | + counts[token] = counts[token] + 1 |
| 64 | + else: |
| 65 | + counts[token] = 1 |
| 66 | + |
| 67 | + return counts |
| 68 | + |
| 69 | + |
| 70 | +def _find_best_pair(pair_counts): |
| 71 | + """在 pair_counts 里找出最该合并的 pair。 |
| 72 | +
|
| 73 | + 规则: 频率最高; 频率相同时取字典序更大的 pair。 |
| 74 | + 如果没有任何有效 pair(频率都 <= 0), 返回 None。 |
| 75 | + """ |
| 76 | + best_pair = None |
| 77 | + best_count = 0 |
| 78 | + |
| 79 | + for pair in pair_counts: |
| 80 | + count = pair_counts[pair] |
| 81 | + |
| 82 | + # 频率 <= 0 的是"已经失效"的残留记录, 跳过。 |
| 83 | + if count <= 0: |
| 84 | + continue |
| 85 | + |
| 86 | + if best_pair is None: |
| 87 | + best_pair = pair |
| 88 | + best_count = count |
| 89 | + elif count > best_count: |
| 90 | + best_pair = pair |
| 91 | + best_count = count |
| 92 | + elif count == best_count and pair > best_pair: |
| 93 | + # bytes 之间的比较就是字典序, tuple 会逐个元素比较。 |
| 94 | + best_pair = pair |
| 95 | + best_count = count |
| 96 | + |
| 97 | + return best_pair |
| 98 | + |
| 99 | + |
| 100 | +def _remove_word_pairs(word, word_index, word_count, pair_counts, pair_to_words): |
| 101 | + """把一个 word 当前贡献的所有相邻 pair 从统计里扣掉。""" |
| 102 | + i = 0 |
| 103 | + while i < len(word) - 1: |
| 104 | + pair = (word[i], word[i + 1]) |
| 105 | + |
| 106 | + pair_counts[pair] = pair_counts[pair] - word_count |
| 107 | + |
| 108 | + if pair in pair_to_words: |
| 109 | + pair_to_words[pair].discard(word_index) |
| 110 | + |
| 111 | + i = i + 1 |
| 112 | + |
| 113 | + |
| 114 | +def _add_word_pairs(word, word_index, word_count, pair_counts, pair_to_words): |
| 115 | + """把一个 word 贡献的所有相邻 pair 加进统计。""" |
| 116 | + i = 0 |
| 117 | + while i < len(word) - 1: |
| 118 | + pair = (word[i], word[i + 1]) |
| 119 | + |
| 120 | + if pair in pair_counts: |
| 121 | + pair_counts[pair] = pair_counts[pair] + word_count |
| 122 | + else: |
| 123 | + pair_counts[pair] = word_count |
| 124 | + |
| 125 | + if pair not in pair_to_words: |
| 126 | + pair_to_words[pair] = set() |
| 127 | + pair_to_words[pair].add(word_index) |
| 128 | + |
| 129 | + i = i + 1 |
| 130 | + |
| 131 | + |
| 132 | +def _merge_in_word(word, pair, new_token): |
| 133 | + """在单个 word 内把所有出现的 pair 合并成 new_token, 返回新的 word 列表。""" |
| 134 | + new_word = [] |
| 135 | + i = 0 |
| 136 | + while i < len(word): |
| 137 | + # 当前位置和下一个正好构成要合并的 pair, 就合并成一个 token。 |
| 138 | + if i < len(word) - 1 and word[i] == pair[0] and word[i + 1] == pair[1]: |
| 139 | + new_word.append(new_token) |
| 140 | + i = i + 2 |
| 141 | + else: |
| 142 | + new_word.append(word[i]) |
| 143 | + i = i + 1 |
| 144 | + return new_word |
| 145 | + |
| 146 | + |
| 147 | +def train_bpe(input_path, vocab_size, special_tokens): |
| 148 | + """训练一个字节级 BPE tokenizer, 返回 (vocab, merges)。 |
| 149 | +
|
| 150 | + 性能关键: 每次合并后不重新扫描全部数据。我们额外维护 |
| 151 | + - pair_counts: pair -> 当前的加权总频率 |
| 152 | + - pair_to_words: pair -> 含有该 pair 的 word 下标集合 |
| 153 | + 这样每轮只需要更新"含有被合并 pair"的那几个 word。 |
| 154 | + """ |
| 155 | + # ---- 第一步: 初始化词表 ---- |
| 156 | + # 先放 256 个单字节, 再把特殊 token 追加在后面。 |
| 157 | + vocab = {} |
| 158 | + for i in range(256): |
| 159 | + vocab[i] = bytes([i]) |
| 160 | + for tok in special_tokens: |
| 161 | + vocab[len(vocab)] = tok.encode("utf-8") |
| 162 | + |
| 163 | + # ---- 第二步: 读文件, 预分词, 拆成两个并行列表 ---- |
| 164 | + with open(input_path, encoding="utf-8") as f: |
| 165 | + text = f.read() |
| 166 | + |
| 167 | + counter = _pretokenize(text, special_tokens) |
| 168 | + |
| 169 | + # words[i] 是可修改的 list, counts[i] 是它出现的次数。 |
| 170 | + words = [] |
| 171 | + counts = [] |
| 172 | + for word_tuple in counter: |
| 173 | + words.append(list(word_tuple)) |
| 174 | + counts.append(counter[word_tuple]) |
| 175 | + |
| 176 | + # ---- 第三步: 建立初始索引 ---- |
| 177 | + pair_counts = {} |
| 178 | + pair_to_words = {} |
| 179 | + i = 0 |
| 180 | + while i < len(words): |
| 181 | + _add_word_pairs(words[i], i, counts[i], pair_counts, pair_to_words) |
| 182 | + i = i + 1 |
| 183 | + |
| 184 | + # ---- 第四步: 反复合并, 直到词表达到目标大小 ---- |
| 185 | + merges = [] |
| 186 | + num_merges = vocab_size - len(vocab) |
| 187 | + |
| 188 | + merge_round = 0 |
| 189 | + while merge_round < num_merges: |
| 190 | + best_pair = _find_best_pair(pair_counts) |
| 191 | + if best_pair is None: |
| 192 | + break |
| 193 | + |
| 194 | + # 记录这次合并, 并把合并出的新 token 加入词表。 |
| 195 | + new_token = best_pair[0] + best_pair[1] |
| 196 | + merges.append(best_pair) |
| 197 | + vocab[len(vocab)] = new_token |
| 198 | + |
| 199 | + # 只处理含有 best_pair 的 word。 |
| 200 | + # 先拷成 list, 因为下面的循环会改动 pair_to_words 里的集合。 |
| 201 | + if best_pair in pair_to_words: |
| 202 | + affected = list(pair_to_words[best_pair]) |
| 203 | + else: |
| 204 | + affected = [] |
| 205 | + |
| 206 | + for word_index in affected: |
| 207 | + old_word = words[word_index] |
| 208 | + word_count = counts[word_index] |
| 209 | + |
| 210 | + # 1) 扣掉这个 word 现在贡献的所有 pair |
| 211 | + _remove_word_pairs(old_word, word_index, word_count, pair_counts, pair_to_words) |
| 212 | + |
| 213 | + # 2) 在这个 word 内合并 best_pair |
| 214 | + new_word = _merge_in_word(old_word, best_pair, new_token) |
| 215 | + words[word_index] = new_word |
| 216 | + |
| 217 | + # 3) 加回新 word 贡献的所有 pair |
| 218 | + _add_word_pairs(new_word, word_index, word_count, pair_counts, pair_to_words) |
| 219 | + |
| 220 | + # best_pair 已经在所有 word 里被消灭, 清理它的残留记录。 |
| 221 | + if best_pair in pair_counts: |
| 222 | + del pair_counts[best_pair] |
| 223 | + if best_pair in pair_to_words: |
| 224 | + del pair_to_words[best_pair] |
| 225 | + |
| 226 | + merge_round = merge_round + 1 |
| 227 | + |
| 228 | + return vocab, merges |
0 commit comments