跳到主要内容

数据截至 (上游 commit d5827816baed)

04 · Python 绑定与生态对接

这一章讲什么： Rust 核心怎么变成 pip install tokenizers 之后那个 Tokenizer——PyO3 层的结构、GIL 的释放策略、批编码/异步编码/训练三条入口线，以及它作为 transformers「fast tokenizer」地基的生态位置。读完你会知道为什么 fast 比 slow 快一个数量级，以及 free-threaded Python 下这层绑定做了什么额外功课。

1. 它要解决的小问题​

Rust 核心写得再快，生态在 Python。绑定层要同时满足四个互相拉扯的要求：

  • 快：不能把时间花在跨语言拷贝上；批处理要真能并行（Python 侧 GIL 不放开，Rust 线程池就是摆设）。
  • 稳：Python 用户会 multiprocessing fork、会多线程共享 tokenizer、会随时 tokenizer.pre_tokenizer = ... 改配置——不能 race、不能死锁。
  • 全：Rust 侧每个 trait 的每个实现都要在 Python 侧可构造、可组合、可序列化。
  • 可训练：训练数据在 Python 侧（list、generator、numpy），喂数据时不能一直被 GIL 卡死。

2. 直觉：薄壳 + 大锁 + 放 GIL​

三句话概括绑定层的全部策略：

  1. 每个 Python 对象就是 Rust 对象的智能指针——PyTokenizer 内部是 Arc<RwLock<Tokenizer>>，方法调用 = 拿锁、调 Rust、转错误；
  2. 耗时操作先 py.detach（放开 GIL）再进 Rust——Rust 内部的 Rayon 并行才真正生效；
  3. 跨边界的数据早拷贝、一次成型——Python 字符串进 Rust 就转成 owned，Rust 结果出 Python 前才装箱。

3. 图示：一次 tokenizer.encode_batch 的旅程​

Python:  tokenizer.encode_batch(["a", "b", ...])
            │
            ▼
绑定层  bindings/python/src/tokenizer.rs:1312
   ① Vec<PyAny> 逐个 extract 成 tk::EncodeInput（持 GIL，纯拷贝）
            │
            ▼
   ② py.detach(|| { ... })          ← GIL 在此释放
            │
            ▼
Rust 核  encode_batch_char_offsets（tokenizers/src/tokenizer/mod.rs:1360）
   ③ maybe_par_iter: Rayon 多线程，每条输入独立走第 1 章的管线
   ④ pad_encodings 统一长度
            │
            ▼
   ⑤ Vec<Encoding> 逐个转 PyEncoding（重新持 GIL）
            │
            ▼
Python:  list[Encoding]

4. 原理演示（示意代码）​

# 示意，非源码：绑定方法的通用形态
def encode_batch(self, inputs):
    owned = [to_rust_encode_input(i) for i in inputs]   # 持 GIL：拷贝
    with gil_released():                                 # py.detach
        results = self.inner.encode_batch(owned)         # Rust + Rayon
    return [PyEncoding(r) for r in results]              # 持 GIL：装箱

5. 真实实现​

5.1 模块根部：gil_used = false 与 fork 保命​

#[pymodule(gil_used = false)]（bindings/python/src/lib.rs:50）声明本扩展在 free-threaded Python（无 GIL 构建）下安全使用——底气来自所有共享状态都在 Rust 侧的 Arc<RwLock> 里，而不是依赖 GIL 保护。

同文件 36-46 行注册了一个 pthread_atfork 回调：

// 摘自 bindings/python/src/lib.rs:39-46
extern "C" fn child_after_fork() {
    use tk::parallelism::*;
    if has_parallelism_been_used() && !is_parallelism_configured() {
        set_parallelism(false);
    }
}

原因写在注释里：Python 用户 fork 进程很常见，而 fork 后 Rayon 全局线程池的状态是坏的，子进程一用就死锁——所以子进程里直接把并行关掉。这是典型的「被用户 issue 教育过」的防御代码。

5.2 PyTokenizer：Arc<RwLock<Tokenizer>>​

PyTokenizer（bindings/python/src/tokenizer.rs:510-524）：

