LoginSignup
18
9

More than 1 year has passed since last update.

スタンバイAdvent Calendar 2021Day 7stanby
@takatori

ElasticsearchのANN Searchを試す

Last updated at Posted at 2021-12-06

この記事はスタンバイAdvent Calaendar 2021の7日目の記事です。

はじめに

ANNとは

与えられた集合の中で与えられた点に最も近い点を見つけるという最適化問題である最近傍探索(Nearest Neighbor Search)というものがあります。また、近い順にk個の解をみつける場合は、kNN Searchと言います。ANN(Approximate Nearest Neighbor: 近似最近傍探索)とは、kNN Searchの一種で、厳密な最近傍を求めず、ある程度の精度で解を見つけることを許容することで、高速に解を求める手法です。

LuceneによるANN Searchのサポート

最近になってANN SearchがApache Luceneでサポートされました1。
LuceneはANNを今後のリリース予定のバージョン9.0で提供する予定です。

なぜ、全文検索エンジンであるApache LuceneでANN Searchがサポートされるようになったのでしょうか？Luceneのissueには以下のような記載があります1。

"Semantic" search based on machine-learned vector "embeddings" representing terms, queries and documents is becoming a must-have feature for a modern search engine

自然言語処理における、Embedding(埋め込み)とは文章や単語などをベクトル表現に変換することです。具体的にはニューラルネットワークなどの機械学習によりベクトル空間にマッピングします。このようにして得られたベクトルは、単語や文章の意味を反映していることが期待されます。すなわち、似た意味をもつ文章や単語のベクトルは類似度が高くなります。この性質を用いることでSemantic Searchを実現することができます。

EmbeddingによるSemantic SearchとANN Search

Semantic Searchとはユーザの検索クエリの意図や意味を理解した上で結果を提供する検索です。通常の検索では文字列のマッチングで検索結果を返します。この場合、検索対象のドキュメントにクエリで指定された単語が含まれていなければマッチしません。例えば、検索クエリがバイトの場合、同義語辞書などの工夫をしないと似た意味のパートを含むドキュメントにマッチさせることができません。Semantic Searchにより検索クエリの意図を理解できれば、クエリで表現されていないユーザの意図にあった検索結果を返すことができます。

先ほど、Embeddingにより得られたベクトル表現は、似た意味のものであれば類似度が高くなると説明しました。上の例で言えば、意味的に近いバイトとパートはベクトル表現でも近くなると考えられます。そこで、検索対象のドキュメントとクエリをそれぞれベクトル表現にし、その類似度を計算することでクエリとドキュメントが意味的にマッチしているものを取得することができます。

ここでANNの話に戻りますが、ドキュメントのベクトルごとにクエリとの類似度を計算していくと、全てのドキュメントに対して計算をする必要があるため時間がかかります。検索エンジンの多くのユースケースでは検索にかかかる時間は短い方が好ましいです。そのため、高速に検索結果を取得できるようにANNが用いられることになります。

実際に、このような仕組みでSemantic Searchを実現している企業が複数存在し、論文も公開されています23。

Elasticsearchにおける実装

現在Luceneで実装されたANN SearchをElasticsearchでも利用できるように開発が進められています4。Luceneは次期9.0でANNをリリースするため、この機能はElasticsearch 8.xのみが対象となっています。
ANN Searchにはさまざまなアルゴリズムが存在します5が、LuceneではHNSW6というアルゴリズムを使用されているため、ElasticsearchでもHNSWでの実装が進められています。
ANN Searchでは高速化のために、検索に適した形でデータを保持することがあります。HNSWではデータを階層的なグラフとして表現することで高層化を実現しています。

HNSWの詳細や他のアルゴリズムを知りたい方は、以下の資料がわかりやすいです。
https://speakerdeck.com/matsui_528/jin-si-zui-jin-bang-tan-suo-falsezui-qian-xian

ここからは、実際にElasticsearchでANN Searchを行う方法を紹介します。
以下の内容は主にこちらのissueの内容を参考にして作成しました。
https://github.com/elastic/elasticsearch/issues/78473

Elasticsearchの設定をする

はじめに、masterブランチのElasticsearchをビルドします。
2021/11月現在、ElasticsearchのANN Searchはリリース前なので、機能を試したい場合、masterブランチで動かす必要があります。

git clone https://github.com/elastic/elasticsearch.git
cd elasticsearch
./gradlew localDistro
cd distribution/local/elasticsearch-8.1.0-SNAPSHOT/bin
./elasticseaerch

