elasticsearch で cosine類似度検索する
全文検索エンジンで cosine 類似度検索できるらしいというのを bert × elasticsearch の記事で見かけてとてもたのしそうだったので、自分でも環境作るところからやってみました。
やっているのは以下の内容です
- docker/docker-compose で elasticsearch x kibana x jupyter の立ち上げ
- jupyter の python から実際に特徴量ベクトルの登録 + 検索の実行
- kibana での可視化 (ちょっとだけ)
紹介するコードはすべて以下のリポジトリから参照できます。
Requirements
- docker
- docker-compose
Setup
docker-compose build docker-compose up -d
以上実行するとサーバーが3つ立ち上がります。docker-compose ps をやって以下のような感じになってれば成功です。
23:00:17 in start-elastic-search on master on 🐳 v18.06.1 took 24s
➜ docker-compose ps
Name Command State Ports
----------------------------------------------------------------------------------------------------------------------------
startelasticsearch_elasticsearch_1 /usr/local/bin/docker-entr ... Up 0.0.0.0:9200->9200/tcp, 0.0.0.0:9300->9300/tcp
startelasticsearch_jupyter_1 /usr/bin/tini -- jupyter n ... Up 0.0.0.0:8888->8888/tcp
startelasticsearch_kibana_1 /usr/local/bin/dumb-init - ... Up 0.0.0.0:5601->5601/tcp
elasticsearch
今回の主題である elasticsearch は localhost:9300 でつながります。これを直接 REST API で叩いても当然つかうことが出来ます。
定義ファイルは ./elasticsearch ディレクトリに入っています。
詳しくは elasticsearch 公式 docker document を参考にして下さい
https://www.elastic.co/guide/en/elasticsearch/reference/current/docker.html
Production Usage の時の注意点 https://www.elastic.co/guide/en/elasticsearch/reference/current/docker.html#_notes_for_production_use_and_defaults とかもあります。
Kibana
Kibana は elasticsearch 用のクライアントアプリケーションです。 localhost:5601 で kibana container とつながっています。開いて適当なサンプルデータを突っ込むとわかりますがとにかくめっちゃいろんなことができることはわかります。開いて色々動かすだけでもわりと楽しいです。
設定について
kibana の container はデフォルトで elasticsearch:9300 につなごうとするので elasticsearch も default の 9300 で起動するよう特に何もいじっていません。
設定ファイルを弄りたい場合には environment に記述する若しくは kibana.yml を bind して up すると良いです。詳しくは kibana の docker document を参考にして下さい。
http://localhost:5601 にアクセスすると kibana server にアクセスできます。
jupyter
python から elasticsearch 使いたいよね、ということで jupyter notebook server を用意しています。
jupyter には http://localhost:8888/tree/notebooks からアクセスできます。 password は "dolphin" です。
python 環境には elasticsearch が pip で入っていますので python client はデフォルトで使うことが出来ます。
Sample notebook
いくつかサンプルのコードが入っていますので参考にして下さい。
特徴ベクトルでの類似度検索
~/client/notebooks/sample_search_by_cosine-sim.ipynb に特徴ベクトルでの類似度検索を行うサンプルが入っています。
128 次元のベクトルを "matching" index に 100000 件登録して、クエリのベクトルとコサイン類似度の意味で近い順にソートする、ということを行っています。
例えば以下のような例が応用でしょうか。
- 画像の特徴ベクトルでの類似度検索
(クエリ画像と、マスタ画像の類似度を arcface や center-loss などの metric learning で得たモデルの特徴量でソート) - 文章を表す特徴量でのソート
(SWEM など文章を埋め込む(embedding)手法)
Result
検索の時間は以下のようになりました。 elasticsearch 爆足だぜ :D
| method | %%timeit |
|---|---|
| scipy.spatial.distance & list内包記法 | 3.27 s ± 17.8 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) |
| scipy.spatial.distance & joblib parallel | 5.2 s ± 108 ms per loop (mean ± std. dev. of 7 runs, 1 loop each) |
| elasticsearch | 43.7 ms ± 661 µs per loop (mean ± std. dev. of 7 runs, 10 loops each) |
Kibana でも見る
上記の notebook を実行すると実際に index が出来上がるので, Kibana からもデータを見ることが出来ます。
インタラクティブに見る
http://localhost:5601/app/kibana#/management/kibana/index_pattern?_g=() にアクセスすると index の pattern をいれる画面が立ち上がります。ここに作成した "matching" と入力して Next Step > create index patterns と進みましょう。
その後左側のナビゲーションの上から二番目、方位磁石っぽいアイコンを押します。するとずらずらっとデータが 100001 件出てきます。これが今 elasticsearch に "matching" index として登録されているドキュメントです。
この画面もまあまあできることが多いのですが一番わかり易い使い方はうえの search toolbar だと思います。ここにカーソルを当てるといい感じにクエリをサジェストしてくれます。なかなか楽しいです。
また http://localhost:5601/app/kibana#/dev_tools/console?_g=() からクエリを叩くことも出来ます。たとえば
GET matching/_search { "query": { "match_all": {} } }
とかやると全件取得が出来ます。楽しいですね。
