はじめに

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

Rerank

rerankは一度出た結果を別の方法で並べ替えることです。生成AIが出る前からrerankは使用されています。

元々大量のドキュメントから適切なドキュメントを探してくるのは大変なことですし、検索する人や目的、タイミングによって欲しい結果は変わったりします。
なので、まず大量のドキュメントを素早く検索できる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の中で全て閉じるのでアプリ側での処理がほとんどなくなります。
動きとしては以下のようになります。

クエリ文をElasticsearchに発行
検索結果の上位（数は指定可能）を、Inference/rerank APIを使ってクエリ文と一緒に機械学習ノードにあるCross Encoderに送る
Cross Encoderが結果を返す

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

手順

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

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.をクリックする。

右側に確認画面が出るのでSynchronizeをクリック

アップロードしたモデルが表示されています。

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の効果を見るために以下のようなデータを用意しています。皆さんの方で適宜追加・変更してご利用ください。

[
    {"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

# キーワード検索のみ
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だけです。