Node.jsの@dqbd/tiktokenとtiktoken-nodeのトークン数の算出を比較した
Node.jsでOpenAIの言語モデルに利用するBPEトークン数を計算するにはサードパーティの@dqbd/tiktokenかtiktoken-nodeのNPMモジュールが使えるのですが2つの実装はことなります。
@dqbd/tiktokenはRustで書いたコア実装をwasm-packでWasmにしてNode.jsから呼び出します。
tiktoken-nodeはNAPI-RSを使いピュアRust実装のtiktoken-rsを呼び出します。
Wasmを通さないぶんtiktoken-nodeのが早いのでは? と思ったので確かめてみることにしました。
@dqbd/tiktoken/lite
@dqbd/tiktokenにはブラウザライクな環境から呼び出すことを想定したliteバージョンがあります(TensorFlow Liteのようなものと解釈)。
In constrained environments (eg. Edge Runtime, Cloudflare Workers), where you don't want to load all the encoders at once, you can use the lightweight WASM binary via @dqbd/tiktoken/lite.
https://github.com/dqbd/tiktoken/tree/26bf591198b1077361f7a9c47e2ca826d32cccc9/js)
説明を読むとエンコーダー設定を明示的に与えるぐらいしか違いがなさそうですが、一応分けて検証します。
ベンチマーク
- tiktoken-node
- @dqbd/tiktoken
- @dqbd/tiktoken/lite
のトークン化のエンコード処理を10文字、100文字、1000文字で10000回づつ繰替えてしてみて平均速度を測ります。
const { performance } = require('node:perf_hooks');
const { getEncoding: getNAPIRSBasedEncoding } = require('tiktoken-node');
const { get_encoding: getWasmPackBasedEncoding } = require('@dqbd/tiktoken');
const { Tiktoken: TiktokenLite } = require("@dqbd/tiktoken/lite");
const cl100kConfig = require("@dqbd/tiktoken/encoders/cl100k_base.json");
// 対象のテキスト(文字数を。。。で調整する)
const text = '今日はいい天気ですね'.padEnd(0, '。');
// ベンチマークする回数
const iterations = 10000;
const tokenizer = 'cl100k_base';
function doEncode(enc, text) {
enc.encode(text);
}
// tiktoken-node(NAPI-RS)
const startTimeNAPIRS = performance.now();
const NAPIRSBasedEncoding = getNAPIRSBasedEncoding(tokenizer);
for (let i = 0; i < iterations; i++) {
doEncode(NAPIRSBasedEncoding, text);
}
const endTimeNAPIRS = performance.now();
// @dqbd/tiktoken(WasmPack)
const startTimeWasm = performance.now();
const wasmPackBasedEnc = getWasmPackBasedEncoding(tokenizer);
for (let i = 0; i < iterations; i++) {
doEncode(wasmPackBasedEnc, text);
}
wasmPackBasedEnc.free();
const endTimeWasm = performance.now();
// @dqbd/tiktoken/lite(WasmPack)
const startLiteWasmPackBasedEnc = performance.now();
const liteWasmPackBasedEnc = new TiktokenLite(
cl100kConfig.bpe_ranks,
cl100kConfig.special_tokens,
cl100kConfig.pat_str
);
for (let i = 0; i < iterations; i++) {
doEncode(liteWasmPackBasedEnc, text);
}
liteWasmPackBasedEnc.free()
const endLiteWasmPackBasedEnc = performance.now();
// Result
const avgTimeNAPIRS = (endTimeNAPIRS - startTimeNAPIRS) / iterations;
const avgTimeWasm = (endTimeWasm - startTimeWasm) / iterations;
const avgTimeWasmLite = (endLiteWasmPackBasedEnc - startLiteWasmPackBasedEnc) / iterations;
console.log(`=== text.length=${text.length}, iterations=${iterations}, tokenizer=${tokenizer} ===`)
console.log(`tiktoken-node(NAPI-RS) の平均実行時間: ${avgTimeNAPIRS} ms`);
console.log(`@dqbd/tiktoken(WasmPack) の平均実行時間: ${avgTimeWasm} ms`);
console.log(`@dqbd/tiktoken/lite(WasmPack) の平均実行時間: ${avgTimeWasmLite} ms`);
=== text.length=10, iterations=10000, tokenizer=cl100k_base ===
tiktoken-node(NAPI-RS) の平均実行時間: 0.009391658300161362 ms
@dqbd/tiktoken(WasmPack) の平均実行時間: 0.1439790875002742 ms
@dqbd/tiktoken/lite(WasmPack) の平均実行時間: 0.14135069580078125 ms
== text.length=100, iterations=10000, tokenizer=cl100k_base ===
tiktoken-node(NAPI-RS) の平均実行時間: 0.04801496669948101 ms
@dqbd/tiktoken(WasmPack) の平均実行時間: 0.21088285830020906 ms
@dqbd/tiktoken/lite(WasmPack) の平均実行時間: 0.2213586500003934 ms
=== text.length=1000, iterations=10000, tokenizer=cl100k_base ===
tiktoken-node(NAPI-RS) の平均実行時間: 3.5001618667006493 ms
@dqbd/tiktoken(WasmPack) の平均実行時間: 6.272181949999928 ms
@dqbd/tiktoken/lite(WasmPack) の平均実行時間: 6.3225741499990225 ms
以下のことが分かりました
- tiktoken-node(NAPI-RS)が2〜15倍程度早かった
- 文字数を増やすごとに@dqbd/tiktokenとの差は縮まる
- @dqbd/tiktoken/liteの実行速度に差はなかった
(※平均実行時間だけじゃなくてグラフで分布を出したい)
環境
❯ system_profiler SPHardwareDataType
Hardware:
Hardware Overview:
Model Name: MacBook Pro
Model Identifier: MacBookPro18,1
Model Number: MK1F3TH/A
Chip: Apple M1 Pro
Total Number of Cores: 10 (8 performance and 2 efficiency)
Memory: 16 GB
Tokenizerの選択
text-embedding-ada-002モデルのTokenizerがcl100k_baseなのでGPT-3.5, 4用途なら基本これを使えばよさそう。
戻り値の違い
@dqbd/tiktokenの方は返り値がUint32Arrayになっている。
> enc.encode(text) # tiktoken-node
(5) [57933, 71634, 57933, 62903, 71634]
> enc.encode(text) # @dqbd/tiktoken
Uint32Array(5) [57933, 71634, 57933, 62903, 71634, buffer: ArrayBuffer(20), byteLength: 20, byteOffset: 0, length: 5]
Discussion