在软件开发领域，需求管理一直是项目成功的核心关键。随着项目复杂度提升和团队规模扩大，传统依赖文档、邮件和会议的需求管理方式显露出明显短板：版本混乱、协作困难、知识难以沉淀。更值得注意的是，行业内能够真正实现需求结构化、资产化，并结合AI技术进行智能化辅助的系统并不多见。我们公司是一家垂直领域专攻企业级需求与非企业级需求管理的公司， 我们公司的大模型应用连接：http://aipoc.chtech.cn:8880/#/login 欢迎试用。

下面我们讲解一下nanochat中的分词器和推理引擎：

"""
BPE Tokenizer in the style of GPT-4.
GPT-4风格的BPE分词器。

Two implementations are available:
有两种实现可用：
1) HuggingFace Tokenizer that can do both training and inference but is really confusing
   HuggingFace分词器，可以同时进行训练和推理，但真的很混乱
2) Our own RustBPE Tokenizer for training and tiktoken for efficient inference
   我们自己的RustBPE分词器用于训练，tiktoken用于高效推理
"""

import os
import copy
from functools import lru_cache

# 特殊token定义
SPECIAL_TOKENS = [
    # every document begins with the Beginning of Sequence (BOS) token that delimits documents
    # 每个文档都以序列开始(BOS)token开头，用于分隔文档
    "<|bos|>",
    # tokens below are only used during finetuning to render Conversations into token ids
    # 下面的token仅在微调期间使用，用于将对话渲染为token id
    "<|user_start|>", # 用户消息开始
    "<|user_end|>",   # 用户消息结束
    "<|assistant_start|>", # 助手消息开始
    "<|assistant_end|>",   # 助手消息结束
    "<|python_start|>",    # 助手调用Python REPL工具
    "<|python_end|>",      # Python代码结束
    "<|output_start|>",    # Python REPL输出返回给助手
    "<|output_end|>",      # 输出结束
]