Elasticsearchが起動したら、インデックスの設定をおこないます。
以下のようなHTTPリクエストをElasticsearchに送ることで設定ができます。

PUT http://127.0.0.1:9200/ann_sample
Content-Type: application/json

{
  "mappings": {
    "properties": {
    "text_vector": {
      "type": "dense_vector",
      "dims": 768,
      "index": true,
      "similarity": "l2_norm",
      "index_options": {
        "type" : "hnsw",
        "m" : 15,
        "ef_construction" : 50
      }
    },
    "text" : {
        "type" : "keyword"
      }
    }
  }
}

このインデックスはtextとtext＿vectorという二つのフィールドを持ちます。
textは検索したい文章をそのまま保持するフィールドです。text_vectorは文書ベクトルを保持するフィールドです。

text_vectorフィールドはdense_vectorとして表現されます。
dense_vectorフィールドは、現在のバージョンのElasticsearchで既にリリースされているフィールドです。
しかしながら、これまでのdense_vectorフィールドは、単にベクトルを保存できるだけであり、クエリ、ソートをサポートされておらず、スクリプト内で専用のベクトル関数を使ってのみアクセスが許可されていたため、cosine類似度をスクリプトで計算するといった形で使用することしかできませんでした。
https://www.elastic.co/guide/en/elasticsearch/reference/7.15/dense-vector.html

ANN Searchの機能はdense_vectorの拡張という形で提供されています。
既存のdense_vectorフィールドではdimsというベクトルの次元数を表すパラメータのみが存在していましたが、
ANNをサポートするためにdense_vectorフィールドに新たなパラメータが追加されています。

  • index
    • trueに設定すると、そのフィールドをANNインデックスに追加することになります
    • デフォルトでは互換勢担保のためfalseになっています
  • similarity
    • indexがtrueの場合、ベクトルのインデックスを作成する際に使用する類似性を指定する必要があります。今のところ、l2_normとdot_productが指定できます。これはLuceneのオプションと一致しています。
  • index_options
    • type
      • 使用するkNNアルゴリズム
      • 現状はhnswのみサポート
    • m
      • HNSWグラフにおける各ノードの最大接続数
    • ef_construction
      • HNSW用のパラメータ
      • 新たに挿入された各ノードのグラフを検索する際に追跡する候補数
      • デフォルトは100

データをElasticsearchにインデックスする

データセットのダウンロード

今回はデータセットとしてlivedoorニュースコーパスを使用しました。
livedoorニュースコーパスはlivedoor NEWSに掲載されたニュース記事を収集したもので、株式会社ロンウィットにより公開されています7。

# ダウンロード
wget https://www.rondhuit.com/download/ldcc-20140209.targ.gz

# ファイルの解答
tar -zxf ldcc-20140209.tar.gz

文書ベクトルの生成とインデクシング

今回文章をベクトル表現にするに当たってBERTを使用しました。BERTはGoogleによって提案されたモデルで、さまざまな言語タスクで高い性能を発揮しています8。

まず、BERTを利用するにあたって必要なライブラリをインストールします。

pip install -r requirements.txt
requirements.txt
transformers==4.5.0
fugashi==1.1.0
ipadic==1.0.0
torch==1.8.0
elasticsearch==7.15.2

今回はPyTorchとTransformersを使用します。
PyTorchはFacebookが開発した深層学習のフレームワークで、Transformersはさまざまなニューラルネットワークを用いた言語モデルが実装されているライブラリです。Transformersでは、さまざまな言語の事前学習モデルが利用でき、日本語のモデルもあります。

以下のindexer.pyは文書ベクトルの生成とインデクシングを行うコードです。

indexer.py
import glob

from elasticsearch import Elasticsearch
from elasticsearch.helpers import bulk

import torch
from transformers import BertJapaneseTokenizer, BertModel
from joblib import Parallel, delayed

model_name = 'cl-tohoku/bert-base-japanese-whole-word-masking'
bert = BertModel.from_pretrained(model_name)
tokenizer = BertJapaneseTokenizer.from_pretrained(model_name)
max_length = 256

client = Elasticsearch()
INDEX_NAME = "ann_sample"

def index_batch(doc):
    requests = get_request(doc)
    print(requests)
    bulk(client, requests)

def get_request(doc):
    return [{
        "_op_type": "index",
        "_index": INDEX_NAME,
        "text": doc,
        "text_vector": embedding(doc)
    }]