// 摘自 bindings/python/src/tokenizer.rs:517-524
pub struct PyTokenizer {
    /// Arc 让 clone 只是引用计数 +1（对齐 RwLock 之前共享语义）
    /// RwLock 让并发 setter/encoder 在 free-threaded Python 下不撞
    /// PyO3 的 per-pyclass 借用检查
    pub(crate) tokenizer: Arc<RwLock<Tokenizer>>,
}

读锁走 read_inner()（同文件 553-563 行），锁被毒化时转成 PyException 而不是 panic。所有 setter（换 normalizer、开 truncation）拿写锁，所有编码拿读锁——多只线程同时 encode 是并行的，改配置会互斥。

5.3 encode：坐标系是字符 offset​

Python 的 Tokenizer.encode（bindings/python/src/tokenizer.rs:1182-1213）做的是三件事：按 is_pretokenized 把 PyAny extract 成 TextInputSequence 或 PreTokenizedInputSequence；拼成 EncodeInput::Single/Dual；调 encode_char_offsets（同文件 1208-1211 行）。

这是与 Rust 默认 encode（字节 offset）的关键差异——Python 字符串按下标取的是字符，所以绑定层统一给字符坐标（换算机制见第 3 章 BytesToCharOffsetConverter）。

5.4 encode_batch 与 async：GIL 释放的三种形态​

入口释放方式并行机制位置
encode_batchpy.detach 包住整个调用Rust 侧 Rayon（into_maybe_par_iter）bindings/python/src/tokenizer.rs:1312-1336
async_encode先 owned 化输入，再 detachTokio spawn_blocking同文件 1244 起
async_encode_batch同上Tokio runtime + Rayon同文件 1364 起

异步版本共用模块级的全局 Tokio runtime TOKIO_RUNTIME（bindings/python/src/lib.rs:13-20，Lazy<Arc<Runtime>> 首次使用时建）。注意异步版在离开 GIL 前把输入完全 owned 化（build_owned_encode_inputs）——PyAny 不能带着 'py 生命周期进别的线程，这是 PyO3 异步的固定动作。

批编码的 Rust 侧（tokenizers/src/tokenizer/mod.rs:1337-1357）还有两个细节：padding 在并行编码之后统一做（保证 BatchLongest 拿到全批长度）；encode_batch_char_offsets（同文件 1360 行）是字符坐标版，Python 用的是它。

5.5 训练入口：PyBufferedIterator 的 256 缓冲​

train_from_iterator（bindings/python/src/tokenizer.rs:1784-1823）接受任意 Python 迭代器（list、generator、numpy 都行）。难点：训练在 py.detach 里跑，但取下一个元素必须持 GIL。解法是 PyBufferedIterator（bindings/python/src/utils/iterators.rs:44-75）：

  • 一次性持 GIL 把缓冲填满（默认 256 个元素，调用处传入）；
  • 填充时顺手用 converter 把 PyAny 转成 owned Rust 值（字符串或字符串列表）——注释明说动机：「转换成不携带 'py token 的值，GIL 才能放开」；
  • Rust 训练循环消费完一批再回来持 GIL refill。

于是 GIL 的持有粒度从「每元素一次」变成「每 256 元素一次」，训练期间 Python 主线程几乎不被卡。文件训练走另一条路：train（bindings/python/src/tokenizer.rs:1748-1763）直接在 detach 里调 train_from_files，全程不碰 Python 对象。PyTrainer 本身也是 Arc<RwLock<...>> 的转发壳（bindings/python/src/trainers.rs:91-110）。

5.6 序列化：一个 JSON 装下整条流水线​

Rust 侧 Serialize/Deserialize 为 TokenizerImpl 手写（tokenizers/src/tokenizer/serialization.rs:15、:50），五个部件经各自的 *Wrapper 枚举 tagged 序列化。于是：

  • tokenizer.save("tokenizer.json") → TokenizerImpl::save（tokenizers/src/tokenizer/mod.rs:1631-1640）；
  • Tokenizer.from_file → from_file（同文件 1566-1570 行，serde_json::from_str）；
  • Python 的 Tokenizer.from_pretrained（bindings/python/src/tokenizer.rs:789）从 Hub 拉下这个 JSON 再 from_file。

