Reference
Setup
下载并安装环境
git clone https://github.com/stanford-cs336/assignment1-basics.git
cd assignment1-basics
uv self update # optional
uv sync # 如果没安装uv，就去看我的uv文章
source .venv/bin/activate
先测试一下，应该是全部fail
uv run pytest
下载数据
- TinyStories [Eldan and Li, 2023]
- OpenWebText [Gokaslan et al., 2019]
新建一个download_data.sh文件，把下面的代码复制进去
#!/bin/bash
# Create data directory and enter it
mkdir -p data
cd data
# Download TinyStories dataset
wget https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-train.txt
wget https://huggingface.co/datasets/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-valid.txt
# Download and extract OWT sample dataset
wget https://huggingface.co/datasets/stanford-cs336/owt-sample/resolve/main/owt_train.txt.gz
gunzip -f owt_train.txt.gz
wget https://huggingface.co/datasets/stanford-cs336/owt-sample/resolve/main/owt_valid.txt.gz
gunzip -f owt_valid.txt.gz
# Go back to previous directory
cd ..
然后
chmod +x download_data.sh # 改成可以执行的文件
./download_data.sh # 运行它
第二章 BPE分词器
2.1 Unicode standard
a. chr(0) returns '\x00' (其实代表的字符是null）
chr()会返回integer(unicode十进制编码）所对应的字符，比如chr(100)会返回d
ord('牛') 会返回字符的Unicode 编码。但用string.encode('utf-8')会返回bytes
ord() 和 .encode('utf-8')的区别
添加图片注释，不超过 140 字（可选）
ord会返回unicode编码，string.encode('utf-8')会返回utf8编码的三个bytes对应的数字。
bytes([raw byte value])会返回这个单独byte对应的字符。
bytes([0]) # []放的是raw byte value， 返回0所对应的十六进制bytes
注意bytes([255])和chr(255)返回的完全不一样。255在bytes里是ff, 在unicode里如果用utf-8编码，可以用两个bytes表示：
添加图片注释，不超过 140 字（可选）
建一个vocab dict
vocab = {i:bytes([i]) for i in range(256)}
添加图片注释，不超过 140 字（可选）
b. chr(0) is invisible when printed but its repr() shows it as the escape sequence '\x00'
Control character（控制字符） 是一种在 Unicode 或 ASCII 编码中用于控制文本的显示、格式或传输行为，而不是表示可见字符的特殊字符。
添加图片注释，不超过 140 字（可选）
c.
"this is a test" + a + "string" # return this is a test\x00string
print("this is a test" + chr(0) + "string") # return print("this is a test" + chr(0) + "string")
这个行为像是class里的repr和str
class MyClass:
def __repr__(self):
return "repr output"
def __str__(self):
return "str output"
b = MyClass()
b # return repr output
print(b) # return str output
2.2 Unicode encodings
a) UTF-8 is preferred for tokenizer training because it’s compact for common characters (like ASCII), avoids surrogate complexities of UTF-16, and allows modeling with a fixed 256-token byte vocabulary. It also aligns with web standards and real-world data encoding.
下面是一个不同UTF之间的比较：
添加图片注释，不超过 140 字（可选）
b.
decode_utf8_bytes_to_str_wrong("€".encode("utf-8"))
This function is incorrect because it decodes each byte individually rather than decoding the full multibyte UTF-8 sequence (e.g. '€' is encoded as three bytes: 0xE2 0x82 0xAC), which leads to a UnicodeDecodeError or incorrect characters.
c. Example: b'\xc3\x28' . This is an invalid UTF-8 sequence because 0xC3 indicates the start of a 2-byte character, but 0x28 is not a valid continuation byte (it doesn't start with 10xxxxxx).
连续的byte需要二进制以10为开头
添加图片注释，不超过 140 字（可选）
以10为开头的有8,9,A,B，这些是有效的连续字节
添加图片注释，不超过 140 字（可选）
Unicode
bit就是二进制里的最小单位，全称是 binary digit（二进制数字）。
一个byte是8个bit
8 bit （8个0或1的数） 可以表达256 个值，因为2^8=256
添加图片注释，不超过 140 字（可选）
每个字符在 Unicode 表里都有一个唯一编号（code point），例如：
'A' → U+0041 → 十进制 65
'你' → U+4F60 → 十进制 20320
但仅靠这个整数编号，计算机并不知道怎么把它“存成字节”或“传给别人”，还需要一套编码方式把它转换为具体的字节序列，比如 UTF-8、UTF-16、UTF-32。
相比string，bytes of letters占用的内存更少。
添加图片注释，不超过 140 字（可选）
十六进制
ord('你') # 20320
hex(ord('你')) # '0x4f60'
# 对应 Unicode 标准写法：U+4F60
十六进制是以 16 为基数的数制，用 0–9 和 A–F 来表示数字
添加图片注释，不超过 140 字（可选）
两个十六进制数字正好等于一个byte
十六进制格式，以 0x 开头
hex(10) # 输出 '0xA' 或 '0xa'
hex(65) # 输出 '0x41'
hex(20320) # 输出 '0x4f60'
0x 是 Python（和大多数编程语言）中表示“十六进制数”的前缀
U+是Unicode 标准的官方写法
\x 是 “十六进制转义符”，在 Python 字符串或字节串中用来表示一个字节（8 位），后面必须紧跟 两位十六进制数。格式：\xHH
b'\x41' # 表示 ASCII 字母 'A'，因为 0x41 = 65
b'\xe4' # 表示 UTF-8 编码中一个字节（汉字的一部分）
一些例子：
0x4F60
4F60 = 4×16³ + F×16² + 6×16¹ + 0×16⁰
= 4×4096 + 15×256 + 6×16 + 0
= 16384 + 3840 + 96 + 0
= 20320
0xFFFF
15×16e3+15×16e2+15×16e1+15=65535
0xFFFF 是 Basic Multilingual Plane（基本多文种平面，BMP）中的最后一个码点。 它常被视为普通 Unicode 字符的最大值。
十六进制转十进制
int('0x4f60', 16) # 输出 20320
int('BD', 16) # 输出 189
十六进制范围：
添加图片注释，不超过 140 字（可选）
UTF-8
UTF-8中的8指的是：
每个编码单元（code unit）是 8 位（bit），也就是 1 字节（byte）
添加图片注释，不超过 140 字（可选）
ASCII
128 characters: letters (A–Z, a–z), digits (0–9), punctuation, and control characters.
Each character is represented using 1 byte， 比如'A' → 0x41
前 128个 Unicode codepoints (U+0000 to U+007F) are identical to ASCII
所以 characters like 'h', 'e', 'l', 'o', etc., are 1 byte in both ASCII and UTF-8.
但对于'€' (U+20AC) needs 3 bytes in UTF-8: 0xE2 0x82 0xAC 如果每个单独decode会报错，需要一起。
UTF-8字节规则：
添加图片注释，不超过 140 字（可选）
2.5a Train BPE
a. BPE Tokenizer Training
要求
- vocab 从256 个byte value开始
- Pre-tokenization： 不用string.split(" "), 而用gpt2的pre-tokenizer，下面是一个执行的例子：
import regex as re
PAT = r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
# 会把每个词放在一个list里占内存
re.findall(PAT, "some text that i'll pre-tokenize")
# 用re.finditer
for m in re.finditer(PAT, text):
word = m.group(0)
用regex package （不是import re哦，regex在re的基础上多了更多功能）
- 当遇到多个字符对（token pair）并列为最高频率时，选“字典序更大的”那一对
max([("A", "B"), ("A", "C"), ("B", "ZZ"), ("BA", "A")])
-
special token
-
vocab里initialize的时候除了256个bytes，还需要加入特殊token，比如<|endoftext|>，然后有它们自己的token_id
-
Parallel pre-tokenization
-
时间的瓶颈就是pre-tokenization step, 可以用pretokenization_example.py的function来chunk text
-
pre-tokenization前把special token移除
-
re.split "|".join(special_token) , 小心用re.escape
-
用test_train_bpe_special_tokens来测试
-
优化merging这一步
-
每次merge，都要算每个pair的count，这其实会计算冗余，可以把pair count存在一个cache上，只更新那些和merged pair overlap的
-
用cProfile 或者scalene来找到时间瓶颈在哪里
-
厉害的同学可以考虑用C++ (cppyy) 或者 Rust (PyO3)来写
-
Function 里：
-
input args: input_path (text_path), vocab_size (int that defines max vocab size), special_tokens (list of strings)
-
return: vocab (dict {token_id, bytes}, merges (list of tuples of bytes)
答案
在cs336_basics下面建一个bpe.py
import regex as re
from collections import defaultdict
from tqdm.contrib.concurrent import process_map
PAT = re.compile(r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""")
def read_text(input_path):
with open(input_path, "r", encoding="utf-8") as f:
text = f.read()
return text
def split_by_special(text, special_tokens, drop_special=True):
if not special_tokens:
return [text]
# Sort by descending length to prioritize longer tokens (e.g., "<|endoftext|><|endoftext|>" before "<|endoftext|>")
special_tokens = sorted(special_tokens, key=len, reverse=True)
pattern = "|".join(re.escape(tok) for tok in special_tokens)
if not drop_special: pattern = f"({pattern})"
pattern = re.compile(pattern)
chunks = pattern.split(text)
return [c for c in chunks if c] # remove empty strings
def word2bytes(word):
"Convert word string to tuple of bytes"
a = list(word.encode('utf-8'))
return tuple(bytes([i]) for i in a)
def count_word(text):
"Split text into word bytes using GPT2 pattern and count word bytes frequency."
word_cnt = defaultdict(int)
for m in PAT.finditer(text):
word = m.group(0)
word_bytes = word2bytes(word)
if len(word_bytes)>=2:
word_cnt[word_bytes]+=1
return word_cnt
def merge_dicts(dicts):
merged = defaultdict(int)
for d in dicts:
for k, v in d.items():
merged[k] += v
return merged
def count_pair(word_cnt):
pair_cnt = defaultdict(int)
for word_bytes,cnt in word_cnt.items():
for pair in zip(word_bytes[:-1],word_bytes[1:]):
pair_cnt[pair]+=cnt
return pair_cnt
def get_max_pair(pair_cnt):
max_pair, _ = max(pair_cnt.items(), key=lambda x: (x[1], x[0])) # lexicographic tie-breaker
return max_pair
def get_basic_vocab(special_tokens):
vocab={token:bytes([token]) for token in range(256)}
for i,token in enumerate(special_tokens):
token_id = 256+i
vocab[token_id] = token.encode("utf-8")
return vocab
def apply_merge(word_bytes,merge):
merged = merge[0]+merge[1]
i = 0
new_word_bytes = []
while i < len(word_bytes):
# Check for match
if i < len(word_bytes) - 1 and word_bytes[i] == merge[0] and word_bytes[i+1] == merge[1]:
new_word_bytes.append(merged)
i += 2
else:
new_word_bytes.append(word_bytes[i])
i += 1
return tuple(new_word_bytes)
def update_cnt(word_cnt,pair_cnt, merge_pair):
new_word_cnt = defaultdict(int)
new_pair_cnt = defaultdict(int, pair_cnt) # copy with defaultdict
for word_bytes,cnt in word_cnt.items():
#----------for word cnt ---------------
old_pairs = list(zip(word_bytes[:-1], word_bytes[1:]))
# Keep the original count if the merge not appear in the key
if merge_pair not in old_pairs:
new_word_cnt[word_bytes]+=cnt
continue
# Use updated key if merge appear
new_word = apply_merge(word_bytes,merge_pair)
new_word_cnt[new_word]+=cnt
#--------for pair cnt ----------------
# Decrease all old pair counts
for pair in old_pairs:
new_pair_cnt[pair]-=cnt
if new_pair_cnt[pair] ==0:
del new_pair_cnt[pair]
# Count new pairs in the new word
new_pairs = list(zip(new_word[:-1], new_word[1:]))
for p in new_pairs:
new_pair_cnt[p] += cnt
return new_word_cnt,new_pair_cnt
def train_bpe(input_path,vocab_size,special_tokens):
text = read_text(input_path)
chunks = split_by_special(text,special_tokens)
# Only parallelize if chunk count is big enough
if len(chunks) < 4: word_dicts = list(map(count_word, chunks))
else: word_dicts = process_map(count_word, chunks, chunksize=1)
word_cnt = merge_dicts(word_dicts)
pair_cnt = count_pair(word_cnt)
vocab = get_basic_vocab(special_tokens)
base_vocab_size = len(vocab)
n_merges=vocab_size-base_vocab_size
merges = []
for i in range(n_merges):
max_pair = get_max_pair(pair_cnt)
vocab[base_vocab_size+i] = max_pair[0]+max_pair[1]
merges.append(max_pair)
word_cnt, pair_cnt = update_cnt(word_cnt,pair_cnt,max_pair)
return vocab, merges
adapters.py里
from cs336_basics.bpe import *
def run_train_bpe(
input_path: str | os.PathLike,
vocab_size: int,
special_tokens: list[str],
**kwargs,
) -> tuple[dict[int, bytes], list[tuple[bytes, bytes]]]:
vocab, merges = train_bpe(
input_path=input_path,
vocab_size=vocab_size,
special_tokens=special_tokens,
)
# vocab, merges = train_bpe(input_path, vocab_size, special_tokens)
return vocab, merges
测试
uv run pytest -k test_train_bpe_special_tokens
uv run pytest -k test_train_bpe
PASS
colab notebook:
2.5b TinyStories
要求
- 在tinystories里训练BPE tokenizer
Tiny Stories Overview
- 用<|endoftext|> 作为special token
- 时间应小于2分钟
- max vocab size=10,000
- 保存vocab和merges
代码
建一个run.py或者.ipynb
from cs336_basics.bpe import *
vocab, merges = train_bpe(
input_path='data/TinyStoriesV2-GPT4-valid.txt',
vocab_size=10_000,
special_tokens=["<|endoftext|>","<|endoftext|><|endoftext|>"],
)
整个过程持续1分50s。 另外version2 （heap+lazy update, 文章最下面的）也差不多这个时长 （所以这个优化不work）
保存vocab 和merges，另外load 的func。这里用了YAML，方便查看。
import yaml
def save_tokenizer_yaml(vocab, merges, fname):
"Save vocab and merges to a YAML file with UTF-8 decoding for readability."
# Convert bytes → string for readability
vocab_serializable = {
k: v.decode("utf-8", errors="replace") if isinstance(v, bytes) else v
for k, v in vocab.items()
}
merges_serializable = [
(a.decode("utf-8", errors="replace"), b.decode("utf-8", errors="replace"))
for a, b in merges
]
with open(fname, "w", encoding="utf-8") as f:
yaml.dump(
{"vocab": vocab_serializable, "merges": merges_serializable},
f,
allow_unicode=True,
sort_keys=False
)
def load_tokenizer_yaml(fname):
"Load vocab and merges from a YAML file, converting strings back to bytes."
with open(fname, "r", encoding="utf-8") as f:
data = yaml.safe_load(f)
vocab_loaded = {
int(k): v.encode("utf-8") if isinstance(v, str) else v
for k, v in data["vocab"].items()
}
merges_loaded = [
(a.encode("utf-8"), b.encode("utf-8")) for a, b in data["merges"]
]
return vocab_loaded, merges_loaded
save和load
save_tokenizer_yaml(vocab,merges,'tokenizer_owt_valid.yaml')
# to load
# vocab,merges=load_tokenizer_yaml('tokenizer_owt_valid.yaml')
时间和memory各是多少？
profile run.py，查看bottleneck
uv add scalene
uv run scalene run.py
打开profile.html
添加图片注释，不超过 140 字（可选）
memory
占用的总memory是左下的浅绿加深绿，198M加一个什么数。然后右上的max:114M指的是在某个时段内存占用的最大值。
time
点击时间column，按照时间排序
添加图片注释，不超过 140 字（可选）
哪个部分take了最长的时间？
时间take最长的是parallel processing chunk （GPT2 & count word）
这里可以把chunksize改到100可以更快（8.49s-->6.5s 在test train bpe里)
排名第二的是list(zip(word...))，这里可以优化的部分是可以用:
- pairwise from itertools代替list(zip(...))，速度稍快，没差太多
- 另外是先scan一边，用continue，pair_cnt因为是copy，只更新那些和merge_pair相关的count
from itertools import pairwise
def update_cnt_fast(word_cnt, pair_cnt, merge_pair):
a, b = merge_pair
new_word_cnt = defaultdict(int)
new_pair_cnt = defaultdict(int, pair_cnt) # copy
for wbytes, cnt in word_cnt.items():
# cheap presence check (no list/zip)
has = False
i, n = 0, len(wbytes) - 1
while i < n:
if wbytes[i] == a and wbytes[i+1] == b:
has = True
break
i += 1
if not has:
new_word_cnt[wbytes] += cnt
continue
# decrement old pairs (iterator, no list)
for p in pairwise(wbytes):
v = new_pair_cnt[p] - cnt
if v: new_pair_cnt[p] = v
else: new_pair_cnt.pop(p, None)
# merge & add new pairs
new_w = apply_merge(wbytes, merge_pair)
new_word_cnt[new_w] += cnt
for p in pairwise(new_w):
new_pair_cnt[p] += cnt
return new_word_cnt, new_pair_cnt
4.27s 在test train bpe里
57.5s 在tinystories里
还有一个方法是只update merge相关的：
Changed pairs per occurrence at index i (merging a,b -> m):
- Decrement: (L, a), (a, b), (b, R)
- Increment: (L, m), (m, R)
def find_merge_positions(wbytes, a, b):
# left-to-right, non-overlapping (typical BPE behavior)
pos = []
i, n = 0, len(wbytes)-1
while i < n:
if wbytes[i] == a and wbytes[i+1] == b:
pos.append(i)
i += 2
else:
i += 1
return pos
def apply_merge_inline(wbytes, a, b, new_tok):
out = []
i, n = 0, len(wbytes)
while i < n:
if i+1 < n and wbytes[i] == a and wbytes[i+1] == b:
out.append(new_tok)
i += 2
else:
out.append(wbytes[i])
i += 1
return tuple(out)
def update_cnt(word_cnt, pair_cnt, merge_pair):
a, b = merge_pair
m = a+b
new_word_cnt = defaultdict(int)
new_pair_cnt = defaultdict(int, pair_cnt)
for wbytes, cnt in word_cnt.items():
pos = find_merge_positions(wbytes, a, b)
if not pos:
new_word_cnt[wbytes] += cnt
continue
# For each merge site, touch only neighbors
for i in pos:
L = wbytes[i-1] if i-1 >= 0 else None
R = wbytes[i+2] if i+2 < len(wbytes) else None
# decrement old local pairs
if L is not None:
p = (L, a)
v = new_pair_cnt[p] - cnt
if v: new_pair_cnt[p] = v
else: new_pair_cnt.pop(p, None)
p = (a, b)
v = new_pair_cnt[p] - cnt
if v: new_pair_cnt[p] = v
else: new_pair_cnt.pop(p, None)
if R is not None:
p = (b, R)
v = new_pair_cnt[p] - cnt
if v: new_pair_cnt[p] = v
else: new_pair_cnt.pop(p, None)
# increment new local pairs
if L is not None:
new_pair_cnt[(L, m)] += cnt
if R is not None:
new_pair_cnt[(m, R)] += cnt
# build merged word once
new_w = apply_merge_inline(wbytes, a, b, m)
new_word_cnt[new_w] += cnt
return new_word_cnt, new_pair_cnt
4.44s 在test train bpe里 （上面那个似乎快一点）
1min 在tinystories里 （同上）
vocab里最长的token是哪个？
sorted(vocab.items(), key=lambda x: len(x[1]), reverse=True)
添加图片注释，不超过 140 字（可选）
accomplishment，比较make sense
2.5c OpenWebText
要求
- max vocab size: 32,000
- 保存vocab和merges
代码
这里我们就训练一下valid set （downscale一下，以学习为主）
vocab, merges = train_bpe(
input_path='data/owt_valid.txt',
vocab_size=32_000,
special_tokens=["<|endoftext|>","<|endoftext|><|endoftext|>"],
)
save_tokenizer_yaml(vocab,merges,'tokenizer_tinystories.yaml')
# to load
# vocab,merges=load_tokenizer_yaml('tokenizer_tinystories.yaml')
用了optimize的代码，训练时长：还在跑
最大vocab是？make sense吗?
和tinystories比如何？
2.6 BPE Tokenizer
要求
用class写，encoder，decoder
答案
from typing import Iterator, Iterable
import json
def split_to_words(text):
"Split text into words."
return PAT.findall(text)
def apply_merges(word_bytes, merges_set, vocab_to_id):
word_bytes = list(word_bytes)
while True:
min_token_id = float('inf')
best_pair_idx = -1
merged = None
for i in range(len(word_bytes) - 1):
pair = (word_bytes[i], word_bytes[i + 1])
if pair in merges_set:
combined = pair[0] + pair[1]
token_id = vocab_to_id.get(combined)
if token_id is not None and token_id < min_token_id:
min_token_id = token_id
best_pair_idx = i
merged = combined
if best_pair_idx == -1:
break
# Apply best merge
word_bytes = (
word_bytes[:best_pair_idx]
+ [merged]
+ word_bytes[best_pair_idx + 2:]
)
return tuple(word_bytes)
def encode_merged(text,merges,vocab_to_id):
word_list = split_to_words(text)
tokens=[]
for word in word_list:
word_bytes=word2bytes(word)
merged_word_bytes = apply_merges(word_bytes,merges,vocab_to_id)
tokens.extend(vocab_to_id[i] for i in merged_word_bytes)
return tokens
# Can merge the above functions into below class
class Tokenizer:
def __init__(self, vocab, merges, special_tokens=None):
self.vocab = vocab
self.merges = set(merges)
self.special_tokens = special_tokens if special_tokens else []
self.special_tokens_bytes = [i.encode('utf-8') for i in self.special_tokens]
self.vocab_to_id={v:k for k,v in vocab.items()}
# Ensure special tokens are in the vocabulary
for token_bytes in self.special_tokens_bytes:
if token_bytes not in self.vocab_to_id:
# Add to vocab if not already present
new_id = len(self.vocab)
self.vocab[new_id] = token_bytes
self.vocab_to_id[token_bytes] = new_id
@classmethod
def from_files(cls, vocab_filepath, merges_filepath, special_tokens=None):
# Load vocab (assumed to be a JSON file: {token_id: byte_string})
with open(vocab_filepath, 'r', encoding='utf-8') as vf:
vocab_data = json.load(vf)
# Optional: convert keys to int if stored as strings
vocab = {int(k): bytes(v, 'latin1') if isinstance(v, str) else bytes(v)
for k, v in vocab_data.items()}
# Load merges (assumed to be a list of pairs like: "a b")
with open(merges_filepath, 'r', encoding='utf-8') as mf:
lines = mf.readlines()
# Optional: skip headers like "#version: 0.2"
merge_pairs = [tuple(line.strip().split()) for line in lines if not line.startswith('#') and line.strip()]
# Convert to byte-pairs
merges = [(a.encode('utf-8'), b.encode('utf-8')) for a, b in merge_pairs]
return cls(vocab=vocab, merges=merges, special_tokens=special_tokens)
def encode(self, text: str) -> list[int]:
chunks = split_by_special(text, self.special_tokens, drop_special=False)
tokens = []
for chunk in chunks:
if self.special_tokens and chunk in self.special_tokens:
tokens.append(self.vocab_to_id[chunk.encode('utf-8')])
else:
tokens.extend(encode_merged(chunk, self.merges, self.vocab_to_id))
return tokens
def encode_iterable(self, iterable: Iterable[str]) -> Iterator[int]:
"""
Given an iterable of strings (e.g., a Python file handle), return a generator that lazily yields token IDs.
This is required for memory-efficient tokenization of large files that we cannot directly load into memory.
"""
for chunk in iterable:
yield from self.encode(chunk)
def decode(self, ids: list[int]) -> str:
"Decode a sequence of token IDs into text."
return b''.join([self.vocab[t] for t in ids]).decode('utf-8',errors='replace')
adapters.py里
def get_tokenizer(
vocab: dict[int, bytes],
merges: list[tuple[bytes, bytes]],
special_tokens: list[str] | None = None,
) -> Any:
return Tokenizer(vocab=vocab, merges=merges, special_tokens=special_tokens)
测试
uv run pytest tests/test_tokenizer.py
PASS
2.7 Experiments
2.5 optimize: Heap + lazy update
可以用heapq最大堆的方法来optimize，配合lazy update
import regex as re
import heapq
from collections import defaultdict
from tqdm.contrib.concurrent import process_map
PAT = re.compile(r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""")
def read_text(input_path):
with open(input_path, "r", encoding="utf-8") as f:
return f.read()
def split_by_special(text, special_tokens, drop_special=True):
if not special_tokens:
return [text]
special_tokens = sorted(special_tokens, key=len, reverse=True)
pattern = "|".join(re.escape(tok) for tok in special_tokens)
if not drop_special:
pattern = f"({pattern})"
return [c for c in re.compile(pattern).split(text) if c]
def word2bytes(word: str):
# tuple of 1-byte bytes objects; consistent with your original representation
return tuple(bytes([b]) for b in word.encode("utf-8"))
def count_word(text):
word_cnt = defaultdict(int)
for m in PAT.finditer(text):
w = m.group(0)
wb = word2bytes(w)
if len(wb) >= 2:
word_cnt[wb] += 1
return word_cnt
def merge_dicts(dicts):
merged = defaultdict(int)
for d in dicts:
for k, v in d.items():
merged[k] += v
return merged
def get_basic_vocab(special_tokens):
vocab = {i: bytes([i]) for i in range(256)}
for i, tok in enumerate(special_tokens):
vocab[256 + i] = tok.encode("utf-8")
return vocab
def apply_merge(word_bytes, merge):
merged = merge[0] + merge[1]
i = 0
out = []
L = len(word_bytes)
while i < L:
if i + 1 < L and word_bytes[i] == merge[0] and word_bytes[i+1] == merge[1]:
out.append(merged)
i += 2
else:
out.append(word_bytes[i])
i += 1
return tuple(out)
# ---------- HEAP + LAZY HELPERS ----------
def _desc_bytes_key(b: bytes):
# For descending lexicographic order on bytes: invert length & each byte.
return (-len(b), tuple(-x for x in b))
def _desc_pair_key(pair):
# pair is (bytes, bytes). Compare first, then second, both descending.
return (_desc_bytes_key(pair[0]), _desc_bytes_key(pair[1]))
def _heap_push(heap, pair, freq):
# Store (-freq, desc_key, pair). Min-heap => max by freq then lexicographically descending.
heapq.heappush(heap, (-freq, _desc_pair_key(pair), pair))
def _heap_pop_best(heap, pair_freq):
# Pop until top matches current freq (lazy invalidation).
while True:
negf, _, pair = heapq.heappop(heap)
f = -negf
if pair_freq.get(pair, 0) == f and f > 0:
return pair, f
# else stale; keep popping
# ---------- BUILD INITIAL INDICES ----------
def _build_word_tables(word_cnt):
"""
Turn {word_tuple -> count} into:
- words: list of word tuples (bytes tokens)
- wcnt: list of counts aligned with words
- wid_of: mapping for quick lookup if needed
"""
words, wcnt = [], []
for wb, c in word_cnt.items():
words.append(wb)
wcnt.append(c)
return words, wcnt
def _count_pairs_and_occ(words, wcnt):
pair_freq = defaultdict(int) # pair -> total weighted count
pair_occ = defaultdict(set) # pair -> set of word_ids containing it
for wid, (wb, c) in enumerate(zip(words, wcnt)):
if len(wb) < 2:
continue
pairs = list(zip(wb[:-1], wb[1:]))
for p in pairs:
pair_freq[p] += c
for p in set(pairs): # occ just needs membership
pair_occ[p].add(wid)
return pair_freq, pair_occ
# ---------- FAST TRAIN LOOP ----------
def train_bpe_fast(input_path, vocab_size, special_tokens=()):
text = read_text(input_path)
chunks = split_by_special(text, special_tokens)
# parallelize the first pass when appropriate
word_dicts = list(map(count_word, chunks)) if len(chunks) < 4 else process_map(count_word, chunks, chunksize=1)
word_cnt = merge_dicts(word_dicts)
# Tables
words, wcnt = _build_word_tables(word_cnt)
pair_freq, pair_occ = _count_pairs_and_occ(words, wcnt)
# Heap (lazy)
heap = []
for p, f in pair_freq.items():
_heap_push(heap, p, f)
# Vocab like yours
vocab = get_basic_vocab(special_tokens)
base_vocab_size = len(vocab)
n_merges = max(0, vocab_size - base_vocab_size)
merges = []
# Train
for step in range(n_merges):
if not heap:
break # no more pairs
best_pair, _ = _heap_pop_best(heap, pair_freq)
merges.append(best_pair)
# Register new token id
new_tid = base_vocab_size + step
vocab[new_tid] = best_pair[0] + best_pair[1]
affected = list(pair_occ.get(best_pair, ()))
if not affected:
# nothing to merge anymore; continue to next
continue
# For each affected word, update in-place
for wid in affected:
old = words[wid]
if len(old) < 2:
continue
c = wcnt[wid]
# ---- remove old pairs from global counts & occ ----
old_pairs = list(zip(old[:-1], old[1:]))
# freq decrements (weighted)
for q in old_pairs:
fq = pair_freq.get(q, 0)
if fq:
fq -= c
if fq <= 0:
pair_freq.pop(q, None)
else:
pair_freq[q] = fq
# push new (maybe smaller) freq; stale entries are lazily discarded later
if fq > 0:
_heap_push(heap, q, fq)
# occ removal (set-based, just membership)
for q in set(old_pairs):
occ = pair_occ.get(q)
if occ is not None:
occ.discard(wid)
if not occ:
pair_occ.pop(q, None)
# ---- apply merge on this word ----
new = apply_merge(old, best_pair)
words[wid] = new
# ---- add new pairs to global counts & occ ----
if len(new) >= 2:
new_pairs = list(zip(new[:-1], new[1:]))
for q in new_pairs:
fq = pair_freq.get(q, 0) + c
pair_freq[q] = fq
_heap_push(heap, q, fq)
for q in set(new_pairs):
pair_occ.setdefault(q, set()).add(wid)
# After processing, the merged pair won't exist anymore in those words,
# so its occ set is now either empty or only for non-updated words
if best_pair in pair_occ and not pair_occ[best_pair]:
pair_occ.pop(best_pair, None)
return vocab, merges
to be continued..