Groongaでのセマンティックサーチの実装

Groongaの開発をしている須藤です。Groongaはもともとキーワードベースの全文検索機能を提供する全文検索エンジンとして開発されましたが、このご時世なので、セマンティックサーチサポートも拡充しています。この記事では現状のGroongaのセマンティックサーチのサポート度合いを紹介します。

セマンティックサーチサポートに必要な機能

セマンティックサーチとはテキストを埋め込み表現（エンべディング、embedding、数値ベクトル）に変換し、埋め込み表現がどのくらい似ているかどうかで検索する機能です。埋め込み表現は、うまく作ることができれば字面ではなくその字面のテキストが持っている意味を表現することができます。そのため、埋め込み表現の類似度を使うことで違う字面だけど似ている意味のテキストを見つけられます。たとえば、「ラーメン」でも「中華そば」を探せるというようなものです。

セマンティックサーチに必要なものは次の2つです。

テキストから埋め込み表現を作る機能
高次元の数値ベクトル（埋め込み表現）間の類似度を高速に計算して、類似度順に並び替える機能

多くのセマンティックサーチのシステムは後者のみを提供しています。前者は別のシステムと組み合わせて実現することが多いです。

類似度関連の処理は、対象の埋め込み表現の数が数万とか数十万とかであれば全データを対象にしても高速に処理できます。

しかし、数百万以上になってくると厳しくなってきます。データ量が増えてきたら高速に検索するための補助データ・アルゴリズムを使います。HNSWやNGTやIVFなどです。これらは正確な結果ではなく近似した結果を返す代わりに高速に検索できます。

Groongaが提供する機能

Groongaは埋め込み表現を作る機能も類似度関連の処理も両方提供する設計にしました。

外部の埋め込み表現作成機能を使う場合は、次のメリットがあります。

埋め込み表現を作成するサービスはたくさんあるので、用途に合わせて好きなサービスと組み合わせることができる
埋め込み表現作成用リソースとセマンティックサーチ用リソースを分離できる
- たとえば、埋め込み表現作成時と検索時で必要なGPUリソースが違うことが多い

一方、次のデメリットがあります。

埋め込み表現をシステム間でやりとりしないといけない
- 最近の埋め込み表現は高次元のためデータがおおきくなりがち。たとえば、512次元の32bit浮動小数点数で表現された埋め込み表現は2KiBになる。1000件あると2MiBで、100万件あると2GiB。クエリーを実行するときも毎回やりとりしないといけない。
アプリケーションまたはライブラリーが埋め込み表現の管理をしないといけない

前者は大きいと大変だよねーというのでなんとなくわかると思います。ネットワーク通信量が多くなるとか必要なストレージやメモリーが多くなるなーとか。

後者はピンとこないかもしれないので簡単な例を出しておきます。

たとえば、従来の全文検索機能だと次のように全文検索できたとします。

connection.execute("select * table where full_text_search(column, 'キーワード1 キーワード2')")

セマンティックサーチをする場合はこんな感じになります。

embeddings = vectorize("キーワード1 キーワード2")
connection.execute("select * table where semantic_search(column, #{embeddings})")

でも、こう書きたいですよね？

connection.execute("select * table where semantic_search(column, 'キーワード1 キーワード2')")

検索システム側が埋め込み表現を管理してくれるとこういう書き方ができるのです。

書き方以外にもメリットがありますが、それはまたいつか紹介します。

今後、外部で作成した埋め込み表現のサポートを強化する可能性もあります（現時点では一部の機能は使えない）が、現時点ではGroonga内の埋め込み表現機能のサポートを優先して開発しています。

Groongaのセマンティックサーチの埋め込み表現機能

Groongaはライブラリーとして使えるLLM推論器llama.cppを組み込んで埋め込み表現を作れるようになっています。名前の通りC++のライブラリーとして使える点と、アクティブに開発が進んでいて新しい言語モデルの対応も積極的なことから採用しました。実装の詳細はわからないので機能改善やバグフィックスはできていませんが、使っている中で見つけたちょっとした問題を直すプルリクエストは送っています。

llama.cppにはHugging Faceにあるモデルを自動でダウンロードする機能がありますが、Groongaでも同様の機能を実装してあります。user-agent: llama.cppでhttps://huggingface.co/v2/${USER}/${MODEL}/manifests/latestにアクセスするとメタデータをダウンロードできて、それを使うとモデルのダウンロード場所がわかるという仕組みなのですが、user-agent: llama.cppが必須なので、Groongaでもuser-agent: llama.cppを使っています。

参考： https://github.com/ggml-org/llama.cpp/blob/945501f5ea4b8ca56b181ecb035e9ee3fb31f432/common/arg.cpp#L416

この機能を使うと、次のようにすることでテキストから埋め込み表現を生成できます。

table_create Data TABLE_NO_KEY
column_create Data text COLUMN_SCALAR ShortText

load Data
[
{"text": "Hello World"}
]

plugin_register functions/language_model

select Data \
  --output_columns ' \
      text, \
      language_model_vectorize("hf:///groonga/all-MiniLM-L6-v2-Q4_K_M-GGUF", text)'

