cedar で最長一致でトークナイズするメモ
RWKV world tokenizer の情報と C++ での実装メモ
で, hat-trie 利用でトークナイズしてみました.
別途, 最長一致ベースでの C++ 高速形態素解析 jagger の python binding に取り組んでいましたので,
高速形態素解析 Jagger の Python binding のメモ
ここで最長一致のための配列ベースの Trie 木(Prefix tree)の実装 cedar
を使ってのトークナイズももやってみます!
(edar のソースコードはとてもつらぽよなので, 理解で結構時間を溶かしたよ... 🥲)
cedar は C++ 例外の利用がないので, WASM/WASI への組み込みもやりやすくなります.
(sandbox 環境でセキュアに文字列処理したいときとか)
cedar
Prefix tree の配列表現で高速に prefix search するライブラリってところです.
余談
trie 木の配列表現を, なぜか NLP 界隈は「ダブル配列」(double-array)と呼んでいますが, これはひじょうにややこしい和製日本語だと思います.
英語では double-array とかで検索してもほぼヒットしないです.
ちなみに CG 界隈(?)ですと, (Q)BVH など, 配列で木構造を表現するのは昔から行われてきており, なにか特定の名前があるわけではなく, フツーに n-ary BVH とか child-sibling tree(N 文木) と呼んでいます.
N 分木(多分木)を二重連鎖木にしてオフセット配列表現するメモ
したがって, ここでも単に Trie 木/Prefix tree の配列表現と呼ぶようにします.
variant?
Jagger https://www.tkl.iis.u-tokyo.ac.jp/~ynaga/jagger/ では, ccedar_core.h と, cedar をちょっと改変(&簡略化?)したものがあります.
API はほぼ同じなので, たぶんどちらを使ってもよいかと思います.
ただ, ccedar_core.h 利用の Jagger では, 内部的に key となる文字列(UTF-8 文字列)を Unicode の code point 値に変換し int(4 bytes)の配列で扱っています.
UTF-8 ですと最大 4 bytes なので, UTF-8 のバイト列を 4 byte に padding してそのまま渡してしまってもよいところですが, Trie 木の圧縮のためでしょうか.
最長一致で取り出す
ソースコードがつらぽよでしかもドキュメントやサンプルコードが皆無に近くてつらぽよなところです.
commonPrefixSearch でいけるかと思いましたが, どうもダメで, commonPrefixPredict を使う必要がありました.
commonPrefixPredict は, 入力文字列に対して, その名の通り共通 prefix を持つキーを返します. vocab 次第で, 返ってくる候補は n 個あります.
例として, RWKV World の vocab を使います. https://zenn.dev/syoyo/articles/d47afddde67886
cedar にはいくつか API があります. まず commonPrefixSearch, commonPrefixPredict, exactMatchSearch を試します.
「である」を, 「で」「であ」「である」「である。」で処理させると, 以下になります.
で
commonPrefixSearch ():
で: found, num=1 9604:3
commonPrefixPredict ():
で: found, num=17 9604:0:19477: 42868:3:96704:あ 58050:6:149753:あっ 63636:9:173147:あった 58051:6:149754:あり 58052:6:149755:ある 42869:3:96719:き 58053:6:149759:きる 42870:3:96725:し 58054:6:149600:しょ 63637:9:173150:しょう 42871:3:96731:す 58055:6:149606:すが 42872:3:96748:の 42873:3:96749:は 58056:6:149616:はな 42874:3:96711:も
exactMatchSearch ():
id = 9604
であ
commonPrefixSearch ():
であ: found, num=2 9604:3 42868:6
commonPrefixPredict ():
であ: found, num=5 42868:0:96704: 58050:3:149753:っ 63636:6:173147:った 58051:3:149754:り 58052:3:149755:る
exactMatchSearch ():
id = 42868
である
commonPrefixSearch ():
である: found, num=3 9604:3 42868:6 58052:9
commonPrefixPredict ():
である: found, num=1 58052:0:149755:
exactMatchSearch ():
id = 58052
である。
commonPrefixSearch ():
である。: found, num=3 9604:3 42868:6 58052:9
commonPrefixPredict ():
である。: not found
exactMatchSearch ():
id = -1
vocab には「で」「であ」「である」があります.
今回は最長一致のがほしいので, commonPrefixPredict で候補がある間は処理を続け, 候補がなくなったら一つ前の状態の id を採用すれば最長一致ができます!
commonPrefixSearch, exactMatchSearch ではダメなのか
たとえば aa, aaaa がのある vocab の場合,
a
commonPrefixSearch ():
a: not found
commonPrefixPredict ():
a: found, num=2 0:1:256:a 1:3:258:aaa
exactMatch ():
id = -1
aa
commonPrefixSearch ():
aa: found, num=1 0:2
commonPrefixPredict ():
aa: found, num=2 0:0:256: 1:2:258:aa
exactMatch ():
id = 0
aaa
commonPrefixSearch ():
aaa: found, num=1 0:2
commonPrefixPredict ():
aaa: found, num=1 1:1:258:a
exactMatch ():
id = -1
a を与えると commonPrefixSearch は not found を返します. 入力文字が aa だった場合に, 最初の文字の処理で not found になってしまいうまく見つけてくれません.
exactMatch も同様です.
commonPrefixPredict ではダメか
commonPrefixPredict は, 候補が多いとバッファを多くとらないといけない & iterate 回数が増えて性能が出ない可能性があります.
また, jagger に付属の ccedar では commonPrefixPredictがありません!
traverse を使う
しょうがないので, 最終的には traverse を使うことにしました.
traverse は trie 木の subtree をたどる機能のようです(ソースコードがややこしすぎて読解できない).
static int lps(trie_t *ptrie, const char *line) {
size_t len = strlen(line);
size_t prev_from{0};
int prev_n{-1};
size_t from{0};
for (size_t i = 0; i < len; i++) {
size_t pos = 0;
// traverse(&line[i], from, pos, 1) でもよい
int n = ptrie->traverse(&line[i], from, pos, pos + 1);
if (n == trie_t::CEDAR_NO_VALUE) {
std::cout << "i " << i << "no value\n";
continue;
}
if (n == trie_t::CEDAR_NO_PATH) {
std::cout << "i " << i << "no path\n";
break;
}
prev_n = n;
prev_from = from;
}
std::cout << "--- lps ---\n";
std::cout << "prev_n " << prev_n << "\n";
std::cout << "prev_from = " << prev_from << "\n";
std::cout << "from = " << from << "\n";
return prev_n;
}
とりあえず 1 文字づつ処理させて, traverse していきます(from(trie 木の subtree(child, sibling index))の index は traverse を呼ぶとアップデートされます.
a
--- lps ---
prev_n -1
prev_from = 0
from = 97
aa
i 0no value
--- lps ---
prev_n 0
prev_from = 256
from = 256
aaa
i 0no value
i 2no value
--- lps ---
prev_n 0
prev_from = 256
from = 352
aaaa
i 0no value
i 2no value
--- lps ---
prev_n 1
prev_from = 258
from = 258
aaaaa
i 0no value
i 2no value
i 4no path
--- lps ---
prev_n 1
prev_from = 258
from = 258
b
i 0no path
--- lps ---
prev_n -1
prev_from = 0
from = 0
prev_n にマッチしているのが返ります!
- CEDAR_NO_PATH になったら共通 prefix を持つキーはなにも無いので, early terminate している
-
from==prev_fromであったら exactMatch しているはずなので, この時点で loop を抜けてもよいであろう.
これで, (たぶん計算量抑えて)最長一致ができます!
さらなる高みへ
exactMatchSearch の利用
LLM でのテキスト生成のためのトークナイズですと, 「ああ」「ああいう」が vocab にあったときにどちらも「ああ」で切り出しても問題が少ないみたいなケースでは, exactMatchSearch 利用でもいいかもしれません(多少コードが簡略化される).
自前実装
構成としてはレイトレでの BVH 構築とかなり同じなので, レイトレで BVH 実装の経験ある方はすんなりいけるでしょう...
レイトレ BVH 始めてのひとは
冬到来! NanoRT ではじめよう C++ レイトレーシング
を参考にしてレイトレ BVH 極めようね.
n 分木処理は IP アドレスのルーティングのやりかたも参考になるよ.
Poptrie: A Compressed Trie with Population Count for Fast and Scalable Software IP Routing Table Lookup (H. Asai et al., ACM SIGCOMM 2015)の日本語解説記事です。気が向いたので書きます。
Discussion