LinderaのFSTをDoubleArrayTrieに変更した話

Posted by johtani on Monday, October 5, 2020

目次

2020/10/06 11:00くらいにマージされました。

@minoru_osuka さんが開発を引き継いだLinderaというKuromojiのRustクローンがあります(リポジトリ) 。 最近趣味でRustを勉強しているので、こちらを少し手伝っています。

Rustの勉強仲間である@takuya_bさんや@ikawahaさんと話をしているときに、FST部分をDouble Array Trieに置き換えると速度が向上するのでは?という話が出まして、@takuya_bさんがDouble Array Trieを作るらしいという話になったので、下準備などをしつつ、作ってもらったライブラリyadaを組み込んでみたという話です。

ベンチマークの追加

下準備として、今のLindera(FST実装)がどのくらいの性能なのか?というのを計っておく必要があります。 幸いにも、Linderaのオリジナルの開発者の方が、criterion.rsというライブラリを使ったベンチマークプログラムを作成してくれていました。

ただ、1種類だけだと少し心もとないなというのと、長い文章やパターンを増やしたほうが良さそうだなということで、 ベンチマーク自体をいくつか追加しました(追加したときのPR)。

種類としては、5種類のベンチマークです。

  1. システム辞書のみのTokenizerのコンストラクタ呼び出し
  2. カスタム辞書ありのTokenizerのコンストラクタ呼び出し
  3. システム辞書のみのTokenizerのtokenize処理の呼び出し
  4. カスタム辞書ありのTokenizerのtokenize処理の呼び出し
  5. 青空文庫の坊っちゃんのテキストをシステム辞書のみのTokenizerでtokenize

1,2はコンストラクタ部分だけの処理をベンチマークテストする目的で作成しました。 LinderaはTokenizerがtokenize処理するのに利用するデータをいくつか内部で保持しています。 これらはファイルにシリアライズされており、Tokenizerのオブジェクト生成時に読み込みやデシリアライズ処理が発生します。 この部分だけも速度を計測したい目的でコンストラクタだけを切り出しました。

3,4はTokenizerのメインの処理です。コンストラクタはベンチマークの対象外にしました。 純粋にtokenizeの処理だけを切り出して計測するためです。 カスタム辞書がある場合、ない場合は念の為切り出した形になっています。

5は長い文章(文章が多いのでバリエーションも増える)を扱いたいために別にしました。

これで、一応下準備が完了です。 ちなみに、Criterionは賢くて、前のベンチマークの結果と最新の結果を比較してくれる機能があります。 どんな感じで出てくるかはベンチマーク結果をご覧ください。

yadaの組み込み

ベンチマークの準備をしていたところyadaがリリースされたので、Linderaへの組み込みを検討し始めました。

中身の理解

lindera-fstを利用して、prefix searchしている処理があるので、そこで利用しているFSTをyadaに置き換えれば良さそうだと判断して、 処理を読んでいきます。

  • Tokenizerは、PrefixDictという構造体でlindera-fstを利用している
    • prefixメソッドが入力文字列を元に、FSTを前方一致検索して、ヒットした単語の情報をIteratorとして取り出せる(単語の情報は「入力文字列の先頭からの文字数」と「ヒットした単語のWordEntry構造体」)
    • PrefixDictのfstは辞書(例:ipadic)ごとにlindera-<辞書名>-builderで生成される
  • システム辞書としては、デフォルトではlindera-ipadic-builderでfstを構築している

という感じです。 また、辞書周りのファイルがそれぞれどんな役割なのか、どんなデータの持ち方をしているのか?といった点を、変更点の調査のついでに書き出してみました。lindera-dictionary/FILES.md。TODOになっている部分も追記が終わっています(PR)

変更点

実際に変更したプログラムの詳細についてはのPRを見ていただくとして、簡単には以下の点になります。

  1. Rustのバージョンを1.46.0に(おもにREADME.md)
    • yadaが利用している機能に1.46.0で導入された機能があるため
  2. lindera-fstをyadaに変更(lindera-core/Cargo.toml, lindera-ipadic-builder/Cargo.toml)
    • 合わせて、dict.fstというファイル名をdict.daに変更
  3. dict.daに関して構築部分と検索部分を変更
    • FSTではFSTから返ってくる値(入力文字列に出てきた単語に関連する値)はu64だったが、yadaのDoubleArrayがu32しか扱えないため、u32に変更。テストの記述はしていないが、扱うデータ的にu32で問題なさそうだったので。
    • 検索部分:PrefixDict構造体のprefixメソッドでDoubleArrayprefix_common_searchを使用
      • DoubleArray自体がprefix_common_searchのメソッドを持っていたので、処理が簡単に置き換え可能だった。FSTはprefixメソッド内で独自で前方一致検索を実装していた。
    • 構築部分:lindera-ipadic-builder/src/lib.rsbuild_dictbuild_user_dictdict.da構築処理
      • ipadicのCSVファイルを読み込んで、見出し語をキーに、辞書にある単語情報のベクタを値とするBTreeMapを生成し、このBTreeMapに基づいてFSTを構築していた部分をDoubleArray構築処理に置き換えた。
      • シフト演算などで、実際の値(dict.vals)へのポインタを作っていたのだが、ここの処理を読み解くためにFILES.mdを書き出した。

