ElasticsearchのANN Searchを試す

この記事は[スタンバイAdvent Calaendar 2021](https://qiita.com/advent-calendar/2021/stanby)の7日目の記事です。
## はじめに
### ANNとは
与えられた集合の中で与えられた点に最も近い点を見つけるという最適化問題である最近傍探索(Nearest Neighbor Search)というものがあります。また、近い順にk個の解をみつける場合は、kNN Searchと言います。ANN(Approximate Nearest Neighbor: 近似最近傍探索)とは、kNN Searchの一種で、厳密な最近傍を求めず、ある程度の精度で解を見つけることを許容することで、高速に解を求める手法です。

### LuceneによるANN Searchのサポート
最近になってANN SearchがApache Luceneでサポートされました[^1]。
[^1]: https://issues.apache.org/jira/browse/LUCENE-9004
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を実現している企業が複数存在し、論文も公開されています[^2][^3]。

[^2]: https://arxiv.org/abs/2006.11632
[^3]: https://arxiv.org/abs/1907.00937


### Elasticsearchにおける実装

現在Luceneで実装されたANN SearchをElasticsearchでも利用できるように開発が進められています[^4]。Luceneは次期9.0でANNをリリースするため、この機能はElasticsearch 8.xのみが対象となっています。
[^4]: https://github.com/elastic/elasticsearch/issues/78473

ANN Searchにはさまざまなアルゴリズムが存在します[^5]が、Luceneでは`HNSW`[^6]というアルゴリズムを使用されているため、Elasticsearchでも`HNSW`での実装が進められています。
[^5]: http://ann-benchmarks.com/
[^6]: https://arxiv.org/abs/1603.09320

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ブランチで動かす必要があります。

```shell
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に送ることで設定ができます。

```shell
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]。

[^8]: https://arxiv.org/abs/1810.04805

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

```bash
pip install -r requirements.txt
```

```text: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`は文書ベクトルの生成とインデクシングを行うコードです。

```python: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

https://www.amazon.co.jp/BERT%E3%81%AB%E3%82%88%E3%82%8B%E8%87%AA%E7%84%B6%E8%A8%80%E8%AA%9E%E5%87%A6%E7%90%86%E5%85%A5%E9%96%80-Transformers%E3%82%92%E4%BD%BF%E3%81%A3%E3%81%9F%E5%AE%9F%E8%B7%B5%E3%83%97%E3%83%AD%E3%82%B0%E3%83%A9%E3%83%9F%E3%83%B3%E3%82%B0-%E3%82%B9%E3%83%88%E3%83%83%E3%82%AF%E3%83%9E%E3%83%BC%E3%82%AF%E6%A0%AA%E5%BC%8F%E4%BC%9A%E7%A4%BE/dp/427422726X/ref=asc_df_427422726X/?tag=jpgo-22&linkCode=df0&hvadid=310844024129&hvpos=&hvnetw=g&hvrand=14030705246167011065&hvpone=&hvptwo=&hvqmt=&hvdev=c&hvdvcmdl=&hvlocint=&hvlocphy=1009308&hvtargid=pla-1265639589305&psc=1&th=1&psc=1



## ANNで検索する

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

```shell
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を行うコードです。

```python: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によりどれだけ高速化されたかをみてみようと思います。



[^7]: https://www.rondhuit.com/download.html