https://huggingface.co/groonga/all-MiniLM-L6-v2-Q4_K_M-GGUF を自動でダウンロードしてHello Worldの埋め込み表現を生成します。

普通に使うときはlanguage_model_vectorize()で明示的に埋め込み表現を生成しないのですが、ここでは例として使いました。

llama.cppは非常にたくさんのオプションがあるのですが、現時点ではほぼ設定できません。ほぼデフォルト値でしか使えません。今後、必要になったものから順次サポートしていきます。

Groongaのセマンティックサーチのシーケンシャル検索機能

Groongaはシーケンシャル検索用の処理を独自で実装しています。SimSIMDを検討したりもしましたが、xsimdを使って、SIMD対応のベクトル間の距離・類似度を計算する機能を実装しています。cosine類似度・内積・L1ノルム（マンハッタン距離）・L2ノルム（ユークリッド距離）が他のプロダクトでもよく使われていたのでこれらを実装してあります。

SimSIMDを採用しなかった理由は忘れてしまいましたが、たしか、Groongaがサポートしている環境用のSIMD実装がなかったからだった気がします。実装したときに、SimSIMDと同じくらいの性能はでることを確認してあるので、独自実装だからといって性能が劣っているということはないはず。

Groongaのセマンティックサーチのインデックス検索機能

Groongaはライブラリーとして使える類似検索ライブラリーFaissを組み込んで大量データでも高速に検索できるようにしています。シーケンシャル検索とは違って近似検索になっています。

DuckDBやClickHouseが使っているUSearchも検討しましたが、mmap()ベースで書き込みできる機能を仕上げられなかったので諦めました。インデックスマージ機能に必要な機能としてmmap()ベースでの書き込み機能も実装していました。

Faissは各種部品を組み合わせて使えるようになっていて、当初想定していたよりも使いやすかったです。現時点では、FaissでいうところのIVFRaBitQとIVFFlat相当のインデックスを実装していますが、転置索引部分にはGroongaの既存の転置索引実装を使ったり、ポスティングリストに量子化した値を埋め込むのではなく別途カラムとして保存したりとGroongaの既存の設計になじませるようにしています。

インデックスを使うには次のようにします。

まず、次のようなデータがあるとします。

table_create Data TABLE_NO_KEY
column_create Data text COLUMN_SCALAR ShortText

load --table Data
[
{"text": "I am a boy."},
{"text": "This is an apple."},
{"text": "Groonga is a full text search engine."}
]

ここに、セマンティックサーチ用のインデックスを作ります。必要なものは次の3つです。

セマンティックサーチ用のプラグイン
各レコードごとに埋め込み表現関連の情報を保存するカラム
セマンティックサーチ用のインデックス

具体的にはこんな感じになります。

plugin_register language_model/knn

column_create Data rabitq_code COLUMN_SCALAR ShortBinary

table_create RaBitQ TABLE_HASH_KEY ShortBinary \
  --default_tokenizer \
    'TokenLanguageModelKNN("model", "hf:///groonga/all-MiniLM-L6-v2-Q4_K_M-GGUF", \
                           "code_column", "rabitq_code")'
column_create RaBitQ data_text COLUMN_INDEX Data text

これで、セマンティックサーチできるようになります。

select Data \
  --filter 'language_model_knn(text, "male child")' \
  --output_columns text

[
  [
    0,
    0.0,
    0.0
  ],
  [
    [
      [
        3
      ],
      [
        [
          "text",
          "ShortText"
        ]
      ],
      [
        "I am a boy."
      ],
      [
        "This is an apple."
      ],
      [
        "Groonga is a full text search engine."
      ]
    ]
  ]
]

検索時に埋め込み表現（数値ベクトル）ではなくテキストそのものを指定できていることがポイントです。Groongaが自動で内部で埋め込み表現にして類似度で検索してくれます。

条件として絞り込むのではなく、他の条件で絞り込んだあとに類似度でソートする場合はこうです。language_model_knn()を--filterではなく--sort_keysで使います。先頭の-は降順という意味です。類似度が高い順に先頭k件を取得したいので降順になるのです。

select Data \
  --filter '_id < 3' \
  --sort_keys '-language_model_knn(text, "male child")' \
  --output_columns text

まとめ

Groongaのセマンティックサーチの実装をあまり詳細に立ち入らないレベルで紹介しました。去年から少しずつ実装を進めていたのですが、インデックスも動くようになったのでまとめました。まだドキュメントがないので、ドキュメントが準備できるまでのつなぎの情報のつもりです。

セマンティックサーチはしたいけど、埋め込み表現を管理したくない！という場合はGroongaサポートサービスを提供していますのでお問い合わせください。設計支援やアプリケーションの開発支援だけでなく、導入支援や運用支援も幅広く対応しています。

ククログ

セマンティックサーチサポートに必要な機能

Groongaが提供する機能

Groongaのセマンティックサーチの埋め込み表現機能

Groongaのセマンティックサーチのシーケンシャル検索機能

Groongaのセマンティックサーチのインデックス検索機能

まとめ

関連情報