Hugging FaceのCross EncoderモデルをElasticsearchの機械学習ノードに入れて完全オンプレミスでSemantic Rerankを試す

# はじめに
生成AIを利用したRAG(Retrieval Augmented Generation)には検索部分の精度が非常に重要であり、その方法について以下のブログで説明しました。
https://qiita.com/takeo-furukubo/items/e5d43fa734e4338b895f
検索精度改善にはいろいろな方法がありますが、rerankは優良な方法の一つです。

## Rerank
rerankは一度出た結果を別の方法で並べ替えることです。生成AIが出る前からrerankは使用されています。
![スクリーンショット 2024-09-26 13.03.23.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2603229/7839ea4e-f227-e834-2899-bb3c210ebb36.png)
元々大量のドキュメントから適切なドキュメントを探してくるのは大変なことですし、検索する人や目的、タイミングによって欲しい結果は変わったりします。
なので、まず大量のドキュメントを素早く検索できるBM25（キーワード検索）やANN（ベクトル検索）を利用して絞り込んでから、その結果を並べ替えます。
その方法には色々あります。下に行くほど処理が重くなります。
- Query rescore
    - rerankではなくrescoreという名前になっています。例えば2024年のドキュメントはスコアを２倍にする、等です
    - ドキュメントはこちらです
https://www.elastic.co/guide/en/elasticsearch/reference/current/filter-search-results.html#rescore
- Learning To Rank
    - 検索している人の属性値（例えば年齢）で、同じ属性値の人が多く選んでいるものに高いスコアを与える、というやり方です
    - 弊社コンサルティングの技術者が記載した記事はこちらです
https://qiita.com/daixque/items/38e24704269adc9910b9
- Semantic Rerank
    - 質問文と検索されたスコアの高い上位ドキュメントをそのままrerankerに渡して、再度並べ替える方法です
    - https://www.elastic.co/search-labs/blog/semantic-reranking-with-retrievers

今回の記事は最後のSemantic rerankについて記載しています。

## Elasticで行うメリット
Elasticの中で全て閉じるのでアプリ側での処理がほとんどなくなります。
動きとしては以下のようになります。
1. クエリ文をElasticsearchに発行
2. 検索結果の上位（数は指定可能）を、Inference/rerank APIを使ってクエリ文と一緒に機械学習ノードにあるCross Encoderに送る
3. Cross Encoderが結果を返す
![スクリーンショット 2024-10-04 11.36.03.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2603229/69a37efe-2a65-a1f5-77d1-9c716ceb199a.png)


semantic rerank自体は多くの検索エンジンでサポートされていますが、全てを検索エンジン内部に置けるのはElasticsearchだけです。
現在は日本語対応のCross Encoderは少ないですが、Cross EncoderをElasticの機械学習ノードに配置することで全ての処理（初段の検索からsemantic rerankまで）をElasticsearchに閉じることができます。

本ブログでは全てをオンプレミスで行います。

# 動作環境
これが一番大変かも知れません。Cross Encoderは生半可なメモリでは動きません。Elastic Cloudのフリートライアルでは機械学習ノードのRAMを4GBまでしか設定できないので、Cross Encoderが動作しません。（有償版ではもちろん問題ないです）
今回試したのは以下の環境です。
- VM
Azure Standard D4s v3 (4 vcpu 数、16 GiB メモリ)
- OS
Ubuntu 20.04.6 LTS

# 手順
:::note info
2024/10/04現在の最新バージョンである8.15.2をベースに記載
:::
## ソースコード
以下git cloneします。
https://github.com/legacyworld/semantic_reranking

## Elasticsearchクラスタ
ベースとなる`docker-compose.yml`はこちらです。
https://github.com/elastic/elasticsearch/blob/8.15/docs/reference/setup/install/docker/docker-compose.yml
1ノードで全て用意するのが楽なのですが、普通データノードと機械学習ノードは混在させないので、分ける構成を取っています。
### .env
`.env`ファイルに必要な情報を記入します。
passwordはご自由に変更してください。
`MEM_LIMIT`がデータノード（2台）とKibanaのメモリ量です。
`ML_MEM_LIMIT`が機械学習ノードのメモリ量です。
もう少し少なくても動くかも知れません。
```
STACK_VERSION = 8.15.2
ELASTIC_PASSWORD = elastic
KIBANA_PASSWORD = elastic
ES_PORT = 9200
CLUSTER_NAME = test
LICENSE = trial
MEM_LIMIT = 1610612736
ML_MEM_LIMIT = 8589934592
KIBANA_PORT = 5601
```
### 起動
```
docker compose up -d
```
一部buildを行います。
立ち上げ時に証明書を作ったりしますので結構時間かかります。
全部で5コンテナ立ち上がります。
- es01/02
    - データノード
