Zenn
🎃

Node.jsの@dqbd/tiktokenとtiktoken-nodeのトークン数の算出を比較した

2023/04/25に公開
Node.js
OpenAI
#
tiktoken
tech

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)

説明を読むとエンコーダー設定を明示的に与えるぐらいしか違いがなさそうですが、一応分けて検証します。

ベンチマーク

  1. tiktoken-node
  2. @dqbd/tiktoken
  3. @dqbd/tiktoken/lite

のトークン化のエンコード処理を10文字、100文字、1000文字で10000回づつ繰替えてしてみて平均速度を測ります。

benchmark.js
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用途なら基本これを使えばよさそう。

https://platform.openai.com/docs/guides/embeddings/what-are-embeddings

戻り値の違い

@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