def embedding(text):

    encoding = tokenizer(
        text,
        max_length=max_length,
        padding='max_length',
        truncation=True,
        return_tensors='pt'
    )
    attention_mask = encoding['attention_mask']

    # 文章ベクトルを計算
    # BERTの最終層の出力の平均を計算する。(ただし、[PAD]は除く)
    with torch.no_grad():
        output = bert(**encoding)
        last_hidden_state = output.last_hidden_state
        averaged_hidden_state = (last_hidden_state * attention_mask.unsqueeze(-1)).sum(1) \
            / attention_mask.sum(1, keepdim=True)

    return averaged_hidden_state[0].tolist()


if __name__ == '__main__':

    BATCH_SIZE = 1000
    docs = []
    count = 0

    category_list = [
        'dokujo-tsushin',
        'it-life-hack',
        'kaden-channel',
        'livedorr-homme',
        'movie-enter',
        'peachy',
        'smax',
        'sports-watch',
        'topic-news'
    ]


    for category in category_list:
        for file in glob.glob(f"../data/text/{category}/{category}*"):
            lines = open(file).read().splitlines()
            print(lines[0])
            text = '\n'.join(lines[3:])
            index_batch(text)

def embedding(text):で文章のベクトル表現を生成しています。

BERTの利用にあたっては、以下の本のコードを参考にさせていただきました。
https://github.com/stockmarkteam/bert-book/blob/master/Chapter10.ipynb

ANNで検索する

最後に、作成したインデックスを使ってANN Searchを行ってみます。
次のようなリクエストを/_knn_searchというエンドポイントに投げることで検索ができます。

GET http://127.0.0.1:9200/ann_sample/_knn_search
Content-Type: application/json

{
    "knn": {
        "field": "text_vector",
        "query_vector": [0.22..., ....]
        "k": 10,
        "num_candidates": 100
    },
    "_source": false,
    "fields": ["text"]
}

各パラメータの説明は以下になります。

  • query_vector
    • クエリのベクトル
  • k
    • 取得する結果の件数
  • num_candidates
    • 各シャード上ごとに候補ベクトルを何件集めるか

kとnum_candidatesの関係ですが、まず結果を収集するために、
各シャード上の近似最近傍候補のベクトルをnum_candidates個見つけます。これらの候補ベクトルとクエリベクトルの類似性を計算し、各シャードからkこのもっとも類似した結果を選択します。最後に各シャードからの結果をマージしてグローバルなトップk個の候補を返却します。

以下のsearcher.pyはann_searchを行うコードです。

searcher.py

def handle_query():
    query = input("Enter query: ")

    # script queryでdense_vectorを使う場合のクエリ
    # script_query = {
    #     "script_score": {
    #         "query": {"match_all": {}},
    #         "script": {
    #             "source":"cosineSimilarity(params.query_vector, doc['text_vector']) + 1.0",
    #             "params":{"query_vector": query_vector}
    #         }
    #     }
    # }

    es_query = {
        "knn": {
            "field": "text_vector",
            "query_vector": embedding(query),
            "k": 10,
            "num_candidates": 100
        },
        "_source": False,
        "fields": ["text"]
    }

    url = 'http://127.0.0.1:9200/livedoor/_knn_search'
    headers = {
        'Content-Type': 'application/json',
    }

    req = urllib.request.Request(url, json.dumps(es_query).encode(), headers)

    with urllib.request.urlopen(req) as res:
        body = res.read().decode()
        print(body)

検索結果のレスポンスのフォーマットは通常の_searchと同じになります。

まとめ

Elasticsearchに新しく追加された ANN Searchを試してみました。
ここで紹介した機能は、まだ開発中であり変わる可能性があるため、ご注意ください。
今回は単純に機能を使ってみるだけでしたが、今後はEmbeddingの方法を変更したり、パラメータを調整したりして精度を改善できないか試してみたいです。また、性能比較をしてANNによりどれだけ高速化されたかをみてみようと思います。

  1. https://issues.apache.org/jira/browse/LUCENE-9004 ↩

  2. https://arxiv.org/abs/2006.11632 ↩

  3. https://arxiv.org/abs/1907.00937 ↩

  4. https://github.com/elastic/elasticsearch/issues/78473 ↩

  5. http://ann-benchmarks.com/ ↩

  6. https://arxiv.org/abs/1603.09320 ↩

  7. https://www.rondhuit.com/download.html ↩

  8. https://arxiv.org/abs/1810.04805 ↩

18
9
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
Sign upLogin
18
9