- es03
    - 機械学習ノード
- kibana
- python
    - Pythonプログラム実行用 & Cross Encoderアップロード用
    - https://github.com/elastic/eland
    - elandは専用のコンテナがあるが、8.15.2からセキュリティが厳しくなりvolumeのマウントが現在出来ないため使用していない
```
docker compose ps
NAME                          IMAGE                                                  COMMAND                  SERVICE   CREATED         STATUS                   PORTS
semantic_reranking-es01-1     docker.elastic.co/elasticsearch/elasticsearch:8.15.2   "/bin/tini -- /usr/l…"   es01      7 minutes ago   Up 7 minutes (healthy)   0.0.0.0:9200->9200/tcp, :::9200->9200/tcp, 9300/tcp
semantic_reranking-es02-1     docker.elastic.co/elasticsearch/elasticsearch:8.15.2   "/bin/tini -- /usr/l…"   es02      7 minutes ago   Up 7 minutes (healthy)   9200/tcp, 9300/tcp
semantic_reranking-es03-1     docker.elastic.co/elasticsearch/elasticsearch:8.15.2   "/bin/tini -- /usr/l…"   es03      7 minutes ago   Up 7 minutes (healthy)   9200/tcp, 9300/tcp
semantic_reranking-kibana-1   docker.elastic.co/kibana/kibana:8.15.2                 "/bin/tini -- /usr/l…"   kibana    7 minutes ago   Up 6 minutes (healthy)   0.0.0.0:5601->5601/tcp, :::5601->5601/tcp
semantic_reranking-python-1   semantic_reranking-python                              "/bin/sh -c 'while :…"   python    7 minutes ago   Up 5 minutes    
```

## Cross Encoderモデルアップロード
今回Cross Encoderモデルとしてこちらを利用します。
https://huggingface.co/hotchpotch/japanese-reranker-cross-encoder-xsmall-v1
以下のコマンドを実行します
(`.env`でパスワードを変えている場合は`-p`の後ろを変えてください)
```
docker exec -it semantic_reranking-python-1 \
eland_import_hub_model \
--url https://es01:9200 \
-u elastic -p elastic \
--hub-model-id hotchpotch/japanese-reranker-cross-encoder-xsmall-v1 \
--task-type text_similarity \
--ca-certs /config/ca/ca.crt \
--max-model-input-length 512 \
--start
```
こんな感じの実行結果になります。
```
tokenizer_config.json: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1.20k/1.20k [00:00<00:00, 6.44MB/s]
sentencepiece.bpe.model: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5.07M/5.07M [00:00<00:00, 79.2MB/s]
special_tokens_map.json: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 964/964 [00:00<00:00, 5.49MB/s]
tokenizer.json: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 17.1M/17.1M [00:00<00:00, 250MB/s]
config.json: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 786/786 [00:00<00:00, 4.52MB/s]
model.safetensors: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 428M/428M [00:01<00:00, 257MB/s]
STAGE:2024-10-02 06:32:26 14:14 ActivityProfilerController.cpp:312] Completed Stage: Warm Up
STAGE:2024-10-02 06:32:26 14:14 ActivityProfilerController.cpp:318] Completed Stage: Collection
STAGE:2024-10-02 06:32:26 14:14 ActivityProfilerController.cpp:322] Completed Stage: Post Processing
Asking to pad to max_length but no maximum length is provided and the model has no predefined maximum length. Default to no padding.
2024-10-02 06:32:28,641 INFO : Creating model with id 'hotchpotch__japanese-reranker-cross-encoder-xsmall-v1'
2024-10-02 06:32:30,905 INFO : Uploading model definition
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 409/409 [01:05<00:00,  6.27 parts/s]
2024-10-02 06:33:36,185 INFO : Uploading model vocabulary
2024-10-02 06:33:38,495 INFO : Starting model deployment
2024-10-02 06:33:50,802 INFO : Model successfully imported with id 'hotchpotch__japanese-reranker-cross-encoder-xsmall-v1'
```
`eland`でのアップロードについてのドキュメントはこちら
https://www.elastic.co/guide/en/machine-learning/current/ml-nlp-text-emb-vector-search-example.html
### アップロードされたモデルをKibanaで確認
左上の三本線->Machine learning
![スクリーンショット 2024-10-02 15.41.33.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2603229/b3960248-ae76-2166-4dfa-32a4f02aed1a.png)
Trained Modelをクリック後に、`ML job and trained model synchronization required
Some jobs or trained models are missing or have incomplete saved objects. Synchronize your jobs and trained models.`と上部に表示されているので、`Synchronize your jobs and trained models.`をクリックする。
![スクリーンショット 2024-10-02 15.39.23.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2603229/c587044d-a0cf-24be-7868-1d256784b46e.png)
右側に確認画面が出るので`Synchronize`をクリック
![スクリーンショット 2024-10-02 15.44.57.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2603229/db568339-5b82-4761-6c0c-75676e635df1.png)
アップロードしたモデルが表示されています。
![スクリーンショット 2024-10-02 15.46.36.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/2603229/312b0a52-19f0-ee45-41bf-047694961969.png)

