tiktoken学习

2301_803652742025-05-31 23:01

1.tiktoken是OpenAI编写的进行高效分词操作的库文件。

2.操作过程：

复制代码

enc = tiktoken.get_encoding("gpt2")
train_ids = enc.encode_ordinary(train_data)
val_ids = enc.encode_ordinary(val_data)

以这段代码为例，get_encoding是创建了一个Encoding对象，结构如下：

复制代码

{
    "name": "gpt2",    #Encoding的名称
    "pat_str": r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""",  #分词正则表达式
    "mergeable_ranks": {b"!": 0, b"\"": 1, ...},  # 50,000+ 条目  #存储预加载的分词表
    "special_tokens": {"<|endoftext|>": 50256},   #特殊分词
    "explicit_n_vocab": 50257  #增加的特殊分词
}

encode_ordinary是利用BPE合并来对输入的train_data进行编码。

BPE合并：利用预加载的mergeable_ranks字典，通过最大前缀匹配查找最大字词映射对train_data编码。

复制代码

while current_byte in mergeable_ranks:
    find next byte that forms existing token
    merge if found in ranks

3.为什么说tiktoken高效？

使用高性能语言Rust实现

避免Python解释器开销；直接操作字节数组，避免Python对象的创建开销；并行处理。

基于Trie树的高效查找

复制代码

struct TrieNode {
    children: HashMap<u8, TrieNode>,
    token_id: Option<u32>,  // 匹配成功时返回 token ID
}

优化：Aho-Corasick 自动机，可以通过增加失败指针fail来避免每次失败从头遍历。就相当于这条路走不通，但是不会从头走，而是会走附近的分岔路看看有没有可以走的。

BPE合并的增量处理

复制代码

fn encode_bytes(bytes: &[u8], trie: &Trie) -> Vec<u32> {
    let mut tokens = Vec::new();
    let mut start = 0;
    
    while start < bytes.len() {
        let (end, token_id) = trie.longest_match(&bytes[start..]);
        tokens.push(token_id);
        start += end;
    }
    
    tokens
}

单次遍历：在扫描过程中同时完成匹配和合并

贪心最长匹配：总是选择可能的最长token

预加载mergeable_ranks

不需要实时建立，提高效率