# 注意：这个分割模式与GPT-4不同，我们使用\p{N}{1,2}而不是\p{N}{1,3}
# 我这样做是因为我不想为较小的词汇表大小在数字上"浪费"太多token
SPLIT_PATTERN = r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,2}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+"""

# -----------------------------------------------------------------------------
# 基于HuggingFace Tokenizer的通用GPT-4风格分词器
from tokenizers import Tokenizer as HFTokenizer
from tokenizers import pre_tokenizers, decoders, Regex
from tokenizers.models import BPE
from tokenizers.trainers import BpeTrainer

class HuggingFaceTokenizer:
    """HuggingFace分词器的轻量包装器，提供一些实用功能"""

    def __init__(self, tokenizer):
        self.tokenizer = tokenizer

    @classmethod
    def from_pretrained(cls, hf_path):
        # 从HuggingFace预训练分词器初始化（例如"gpt2"）
        tokenizer = HFTokenizer.from_pretrained(hf_path)
        return cls(tokenizer)

    @classmethod
    def from_directory(cls, tokenizer_dir):
        # 从磁盘上的本地目录初始化（例如"out/tokenizer"）
        tokenizer_path = os.path.join(tokenizer_dir, "tokenizer.json")
        tokenizer = HFTokenizer.from_file(tokenizer_path)
        return cls(tokenizer)

    @classmethod
    def train_from_iterator(cls, text_iterator, vocab_size):
        # 从文本迭代器训练
        # 配置HuggingFace分词器
        tokenizer = HFTokenizer(BPE(
            byte_fallback=True, # 需要！
            unk_token=None,
            fuse_unk=False,
        ))
        # 标准化器：无
        tokenizer.normalizer = None
        # 预分词器：GPT-4风格
        gpt4_split_regex = Regex(SPLIT_PATTERN) # huggingface要求你将其包装在Regex中！！
        tokenizer.pre_tokenizer = pre_tokenizers.Sequence([
            pre_tokenizers.Split(pattern=gpt4_split_regex, behavior="isolated", invert=False),
            pre_tokenizers.ByteLevel(add_prefix_space=False, use_regex=False)
        ])
        # 解码器：ByteLevel（与ByteLevel预分词器配对）
        tokenizer.decoder = decoders.ByteLevel()
        # 后处理器：无
        tokenizer.post_processor = None
        # 训练器：BPE
        trainer = BpeTrainer(
            vocab_size=vocab_size,
            show_progress=True,
            min_frequency=0, # 无最小频率
            initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
            special_tokens=SPECIAL_TOKENS,
        )
        # 开始训练
        tokenizer.train_from_iterator(text_iterator, trainer)
        return cls(tokenizer)

    def get_vocab_size(self):
        return self.tokenizer.get_vocab_size()

    def get_special_tokens(self):
        special_tokens_map = self.tokenizer.get_added_tokens_decoder()
        special_tokens = [w.content for w in special_tokens_map.values()]
        return special_tokens

    def id_to_token(self, id):
        return self.tokenizer.id_to_token(id)

    def _encode_one(self, text, prepend=None, append=None):
        # 编码单个字符串
        # prepend/append可以是特殊token的字符串或直接的token id
        assert isinstance(text, str)
        ids = []
        if prepend is not None:
            prepend_id = prepend if isinstance(prepend, int) else self.encode_special(prepend)
            ids.append(prepend_id)
        ids.extend(self.tokenizer.encode(text, add_special_tokens=False).ids)
        if append is not None:
            append_id = append if isinstance(append, int) else self.encode_special(append)
            ids.append(append_id)
        return ids

    def encode_special(self, text):
        # 通过精确匹配编码单个特殊token
        return self.tokenizer.token_to_id(text)

    def get_bos_token_id(self):
        bos = self.encode_special("<|bos|>")
        return bos

    def encode(self, text, *args, **kwargs):
        if isinstance(text, str):
            return self._encode_one(text, *args, **kwargs)
        elif isinstance(text, list):
            return [self._encode_one(t, *args, **kwargs) for t in text]
        else:
            raise ValueError(f"Invalid input type: {type(text)}")

    def __call__(self, *args, **kwargs):
        return self.encode(*args, **kwargs)

    def decode(self, ids):
        return self.tokenizer.decode(ids, skip_special_tokens=False)

    def save(self, tokenizer_dir):
        # 将分词器保存到磁盘
        os.makedirs(tokenizer_dir, exist_ok=True)
        tokenizer_path = os.path.join(tokenizer_dir, "tokenizer.json")
        self.tokenizer.save(tokenizer_path)
        print(f"Saved tokenizer to {tokenizer_path}")

# -----------------------------------------------------------------------------
# 基于rustbpe + tiktoken组合的分词器
import pickle
import rustbpe
import tiktoken

class RustBPETokenizer:
    """围绕tiktoken的轻量包装器（用于高效推理），但使用rustbpe进行训练"""

    def __init__(self, enc, bos_token):
        self.enc = enc
        self.bos_token_id = self.encode_special(bos_token)

    @classmethod
    def train_from_iterator(cls, text_iterator, vocab_size):
        # 1) 使用rustbpe训练
        tokenizer = rustbpe.Tokenizer()
        # 特殊token稍后在__init__中插入，我们不在这里训练它们
        vocab_size_no_special = vocab_size - len(SPECIAL_TOKENS)
        assert vocab_size_no_special >= 256, f"vocab_size_no_special must be at least 256, got {vocab_size_no_special}"
        tokenizer.train_from_iterator(text_iterator, vocab_size_no_special, pattern=SPLIT_PATTERN)
        
        # 2) 为推理构建相关的tiktoken编码
        pattern = tokenizer.get_pattern()
        mergeable_ranks_list = tokenizer.get_mergeable_ranks()
        mergeable_ranks = {bytes(k): v for k, v in mergeable_ranks_list}
        tokens_offset = len(mergeable_ranks)
        special_tokens = {name: tokens_offset + i for i, name in enumerate(SPECIAL_TOKENS)}
        enc = tiktoken.Encoding(
            name="rustbpe",
            pat_str=pattern,
            mergeable_ranks=mergeable_ranks, # dict[bytes, int] (token字节 -> 合并优先级排名)
            special_tokens=special_tokens,   # dict[str, int] (特殊token名称 -> token id)
        )
        return cls(enc, "<|bos|>")

    @classmethod
    def from_directory(cls, tokenizer_dir):
        pickle_path = os.path.join(tokenizer_dir, "tokenizer.pkl")
        with open(pickle_path, "rb") as f:
            enc = pickle.load(f)
        return cls(enc, "<|bos|>")

    @classmethod
    def from_pretrained(cls, tiktoken_name):
        # https://github.com/openai/tiktoken/blob/eedc8563/tiktoken_ext/openai_public.py
        enc = tiktoken.get_encoding(tiktoken_name)
        # tiktoken称特殊文档分隔符token为"<|endoftext|>"
        # 是的，这很令人困惑，因为这个token几乎总是在文档开头被PREPENDED（前置）
        # 它最常用于在推理期间向LLM发出新序列开始的信号等。
        # 所以在nanoChat中我们总是使用"<|bos|>"，缩写为"beginning of sequence"，但历史上它通常被称为"<|endoftext|>"。
        return cls(enc, "<|endoftext|>")

    def get_vocab_size(self):
        return self.enc.n_vocab

    def get_special_tokens(self):
        return self.enc.special_tokens_set

    def id_to_token(self, id):
        return self.enc.decode([id])

    @lru_cache(maxsize=32)  # 缓存最近使用的特殊token编码
    def encode_special(self, text):
        return self.enc.encode_single_token(text)

    def get_bos_token_id(self):
        return self.bos_token_id

    def encode(self, text, prepend=None, append=None, num_threads=8):
        # text可以是字符串或字符串列表

        if prepend is not None:
            prepend_id = prepend if isinstance(prepend, int) else self.encode_special(prepend)
        if append is not None:
            append_id = append if isinstance(append, int) else self.encode_special(append)

        if isinstance(text, str):
            ids = self.enc.encode_ordinary(text)  # 普通编码，不添加特殊token
            if prepend is not None:
                ids.insert(0, prepend_id) # TODO: 这里有点低效？嗯
            if append is not None:
                ids.append(append_id)
        elif isinstance(text, list):
            ids = self.enc.encode_ordinary_batch(text, num_threads=num_threads)  # 批量编码
            if prepend is not None:
                for ids_row in ids:
                    ids_row.insert(0, prepend_id) # TODO: 同样
            if append is not None:
                for ids_row in ids:
                    ids_row.append(append_id)
        else:
            raise ValueError(f"Invalid input type: {type(text)}")

        return ids

    def __call__(self, *args, **kwargs):
        return self.encode(*args, **kwargs)

    def decode(self, ids):
        return self.enc.decode(ids)

    def save(self, tokenizer_dir):
        # 将编码对象保存到磁盘
        os.makedirs(tokenizer_dir, exist_ok=True)
        pickle_path = os.path.join(tokenizer_dir, "tokenizer.pkl")
        with open(pickle_path, "wb") as f:
            pickle.dump(self.enc, f)
        print(f"Saved tokenizer encoding to {pickle_path}")

    def render_conversation(self, conversation, max_tokens=2048):
        """
        对单个Chat对话（我们在这里称之为"doc"或"document"）进行分词。
        Returns:
        - ids: list[int] 是这个渲染对话的token id列表
        - mask: list[int] 相同长度，mask = 1 表示助手应该训练的token。
        """
        # 我们将返回的ids、masks和一个辅助函数来帮助构建它们
        ids, mask = [], []
        def add_tokens(token_ids, mask_val):
            if isinstance(token_ids, int):
                token_ids = [token_ids]
            ids.extend(token_ids)
            mask.extend([mask_val] * len(token_ids))

        # 有时第一条消息是系统消息...
        # => 只需将其与第二条（用户）消息合并
        if conversation["messages"][0]["role"] == "system":
            # 这里目前需要进行一些对话手术...
            conversation = copy.deepcopy(conversation) # 避免改变原始数据
            messages = conversation["messages"]
            assert messages[1]["role"] == "user", "系统消息后必须是用户消息"
            messages[1]["content"] = messages[0]["content"] + "\n\n" + messages[1]["content"]
            messages = messages[1:]
        else:
            messages = conversation["messages"]
        assert len(messages) >= 1, f"对话消息少于1条: {messages}"

        # 获取我们需要的所有特殊token
        bos = self.get_bos_token_id()
        user_start, user_end = self.encode_special("<|user_start|>"), self.encode_special("<|user_end|>")
        assistant_start, assistant_end = self.encode_special("<|assistant_start|>"), self.encode_special("<|assistant_end|>")
        python_start, python_end = self.encode_special("<|python_start|>"), self.encode_special("<|python_end|>")
        output_start, output_end = self.encode_special("<|output_start|>"), self.encode_special("<|output_end|>")

        # 现在我们可以对对话进行分词
        add_tokens(bos, 0)  # BOS token，不监督
        for i, message in enumerate(messages):

            # 这里围绕假设进行一些健全性检查，以防止误用
            must_be_from = "user" if i % 2 == 0 else "assistant"
            assert message["role"] == must_be_from, f"消息 {i} 来自 {message['role']} 但应该来自 {must_be_from}"

            # content可以是简单字符串或部分列表（例如包含工具调用）
            content = message["content"]

            if message["role"] == "user":
                assert isinstance(content, str), "用户消息应该只是字符串"
                value_ids = self.encode(content)
                add_tokens(user_start, 0)    # 用户开始，不监督
                add_tokens(value_ids, 0)     # 用户内容，不监督
                add_tokens(user_end, 0)      # 用户结束，不监督
            elif message["role"] == "assistant":
                add_tokens(assistant_start, 0)  # 助手开始，不监督
                if isinstance(content, str):
                    # 简单字符串 => 简单地添加tokens
                    value_ids = self.encode(content)
                    add_tokens(value_ids, 1)     # 助手回复，监督
                elif isinstance(content, list):
                    for part in content:
                        value_ids = self.encode(part["text"])
                        if part["type"] == "text":
                            # 字符串部分 => 简单地添加tokens
                            add_tokens(value_ids, 1)  # 文本，监督
                        elif part["type"] == "python":
                            # Python工具调用 => 在<|python_start|>和<|python_end|>内添加tokens
                            add_tokens(python_start, 1)  # Python开始，监督
                            add_tokens(value_ids, 1)     # Python代码，监督
                            add_tokens(python_end, 1)    # Python结束，监督
                        elif part["type"] == "python_output":
                            # Python输出 => 在<|output_start|>和<|output_end|>内添加tokens
                            # 这些token都不受监督，因为token在测试时来自Python
                            add_tokens(output_start, 0)  # 输出开始，不监督
                            add_tokens(value_ids, 0)     # 输出内容，不监督
                            add_tokens(output_end, 0)    # 输出结束，不监督
                        else:
                            raise ValueError(f"未知的部分类型: {part['type']}")
                else:
                    raise ValueError(f"未知的内容类型: {type(content)}")
                add_tokens(assistant_end, 1)  # 助手结束，监督

        # 截断到最多max_tokens个token（有助于防止OOM）
        ids = ids[:max_tokens]
        mask = mask[:max_tokens]
        return ids, mask

    def visualize_tokenization(self, ids, mask):
        """调试中有用的小辅助函数：可视化render_conversation的分词"""
        RED = '\033[91m'
        GREEN = '\033[92m'
        RESET = '\033[0m'
        tokens = []
        for i, (token_id, mask_val) in enumerate(zip(ids, mask)):
            token_str = self.decode([token_id])
            color = GREEN if mask_val == 1 else RED
            tokens.append(f"{color}{token_str}{RESET}")
        return '|'.join(tokens)

    def render_for_completion(self, conversation):
        """
        在强化学习期间使用。在该设置中，我们想要
        渲染对话以为助手完成做准备。
        与Chat SFT情况不同，我们不需要返回mask。
        """
        # 我们需要做一些手术：我们需要弹出最后一条消息（助手的）
        conversation = copy.deepcopy(conversation) # 避免改变原始数据
        messages = conversation["messages"]
        assert messages[-1]["role"] == "assistant", "最后一条消息必须来自助手"
        messages.pop() # 原地删除最后一条消息（助手的）

        # 现在对对话进行分词
        ids, mask = self.render_conversation(conversation)

        # 最后，为助手完成做准备，附加助手开始token
        assistant_start = self.encode_special("<|assistant_start|>")
        ids.append(assistant_start)
        return ids

# -----------------------------------------------------------------------------
# nanochat特定的便利函数

def get_tokenizer():
    """获取分词器实例"""
    from nanochat.common import get_base_dir
    base_dir = get_base_dir()
    tokenizer_dir = os.path.join(base_dir, "tokenizer")
    # return HuggingFaceTokenizer.from_directory(tokenizer_dir)
    return RustBPETokenizer.from_directory(tokenizer_dir)

def get_token_bytes(device="cpu"):
    """获取每个token的字节数，用于计算bits per byte指标"""
    import torch
    from nanochat.common import get_base_dir
    base_dir = get_base_dir()
    tokenizer_dir = os.path.join(base_dir, "tokenizer")
    token_bytes_path = os.path.join(tokenizer_dir, "token_bytes.pt")
    assert os.path.exists(token_bytes_path), f"在 {token_bytes_path} 找不到token字节？它由tok_train.py写入"
    with open(token_bytes_path, "rb") as f:
        token_bytes = torch.load(f, map_location=device)
    return token_bytes

关键概念解释：

1. BPE分词原理:

   • 字节对编码: 从单个字节开始，逐步合并最常见的字节对

   • 词汇表: 包含所有可能的token，包括单个字节和合并的字节对

   • Byte Fallback: 遇到未知字符时回退到字节级别

2. 特殊Token的作用:

   • <|bos|>: 文档开始，用于分隔不同文档

   • 用户/助手标记: 标识对话中的不同角色

   • Python工具标记: 标识代码执行和输出

3. 对话渲染:

   • 监督学习: 只有助手的回复需要计算损失

   • 掩码机制: 用0/1标识哪些token需要监督

   • 工具调用: 支持Python代码执行和结果返回

4. 性能优化:

   • 多线程编码: 批量处理文本时使用多线程

   • 缓存: 对特殊token编码进行缓存

   • 流式处理: 支持大规模文本的流式处理

这个分词器系统支持完整的对话处理流程，包括工具调用和强化学习场景。