注意 Rust 侧的 Tokenizer::from_pretrained 已标记 #[deprecated(since = "0.14.0")]（tokenizers/src/tokenizer/mod.rs:1596-1606），官方建议改用 hf-hub crate 自己下载再 from_file——「联网」被刻意移出核心库职责。

PyTokenizer 也实现了 Serialize（bindings/python/src/tokenizer.rs:532-543），直接透传内部 Tokenizer——所以 json.dumps(tokenizer)、pickle 支持都白捡。

5.7 生态位：transformers 的 fast tokenizer 蹲在这上面​

本仓库自身不带 transformers 依赖（代码里只有 CHANGELOG 提到一次联动修复，bindings/python/CHANGELOG.md:27），但 tokenizer.json + 这层绑定就是 transformers PreTrainedTokenizerFast 的底座（inferred：生态常识，本克隆内无直接代码引用）：transformers 加载 fast tokenizer 时构造的就是这里的 tokenizers.Tokenizer，encode_plus/__call__ 的批处理最终落到 §5.4 的 encode_batch。这也是为什么「slow → fast」转换工具（bindings/python/scripts/convert.py）活在绑定层而不是核心层。

另外绑定层还附带一组 implementations/ 预设（BERT WordPiece、ByteLevel BPE、SentencePiece Unigram 等，bindings/python/py_src/tokenizers/implementations/），以 ByteLevelBPETokenizer（bindings/python/py_src/tokenizers/implementations/byte_level_bpe.py:10）为例，就是把第 1 章的四段管线按经典论文配方填好槽位的 Python 便捷类，基类 BaseTokenizer（bindings/python/py_src/tokenizers/implementations/base_tokenizer.py:13）持有内部 _tokenizer 并转发 padding/truncation 等配置。

6. 坑与注意点​

  • 改配置（setter）与编码并发会互相等锁——批量编码中途另一个线程 enable_truncation，RwLock 语义下安全但不免费；配置应在编码前定型。
  • encode 单条不走 Rayon。想要并行就用 encode_batch；在 asyncio 里用 async_encode_batch 而不是拿 encode 包 to_thread（绑定层已经把 Tokio 桥做好了）。
  • fork 后并行被静默关闭（§5.1）。multiprocessing 子进程里觉得变慢了，先怀疑 TOKENIZERS_PARALLELISM 与这个回调，而不是怀疑机器。
  • 字符 vs 字节坐标系：Python 全是字符 offset（§5.3），与 Rust encode 的默认行为不同；跨语言对齐结果时先统一坐标系（第 3 章 §6 已展开）。
  • 训练大语料优先用文件入口（train），迭代器入口的 256 缓冲对超大 epoch 仍有 GIL 往返开销。
  • from_pretrained（Rust）已废弃；新代码用 hf-hub 下载 + from_file。Python 侧接口仍在，但行为是「拉 tokenizer.json 再 from_file」。
  • free-threaded 支持是声明 + 锁，不是魔法：gil_used = false 成立的前提是所有共享态都在 Arc<RwLock> 后；给绑定层加新状态时沿用同一模式，否则会破坏这个承诺。

7. 代码地图​

主题文件路径符号名
模块根部bindings/python/src/lib.rs#[pymodule(gil_used = false)]、TOKIO_RUNTIME、child_after_fork
Tokenizer 壳bindings/python/src/tokenizer.rsPyTokenizer、read_inner、write_inner
编码入口bindings/python/src/tokenizer.rsencode、encode_batch、async_encode、async_encode_batch
训练入口bindings/python/src/tokenizer.rs、bindings/python/src/trainers.rstrain、train_from_iterator、PyTrainer
迭代器缓冲bindings/python/src/utils/iterators.rsPyBufferedIterator、refill
Encoding 壳bindings/python/src/encoding.rsPyEncoding
预设配方bindings/python/py_src/tokenizers/implementations/ByteLevelBPETokenizer、BertWordPieceTokenizer、BaseTokenizer
核心序列化tokenizers/src/tokenizer/serialization.rs、tokenizers/src/tokenizer/mod.rsimpl Serialize/Deserialize for TokenizerImpl、from_file、save
并行开关tokenizers/src/utils/parallelism.rsmaybe_par_iter、set_parallelism、TOKENIZERS_PARALLELISM