実質無料で気安くベクトル検索を体験する

ベクトル検索

世間は AI 花盛り。パブリッククラウド各社の新機能発表でも AI 関連の機能が盛り沢山です。
AI アシスタントとのチャットや画像の生成といった機能がわかりやすい例ですが、他にも比較的地味ではありつつ強力な機能がいろいろあります。今日はそういった機能の中から、ベクトル検索を取り上げてみたいと思います。

たとえば Google Cloud では AlloyDB や BigQuery がベクトル検索に対応した database として名前があがっています。BigQuery では model の deploy や instance を用意する必要なく SQL から直接利用ができ 非常に便利そうです。

その機能としては、探したいことが書かれている文書を探したり類似する画像を探したりできる、いろいろと応用の夢が広がるものです。しかし、どうしても AI 関連の機能は「でも… お高いのでは?」という印象がつきまといます。実際、BigQuery とそこからリンクがはられている Vertex AI 関係の pricing のページ を見ましたが、実際にやってみる前にある程度正確に見積もるのはかんたんではないという印象を受けました。

そこで今回は、この夢のあるベクトル検索を、費用に心配無く体験してみたいと思います。

事前準備 : database として Redis を用意する

特に汎用キャッシュとして人気があり様々な場所で活用されている Redis ですが、ベクトル検索に対応した database のひとつです。クラウド上のマネージドサービスとしても利用可能な選択肢が多いですが、もちろんローカルで実行もできますので、今回はこれで試してみたいと思います。

ベクトル検索は、Redis の拡張版である Redis Stack で利用可能とのことなので、これをお手軽に Docker で起動してみます。こちら に Docker で起動するためのドキュメントがあります。

docker run -it --rm -d redis/redis-stack

何の追加設定も必要なく、デフォルト設定で問題なく起動しました。

念のため以下のように動作確認だけしておきました。(準備としては不要なステップです)
起動した container の shell に入って、サンプルデータ を load させてみます。

docker exec -it fac10a229db1 bash
root@fac10a229db1:/# redis-cli < /tmp/bicycles.txt
root@fac10a229db1:/# redis-cli
127.0.0.1:6379> FT.SEARCH idx:bicycle "@price:[1000 +inf]"
 1) (integer) 5
 2) "bicycle:1"
 3) 1) "$"
 ...

ちゃんとデータが入っていることが確認できました。正常に動作していそうです。

事前準備 : Python まわり

Redis 公式 document で、ベクトル検索に関するわかりやすい例が用意されています。
Redis as a vector database quick start guide

これを流すことで一通りのフローやできることが体験できるようになっています。
中でいくつかの Python の library や Microsoft の embedding model を使っているので、適宜確認し、必要に応じて差し替えなどしてください。

ちゃんとした運用環境向けに構築する場合は database と application の container はわけるところですが、今回は横着して、先ほど起動した Redis の container でこの python 実行もしてしまいます。

先程のように docker exec で container 内の shell を使ってもよいのですが、後ほど python コードの変更をするときに便利なように、VS Code の Remote extension を使うのをおすすめします。VS Code を開き、Command Palette (Shift + Command + P (Mac) / Ctrl + Shift + P (Windows/Linux)) から “Attach to Running Container…” を選択し、target container として起動していた Redis のものを選択すると、container 内に devcontainer が setup され、Terminal はもちろんのこと Explorer や Editor も利用可能になります。

Redis の container には都合よくすでに python3 も入っていました。使い捨て利用のつもりで、深く考えずに以下のように必要そうなものを入れてしまいましょう。

apt update
apt install pip
pip install numpy pandas redis sentence-transformers tabulate

上記の quick start guide ですが、いくつかあるコードブロックのうちどれでもコピーすると、ひとまとまりのサンプルコードがクリップボードにコピーされるようになっているようです。説明文内では細かく分断されていますが、手元ではひとつの python file にペーストして見るとわかりやすいと思います。

また、ページ前半はアカウント取得についての説明になっています。今回はローカル実行なのでその部分は読み飛ばして問題ありません。

ベクトル検索をやってみる・サンプルの処理の流れ

さて、ようやく実行してみましょう。

上で書いたように、サンプル全体をコピペしたものを python file として実行するだけでおそらく最後まで動くと思います。ただ console への出力はされないので、様子を見ながら実行する場合は中で print() を足したり、interpreter で実行するとよいでしょう。

サンプルコードの中で何をしているか簡単に書いておきます。

• サンプルデータ を download し、key をつけて Redis へ load
• Sentence Transformers を使って embeddings を生成し、元のデータと合わせて保存
• key を使って get() し、embeddings が入ったデータを確認
• index を生成し、それを利用した検索ができることを確認
• 複数の query による検索を実行し、その結果を表に整形

以下に、サンプルコード後半で出力される表の一部を引用します。database にあるデータに対して任意のテキストで検索をかけ、このように出力が得られる具体的な例になっていて、検索機能の効果がイメージしやすいと思います。表の query が検索語、score が類似度で、あとはヒットしたデータです。(データ内容はサンプルコード内で download されたサンプルデータです)

サンプルコードは検索のしかたなどにいくつかのバリエーションがある例になっているので、参考にしつつすこし手を入れながら実行するのもよさそうです。以下に検索の実装を一部引用しておきます。

query = (
    Query('(*)=>[KNN 3 @vector $query_vector AS vector_score]')
     .sort_by('vector_score')
     .return_fields('vector_score', 'id', 'brand', 'model', 'description')
     .dialect(2)
)

...

            client.ft("idx:bikes_vss")
            .search(
                query,
                {
                    "query_vector": np.array(
                        encoded_query, dtype=np.float32
                    ).tobytes()
                }
                | extra_params,
            )

書き方に慣れが必要そうですが、このような表現で「近い」ものの検索ができるなんて、非常に強力な機能ですね！

ここで使われている KNN (k-nearest neighbor, k近傍法) に関しては Redis の document の他にもよい読み物がたくさんありました。末尾の参考のところにもひとつつけておきます。

まとめ

様々なアプリケーションに応用できそうな機能であるベクトル検索を実質無料で体験してみました。
あとはアイディアさえあれば、いろんなことができそうです。それがいちばん難しいのですが 😇

参考

• Redis Docs > … > Vectors
• Sentence Transformers で使用できる model の例 (いずれも apache-2.0 license)
  • paraphrase-MiniLM-L6-v2
  • msmarco-distilbert-base-v4
• k近傍法アルゴリズムとは (IBM)