という感じです。そもそもデータ構造がどうなっているのか?から読み解いて、変更部分を洗い出して変更していった形になります。 取り込み作業中にいくつかyadaに要望(このへん)を上げて、変更を取り込んでもらい、最終的にyadaのバージョン0.3.2で問題なく動きそうだという形になりました。@takuya_bさん、対応ありがとうございました。

エッジケースバグの発見

作っててよかった、テストケースでした(実際にはベンチマークテストですが)。 取り込み作業中に、Lindera本体のcargo testはすべてOKになるが、ベンチマークを取ろうとしたときに、坊っちゃんの文字列を入力にしたベンチマークが失敗するという事象が発生しました(PRのコメント参照)。 切り分けのために、入力の文章のどこでおかしくなるのか?DoubleArrayのbuildメソッドに渡している値がおかしくないか?などをすこしずつ調べていくと次のバグが判明したという感じです。

特定のデータ(ipadicの見出し語一覧)をDoubleArrayに入れて、特定の文字列(「は相」)をcommon_prefix_searchにいれたら、 返ってくる情報(0から何バイト目の文字が一覧に存在した)が、不正な値が返ってくるというバグでした。 @takuya_bさんに見てもらいつつ(DoubleArrayの中身わからん。。。)、修正してもらいました。素早い対応ありがとうございます。

やはり、いろんな文字列入れてテストしてみるの重要ですね。 ということで、ベンチマークだけでなく、テストケースとしても坊っちゃんのファイルを読み込んでトークナイズするようにPRでテストケースを追加しています。

ベンチマーク結果

yadaを利用した変更が終わったので、再度cargo benchを実行して計測です。 計測としては、masterブランチでまずcargo benchを実行し、yadaの実装をしたブランチに切り替えてからcargo benchを実行します。 すると、Criterion? cargo benchが、最終的な結果に前回との差分でどのくらい性能が改善、改悪したかも合わせて出力してくれます。 実行環境と結果は以下のとおりです。

  • MacBook Pro 16インチ
    • CPU:Core i7 6コア 2.6GHz
    • メモリ:32GB

コンストラクタのベンチマークについては10%ほど性能が悪くなっています。 これは、FSTよりもDoubleArrayTrieのほうがデータが大きくなってしまうためだと思われます。 実際にファイルのサイズは次のようになりました。yada(DoubleArrayTrie)のほうが2倍以上大きいことがわかります。 また、このファイル以外にもLinderaが利用しているデータはありますが、それらは今回変更の対象にはなっていません。 なので、単純にこのファイルの読み込みの処理に時間がかかっているのだと想像できます。

  • 2147765 / FST / dict.fst
  • 5425152 / yada / dict.da

tokenizeのベンチマークについては、11%〜28%の改善が見られました。 文章から、内部に保持している辞書に存在する単語を見つけ出す処理に利用されるのがFST、DoubleArrayTrieです。 今回の変更では、この処理に利用しているデータ構造だけを変更しました。 実際には

  • DoubleArrayTrieを用いた単語の検索処理
  • 見つかった単語の持つ値(data.valsのオフセット情報)を元にシフト演算

といった処理が実行されます。シフト演算はu64だったものがu32に変更されたくらいなので、大した処理量ではないかと。 大部分はDoubleArrayTrieを利用したルックアップ処理が速度向上に寄与していると思います。

まとめ

最近Linderaに加えた変更、作ったPRについて少しブログにまとめてみました。 ちなみに、まだPRの段階でレビュー&リリース待ちという感じです。

実際には作ってもらったライブラリを組み込んでみたというだけなのですが、速度が向上した結果が見れたのは面白いです。 また、基本的なデータ構造とかアルゴリズムの勉強にもなりました(2次元配列を1次元配列に押し込むとか)。このへんも今後も勉強していきたいです。

組み込む際に色々と協力していただいた@takuya_bさん、@ikawahaさん、巻き込んでくれた@minoru_osukaさんに改めて感謝いたします。

Rustや形態素解析のプログラムの勉強を兼ねて、今後もなにか改善できる部分がないかなどを見ていこうと思っています。 Rustで形態素解析をしたいという人がどのくらいいるかはわかりませんが、おかしなところや疑問点などあればコメントください。


comments powered by Disqus