从零手写 Byte-Level BPE 分词器：开启通向大语言模型的第一道门

===================================================================================
                   未加预分词 (No Pre-tokenization) 导致的灾难
===================================================================================
语料里有两句话:
  "I eat an apple. You eat an apple."

BPE 统计全局文本时发现: "apple" 后面经常跟着 "."
结果: BPE 强行把它们焊死在了一起，生成了一个新 Token: "apple." (带句号的苹果)

灾难后果:
  下次用户向模型提问: "Do you like apple?" (句末是问号)
  此时模型查遍词表，只认识带句号的 "apple."，根本无法复用 "apple" 的语义！
===================================================================================

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

class ByteBpeTokenizer:
    SPLIT_REGEX = SPLIT_REGEX

    def __init__(self):
        # 1. 必须准备的 4 个特殊标记 (Special Tokens)
        self.special_tokens = ["<pad>", "<s>", "</s>", "<unk>"]
        self.pad_token_id = 0  # 占位补齐符：用于短文本补齐
        self.bos_token_id = 1  # 句子开始符 (Beginning of Sequence)
        self.eos_token_id = 2  # 句子结束符 (End of Sequence)
        self.unk_token_id = 3  # 兜底未知符

        # 2. 映射字典
        self.vocab: Dict[int, bytes] = {}            # 数字 ID -> 原始字节串
        self.token_to_id: Dict[bytes, int] = {}      # 原始字节串 -> 数字 ID
        self.merges: Dict[Tuple[int, int], int] = {} # 记录合并规则: (id_a, id_b) -> 新合并ID

        # 3. 初始化地基
        self._init_base_vocab()

    def _init_base_vocab(self):
        """装入 4 个特殊标记和 256 个基础字节原子"""
        current_id = 0
        # 先存特殊标记 (占用 ID 0, 1, 2, 3)
        for st in self.special_tokens:
            b_st = st.encode("utf-8")
            self.vocab[current_id] = b_st
            self.token_to_id[b_st] = current_id
            current_id += 1

        # 紧接着存入 0 ~ 255 的基础单字节 (占用 ID 4 ~ 259)
        for b in range(256):
            b_byte = bytes([b])
            self.vocab[current_id] = b_byte
            self.token_to_id[b_byte] = current_id
            current_id += 1

    def train(self, texts: List[str], target_vocab_size: int = 4096, verbose: bool = True):
        # 第一步：预分词。把所有句子切成块，并转为基础字节的数字列表
        words = []
        for text in texts:
            for match in self.SPLIT_REGEX.finditer(text):
                chunk = match.group()
                # 将切块编码成字节，并查表转为初始 ID 列表
                token_ids = [self.token_to_id[bytes([b])] for b in chunk.encode("utf-8")]
                if token_ids:
                    words.append(token_ids)

        # 第二步：迭代合并。需要合并的次数 = 目标词表大小 - 当前已有大小
        num_merges = target_vocab_size - self.vocab_size
        for step in range(num_merges):
            pair_counts: Dict[Tuple[int, int], int] = {}
            # 扫描所有单词，统计相邻 pair 出现次数
            for word in words:
                for i in range(len(word) - 1):
                    pair = (word[i], word[i + 1])
                    pair_counts[pair] = pair_counts.get(pair, 0) + 1

            if not pair_counts:
                break # 没有可合并的了，提前结束

            # 找出出现次数最多的冠军 Pair
            best_pair = max(pair_counts, key=pair_counts.get)
            if pair_counts[best_pair] < 2:
                break # 如果最高频的 pair 才出现 1 次，合并毫无价值，停机

            # 产生新 Token 并注册到字典
            new_id = self.vocab_size
            self.merges[best_pair] = new_id
            new_bytes = self.vocab[best_pair[0]] + self.vocab[best_pair[1]]
            self.vocab[new_id] = new_bytes
            self.token_to_id[new_bytes] = new_id

            # 就地替换：在原数据集中把所有的 best_pair 替换为 new_id
            new_words = []
            for word in words:
                i = 0
                new_word = []
                while i < len(word):
                    if i < len(word) - 1 and (word[i], word[i + 1]) == best_pair:
                        new_word.append(new_id)
                        i += 2 # 跨越两位
                    else:
                        new_word.append(word[i])
                        i += 1
                new_words.append(new_word)
            words = new_words

    def encode(self, text: str, add_bos: bool = False, add_eos: bool = False) -> List[int]:
        result = []
        if add_bos:
            result.append(self.bos_token_id) # 句子开头加上 <s>

        for match in self.SPLIT_REGEX.finditer(text):
            chunk = match.group()
            # 1. 拆成单字节 ID
            chunk_ids = [self.token_to_id[bytes([b])] for b in chunk.encode("utf-8")]
            # 2. 按照学到的合并规则逐级合并
            merged_ids = self._merge_token_sequence(chunk_ids)
            result.extend(merged_ids)

        if add_eos:
            result.append(self.eos_token_id) # 句子末尾加上 </s>

        return result

    def decode(self, token_ids: List[int], skip_special_tokens: bool = True) -> str:
        byte_chunks = []
        for tid in token_ids:
            # 忽略掉 <s>, </s>, <pad> 等控制符号
            if skip_special_tokens and tid in [self.pad_token_id, self.bos_token_id, self.eos_token_id, self.unk_token_id]:
                continue
            if tid in self.vocab:
                byte_chunks.append(self.vocab[tid])

        # 核心：将所有散落的字节拼成完整的 bytes 数组，再反解出汉字与 Emoji
        all_bytes = b"".join(byte_chunks)
        return all_bytes.decode("utf-8", errors="replace")

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