## Inference Endpoint作成
`inference/rerank`を使ってCross Encoderに検索結果の上位を送付できるようにします。
以下のコマンドを実行します。
```
docker exec -it semantic_reranking-python-1 curl --cacert /config/ca/ca.crt -u elastic:elastic -X PUT "https://es01:9200/_inference/rerank/my-rerank" -H 'Content-Type: application/json' -d'
{
  "service": "elasticsearch",
  "service_settings": {
    "num_allocations": 1,
    "num_threads": 1,
    "model_id": "hotchpotch__japanese-reranker-cross-encoder-xsmall-v1" 
  }
}
'
```
少し時間がかかりますが(十秒ぐらい)、以下のように返ってくればOKです。
```
{"inference_id":"my-rerank","task_type":"rerank","service":"elasticsearch","service_settings":{"num_allocations":1,"num_threads":1,"model_id":"hotchpotch__japanese-reranker-cross-encoder-xsmall-v1"},"task_settings":{"return_documents":true}}
```

## データ投入
rerankerの効果を見るために以下のようなデータを用意しています。皆さんの方で適宜追加・変更してご利用ください。
```json
[
    {"description": "京都府の観光名所はたくさんありますが、金閣寺や清水寺が有名です"},
    {"description": "東京都の観光名所はたくさんありますが、浅草寺や東京スカイツリーが有名です"},
    {"description": "東京都について教えて下さい"},
    {"description": "京橋という地名は東京都にも大阪府にもあります"},
    {"description": "京都には京橋という観光名所があります"},
    {"description": "大阪では２０２５年に万博が開かれます"},
    {"description": "東京オリンピックは2021年に開かれました"},
    {"description": "兵庫県の観光名所は、姫路城や有馬温泉が有名です"}
]
```
以下のコマンドで投入が簡単に行えます。
```
docker exec -it semantic_reranking-python-1 python /src/ingest.py --index_name=rerank --file=/src/data.json
```
以下のようになればOKです
```
Indexing documents, this might take a while...
100%|█████████████████████████████████████████████████████████████████████████████████████████| 6/6 [00:00<00:00, 10.44documents/s, success=6]
Indexing completed! Success percentage: 100.0%
Done indexing documents!
```
以下のコマンドで確認できます。
```
docker exec -it semantic_reranking-python-1  curl --cacert /config/ca/ca.crt -u elastic:elastic https://es01:9200/_search
```
以下のようになっていればOKです。
```
{"took":5,"timed_out":false,"_shards":{"total":15,"successful":15,"skipped":0,"failed":0},"hits":{"total":{"value":6,"relation":"eq"},"max_score":1.0,"hits":[{"_index":"rerank","_id":"B9sFUJIB4uj14NN9DtZd","_score":1.0,"_source":{"description":"京都府の観光名所はたくさんありますが、金閣寺や清水寺が有名です"}},{"_index":"rerank","_id":"CNsFUJIB4uj14NN9DtZd","_score":1.0,"_source":{"description":"東京都の観光名所はたくさんありますが、浅草寺や東京スカイツリーが有名です"}},{"_index":"rerank","_id":"CdsFUJIB4uj14NN9DtZd","_score":1.0,"_source":{"description":"東京都の観光名所を教えて下さい"}},{"_index":"rerank","_id":"CtsFUJIB4uj14NN9DtZd","_score":1.0,"_source":{"description":"京橋という地名は東京都にも大阪府にもあります"}},{"_index":"rerank","_id":"C9sFUJIB4uj14NN9DtZd","_score":1.0,"_source":{"description":"京都には京橋という観光名所があります"}},{"_index":"rerank","_id":"DNsFUJIB4uj14NN9DtZd","_score":1.0,"_source":{"description":"大阪では２０２５年に万博が開かれます"}}]}}
```

