Elasticsearchの近似近傍探索によるドラえもんのひみつ道具検索エンジンを作ってみた

Elasticsearch 8 系から使用可能になった近似近傍探索¹を使って、ドラえもんのひみつ道具の自然言語検索ができる検索エンジンを作ってみた。

デモ動画のように、検索したいひみつ道具を説明する文章することで近しいひみつ道具が検索されます。

コードはGitHub に公開してあるので、興味のある方は手元で、動かして遊ぶことが出来ます。 poetry と Docker さえあれば動きます。

hurutoriya/doraemon-himitsu-dogu-search: Doraemon Himitsu Dogu Japanese semantic search based on Elascticsearch ANN

システムの概要図は以下

Indexing

graph LR;
text-->|json形式に構造化|json
json-->|説明文|HuggingFace
json-->Elasticsaerch
HuggingFace-->|特徴ベクトル化|Elasticsaerch

Query

graph LR;
Query-->|強くなる|id["HuggingFace(encoder)"]
subgraph Streamlit
    id["HuggingFace(encoder)"]
end
subgraph Elasticsearch
    ANN
end
id["HuggingFace(encoder)"]-->|"[1.2, ... 0.3]"|ANN

所感

文章の特徴抽出は、sonoisa/sentence-bert-base-ja-mean-tokens-v2 · Hugging Faceを使わせていただき、驚くほど簡単に実現できた。 実際はもっと精度を高めるには、fine tune などがいいのだろうが、システム側を作ることに注力したかったので今回は割愛

デモアプリの構築は streamlit を使って20mくらいで作れたので、相変わらず便利すぎて愛用している。今回の検索エンジンはCLIから実行もできるが、こうやってデモアプリがあったほうがそれっぽくて気持ちいい。

Elasticsearch の近似近傍探索は、今回ベクトル同士の近似近傍探索しかやっていないが、それもindexing, query 部分は公式ドキュメントを見れば事足りたので感動。

query = {
		"knn": {
				"field": "vector",
				"query_vector": sentence_embeddings[0],
				"k": 10,
				"num_candidates": 100,
		},
		"fields": ["name", "description"],
}
result = es.search(index=INDEX_NAME, body=query)

だけで近似近傍探索ができる世界になった。

ぶっちゃけて言えば、この規模で近似近傍探索のみやるなら、インメモリでも動くと思うのでElasticsearchを使う意義を考えざるをえないが、自分が触ってみたかったのでやってみた。

やってみたいこと

これで、近似近傍探索周りで遊べるパイプライン+アプリが出来たので暇が出来たら色々試してみたいなと思っている。 例えば、

• 近似近傍探索探索部分は、ベクターデータベースの Milvusでも実現できる内容になっているが、Elasticsearch の真価であるフィルタリングや従来の検索手法と組み合わせる²ことが出来ていないのでやってみたい。

• 特徴抽出のパイプラインは出来たので、もっと大規模なドキュメントを取り扱って課題にぶち当たると楽しそう

• Vespa やQdrant、Milvus など他のライブラリでも近似近傍探索を使った検索をやれる

Elasticsearch 8.4 introduces a hybrid search: ability to combine results from knn search with traditional search features (queries, aggs etc) and all this under a single familiar _search API.

参考資料

https://www.tv-asahi.co.jp/doraemon/tool/a.html

1. スクレイピング
2. sentens bert で学習
3. ベクトル吐き出す
4. ESに入れる
5. ES streamlit で行う

• https://huggingface.co/spaces/mya-mya/HimitsudoguSearch/tree/main
• HimitsudoguSearch - a Hugging Face Space by mya-mya
• 【日本語モデル付き】2020年に自然言語処理をする人にお勧めしたい文ベクトルモデル - Qiita

特徴ベクトル化

sentens bert を使って使って文章埋め込みベクトルにしてみる

• huggingface transformersで使える日本語モデルのまとめ – Yellowback Tech Blog
• sonoisa/sentence-bert-base-ja-mean-tokens-v2 · Hugging Face

ES indexing

• k-nearest neighbor (kNN) search | Elasticsearch Guide [master] | Elastic
• elastic/eland: Python Client and Toolkit for DataFrames, Big Data, Machine Learning and ETL in Elasticsearch
• Elastic Stack 8.0 の NLP で日本語センチメント分析を試してみた - 後編 - Qiita
• Elasticsearchを用いて類似度ベクトル検索をやってみてわかったこと - ログミーTech
• Python Elasticsearch 基本的な使い方まとめ - Qiita
• Elasticsearch v8に導入されたkNN search APIの検索速度を試す | Wantedly Engineer Blog

Footnotes

1. k-nearest neighbor (kNN) search | Elasticsearch Guide [master] | Elastic ↩

2. Elasticsearch 8.4 introduces a hybrid search by Mayya Sharipova ↩