# 标点吸附正则：允许汉字词块末尾吸收一个中文标点（，。！？；）
SPLIT_REGEX = re.compile(
    r"""'(?:[sdmt]|ll|ve|re)| ?\p{L}+(?:[，。！？；])?| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+"""
)

DEFAULT_SPECIAL_TOKENS = [
    "<pad>",         # ID 0: 批次填充占位符
    "<s>",           # ID 1: 文本起始符 (BOS)
    "</s>",          # ID 2: 文本结束符 (EOS)
    "<unk>",         # ID 3: 未知字符标记 (Byte-Level 下作为冗余占位)
    "<|title|>",     # ID 4: 诗词题目引导标记
    "<|author|>",    # ID 5: 作者姓名引导标记
    "<|dynasty|>",   # ID 6: 朝代归属引导标记
    "<|content|>"    # ID 7: 诗词正文起始标记
]

def _update_special_pattern(self):
    """按长度降序贪婪编译特殊标记正则"""
    if self.special_tokens:
        sorted_tokens = sorted(self.special_tokens, key=len, reverse=True)
        pattern_str = "(" + "|".join(re.escape(st) for st in sorted_tokens) + ")"
        self.special_token_pattern = re.compile(pattern_str)

if allowed_special and self.special_token_pattern:
    parts = self.special_token_pattern.split(text)
    for part in parts:
        if not part:
            continue
        # 命中特殊标记：直接转为对应单一原子 ID，严禁拆散！
        if part in self.special_tokens:
            result.append(self.token_to_id[part.encode("utf-8")])
        else:
            # 普通文本才送进 SPLIT_REGEX 走常规 BPE 切分
            self._encode_chunk(part, result)

python tests/test_tokenizer.py

初始词表大小（含特殊标记与基础 256 字节）: 264
[Tokenizer] 开始训练 BPE，初始词表大小: 264，目标词表大小: 320
[Tokenizer] 训练完成！最终词表大小: 276

--- 编解码结果 ---
原始文本: 从零手搓 Transformer 🚀 与 Tokenizer！
Token IDs: [1, 236, 195, 150, ..., 247, 196, 137, 2]
解码还原: 从零手搓 Transformer 🚀 与 Tokenizer！

--- 特殊 Token 原子提取与结构化元数据测试 ---
元数据序列: <|title|>登鹳雀楼<|author|>王之涣<|dynasty|>唐<|content|>白日依山尽，黄河入海流。</s>
元数据 Token IDs: [4, 239, 161, 195, 241, 193, 187, 241, 163, 136]...
跳过特殊标记还原: 登鹳雀楼王之涣唐白日依山尽，黄河入海流。

✅ Tokenizer 纯原生无损编解码与特殊 Token 原子保护校验 100% 通过！

tok = ByteBpeTokenizer.load("data/tokenizer.json")
ids = tok.encode("<|title|>静夜思<|author|>李白<|content|>床前明月光")
print("第 0 个 ID:", ids[0])  # 看看是否等于 4 (<|title|>)？