## 検索
投入したドキュメントに対して「京都の観光名所を教えて」という検索文を入れてみます。
本来は「京都府の観光名所はたくさんありますが、金閣寺や清水寺が有名です」が一番上に来てほしいところですが、そうはならないようにわざとしています。

これは別に、キーワード検索がダメ、というわけではありません。わざと形態素解析も一切設定せずにノイズが多く乗るN-Gramのみで行うようにしているからです。
「東京都」の後ろの「京都」にひっかかりますし、「東京スカイツリー」の「京」もひっかかります。
形態素解析をすればこれはありえません。

では実際に実行してみます。
```
 docker exec -it semantic_reranking-python-1 python /src/rerank.py rerank 京都の観光名所を教えて
 ```
 結果は以下のようになります。
 Semantic Rerankにある小数点の数字はRerankのスコアです。
```
質問文：京都の観光名所を教えて
キーワード検索のみ
1 東京都について教えて下さい
2 東京都の観光名所はたくさんありますが、浅草寺や東京スカイツリーが有名です
3 京都府の観光名所はたくさんありますが、金閣寺や清水寺が有名です
4 京都には京橋という観光名所があります
5 兵庫県の観光名所は、姫路城や有馬温泉が有名です
6 京橋という地名は東京都にも大阪府にもあります
7 東京オリンピックは2021年に開かれました

キーワード検索＋Semantic Rerank
1 1.1773956 京都府の観光名所はたくさんありますが、金閣寺や清水寺が有名です
2 1.0673646 京都には京橋という観光名所があります
3 -0.26459968 兵庫県の観光名所は、姫路城や有馬温泉が有名です
4 -0.28720367 東京都の観光名所はたくさんありますが、浅草寺や東京スカイツリーが有名です
5 -0.31205043 京橋という地名は東京都にも大阪府にもあります
6 -0.5200176 東京都について教えて下さい
7 -0.5813435 東京オリンピックは2021年に開かれました
```
キーワード検索のみの場合、東京都についてのドキュメントが上に来てしまっています。これは「東京都」の後ろの「京都」に引っかかっているからです。
他にも「京橋」も「京」が引っかかっています。

ところが、Semantic Rerankを通すことで、ほしい2つの情報だけが高いスコアで上位に来ています。プログラムでは`min_score`をコメントアウトしていますが、これをある程度の正の値にしておけば、全く関係ないものは排除可能です。

### Rerankを行うプログラム
重要な部分を抜粋します。キーワード検索のみを行う場合と大きく変わらないのが見ていただけると思います。
`retriever`を利用することで、非常に簡潔に記述することが可能になっています。
ドキュメントはこちらです。
https://www.elastic.co/guide/en/elasticsearch/reference/current/retriever.html

```python
# キーワード検索のみ
query_body = {
   "query": {
      "match": {
         "description": query
      }
   }
}

result = es.search(index=search_index, body=query_body)

# キーワード検索＋Semantic Rerank
print("キーワード検索＋Semantic Rerank")
query_body = {
       "retriever": {
      "text_similarity_reranker": {
         "retriever": {
            "standard": {
               "query": {
                  "match": {
                     "description": query
                  }
               }
            }
         },
         "field": "description",
         "inference_id": "my-rerank",
         "inference_text": query,
         "rank_window_size": 10
#         "min_score": 0.5
      }
   }
}

result = es.search(index=search_index, body=query_body)
```

## 形態素解析もやってみる
形態素解析をするとキーワード検索で大体正しくなるのでRerankする意味があまり見出せませんが、一応記載しておきます。

