前言
AI Agent 开发中,上下文的 Token 计数是不得不品的一环。
从文本中统计 Token 数并不简单,具体方案大致可以分为后端 API 请求和 Agent 本地计算。
Anthropic API 提供了统计 Token 数的专用接口,但其他厂商基本都没有。常规交互还能从后端返回的 Usage 中获取 Token 数,但进行上下文压缩后怎么办?总不能压缩后专门发一次请求统计 Token 数吧。
因此,Agent 具有本地统计能力是必须的。
通过字符或字节长度进行的估算是非常粗糙的,常见措施是 Agent 内置与后端 LLM 匹配的分词器(Tokenizer), 专用于 Token 计数。
但对于计数这样的简单任务,传统 Tokenizer 包含解码、编号存储等大量冗余功能,影响性能且容易导致应用程序体积增大。 而 toklen 库专注于 Token 计数,试图保证最小的体积和最佳的性能。
toklen 库是单线程的,这保证了 WebAssembly 的兼容性,因此可以用于 JavaScript/TypeScript 项目。
安装与使用
目前 toklen 是 Rust 库,而非可执行程序。 这意味着你可以将它作为第三方库导入自身的 Agent 项目。 另外,你也可以将它编译成 wasm 对象,用于 JS/TS 项目。
在 Cargo.toml 中导入依赖:
[dependencies]
toklen = "0.2"
首先,你需要创建一个 Tokenizer 对象,它决定了分词的具体规则。
当前的分词器遵循 HuggingFace Tokenizers 标准,这意味这你需要提供一个 tokenizer.json 文件。
你可以从 这里 了解到它的更多信息。
为方便演示,这里给出 Deepseek-V3 的 tokenizer.json 文件,由 Deepseek 官方的 API 文档提供。
那么你可以像这样创建 Tokenizer :
// 读取 json 文件,存入字符串
let json = std::fs::read_to_string("tokenizer.json").unwrap();
// 从 json 字符串(或 u8 数组) 创建分词器
let tokenizer = toklen::Tokenizer::from_json(&json).unwrap();
对于命令行程序,或者需要导出为 wasm 的项目,考虑将配置文件内联:
// 将配置文件数据内联存储
const TOKENIZER_JSON: &[u8] = include_bytes!("tokenizer.json");
// 从内联数据创建分词器
let tokenizer = toklen::Tokenizer::from_json(TOKENIZER_JSON).unwrap();
如果 JSON 解析失败,或存在不支持的 tokenizer.json 配置,构造函数会返回 Err 。
toklen 支持大部分常用的配置,但并非全部支持,你可以从 项目文档 了解具体情况。
当 Tokenizer 创建后,就可以用它统计 Token 数了,这非常简单:
let count = tokenizer.encode_len("Hello, world!").unwrap();
println!("Token count: {count}");
encode_len 函数可能失败,比如遇到非法的 UTF8 字符。失败时返回 Err(len / 4),表示 Token 的估计值,默认是字符串的字节数除以四。
Tokenizer 对象是可复用的,并且是线程安全的。通常推荐将其作为全局单例存储,因为它的初始化开销远大于后续的计数开销。例如:
static TOKENIZER: LazyLock<Tokenizer> = LazyLock::new(|| Tokenizer::from_json(JSON).unwrap());
LazyLock 会在第一次访问时将其初始化,后续访问可以复用同一个对象。无需在意内部的原子变量开销,这和 C++ 函数局部静态变量的隐藏开销基本一致。
注意,toklen 是单线程的,这意味着 encode_len 函数本身不会将任务拆分到多个线程执行。
但你仍然可以在多个线程调用同一个分词器的 encode_len 函数,同时计算多个文本的 Token 数,这不会发生干扰。
性能对比
我们将它与 HuggingFace 的 Tokenizers 库,以及多线程的 Fastokens 库进行了性能比较,结果如下:
| 序号 | token 数 | toklen | tokenizers | fastokens |
|---|---|---|---|---|
| 0 | 9871 | 226.9ms - 43.5k TPS | 177.3ms - 55.6k TPS | 444.4ms - 22.2k TPS |
| 1 | 9871 | 1.380ms - 7152k TPS | 10.227ms - 965k TPS | 1.727ms - 5715k TPS |
| 2 | 9871 | 1.479ms - 6674k TPS | 12.433ms - 793k TPS | 1.855ms - 5321k TPS |
| 3 | 9871 | 1.774ms - 5564k TPS | 11.592ms - 851k TPS | 1.107ms - 8916k TPS |
| 4 | 9871 | 1.305ms - 7564k TPS | 9.853ms - 1001k TPS | 0.978ms - 10093k TPS |
| 5 | 9871 | 1.278ms - 7723k TPS | 9.730ms - 1014k TPS | 0.903µs - 10931k TPS |
| 6 | 9871 | 1.263ms - 7815k TPS | 9.554ms - 1033k TPS | 1.077ms - 9165k TPS |
| 7 | 9871 | 1.260ms - 7834k TPS | 10.69ms - 922k TPS | 0.948µs - 10412k TPS |
| 8 | 9871 | 1.288ms - 7663k TPS | 9.724ms - 1015k TPS | 1.101ms - 8965k TPS |
| 9 | 9871 | 1.279ms - 7717k TPS | 11.03ms - 894k TPS | 1.001ms - 9861k TPS |
每轮测试使用的文本相同,因此 Token 数始终为 9871 。每组数据前面的是耗时,右边的 TPS 是对应的每秒 Token 处理数。第 0 行的速度特别低,因为首次计算包含了分词器的初始化耗时。
整体上,单线程 toklen 的处理速度可以达到单线程 tokenizers 库的 7.5 倍, 可以和多线程的 fastokens 库持平,这正是因为 toklen 去除了大量计数无关逻辑。
具体的测试代码可以参考 Github 仓库的 bench 文件夹。
许可证
MIT or Apache-2.0