### プラグイン導入
`analysis-icu`と`analysis-kuromoji`を3台のelastcsearchノードにインストールして再起動が必要です。
プラグインインストールは下記のコマンドで行います。
```
 docker exec -it semantic_reranking-es01-1 /usr/share/elasticsearch/bin/elasticsearch-plugin install --batch analysis-kuromoji analysis-icu
 ```
 以下のように表示されて再起動が必要であることが分かります。
 ```
 -> Installing analysis-kuromoji
-> Downloading analysis-kuromoji from elastic
-> Installed analysis-kuromoji
-> Installing analysis-icu
-> Downloading analysis-icu from elastic
-> Installed analysis-icu
-> Please restart Elasticsearch to activate any plugins installed
```
以下のコマンドで再起動を行います。
```
docker compose restart es01
```
再起動を立て続けに行うと、複数のmasterノードがいなくなることになります。
以下のコマンドでhealthがgreenになっていることを確認してから順次es02,es03に同様にプラグインをインストールして再起動します。
```
docker exec -it semantic_reranking-python-1 curl --cacert /config/ca/ca.crt -u elastic:elastic https://es01:9200/_cat/health?v=true
```
結果
```
epoch      timestamp cluster status node.total node.data shards pri relo init unassign pending_tasks max_task_wait_time active_shards_percent
1727942863 08:07:43  mac     green           3         2     78  39    0    0        0             0                  -                100.0%
```

### Index Settings
以下を実行します。
```
docker exec -it semantic_reranking-python-1 curl --cacert /config/ca/ca.crt -u elastic:elastic -X PUT "https://es01:9200/rerank_jp" -H 'Content-Type: application/json' -d'
{
  "settings": {
    "analysis": {
      "char_filter": {
        "normalize": {
          "type": "icu_normalizer",
          "name": "nfkc",
          "mode": "compose"
        }
      },
      "tokenizer": {
        "ja_kuromoji_tokenizer": {
          "mode": "search",
          "type": "kuromoji_tokenizer"
        }
      },
      "analyzer": {
        "kuromoji_analyzer": {
          "tokenizer": "ja_kuromoji_tokenizer",
          "char_filter": ["normalize"],
          "filter": [
            "kuromoji_baseform",
            "kuromoji_part_of_speech",
            "cjk_width",
            "ja_stop",
            "kuromoji_stemmer",
            "lowercase"
          ]
        }
      }
    }
  },
  "mappings": {
    "properties": {
      "description": {
        "type": "text",
        "analyzer": "kuromoji_analyzer"
      }
    }
  }
}
'
```
以下が表示されればOKです。
```
{"acknowledged":true,"shards_acknowledged":true,"index":"rerank_jp"}
```

### 形態素解析ありでデータ投入
以下のコマンドを実行します。
```
docker exec -it semantic_reranking-python-1 python /src/ingest.py --index_name=rerank_jp --file=/src/data.json
```

### 形態素解析ありでデータを検索
以下のコマンドを実行します。「京都」「観光名所」「教えて」が入っていないドキュメントは排除されています。「東京都について教えて下さい」が上位に来ているのは、BM25では短い文章でヒットするとスコアが高くなるからです。
```
docker exec -it semantic_reranking-python-1 python /src/rerank.py rerank_jp 京都の観光名所を教えて
```
結果はこうなります。
```
質問文：京都の観光名所を教えて
キーワード検索のみ
1 京都には京橋という観光名所があります
2 京都府の観光名所はたくさんありますが、金閣寺や清水寺が有名です
3 東京都について教えて下さい
4 兵庫県の観光名所は、姫路城や有馬温泉が有名です
5 東京都の観光名所はたくさんありますが、浅草寺や東京スカイツリーが有名です

キーワード検索＋Semantic Rerank
1 1.1773956 京都府の観光名所はたくさんありますが、金閣寺や清水寺が有名です
2 1.0673646 京都には京橋という観光名所があります
3 -0.26459968 兵庫県の観光名所は、姫路城や有馬温泉が有名です
4 -0.28720367 東京都の観光名所はたくさんありますが、浅草寺や東京スカイツリーが有名です
5 -0.5200176 東京都について教えて下さい
```
これであればRerankは必要ないと思いますが、当然現実の検索ではこのような単純なケースはないと思います。

# まとめ
今回はSemantic Rerankという検索精度を向上させる方法をご紹介しました。
キーワード検索やANNのようなベクトル検索は非常に早く大容量のデータを検索できますが、当然ノイズが乗ってきます。
それに対して、質問文とドキュメントを同時に処理するSemantic Rerankは有効な方法ですが、非常に処理が重いです。全てのドキュメントをRerankerに投げる、という方法を取ることは出来ません。
それぞれにメリット・デメリットがあり、うまく使い分けることが重要です。

Elasticsearchではどのような構成でも対応が可能です。オンプレ・クラウド・ハイブリッド、全てに対応できるのはElasticsearchだけです。