ときどき起きる

だいたい寝ている

ElasticsearchのIngest Pipelineでtext embeddingを埋め込む & サクッとKNN+BM25のHybrid Searchを試せるリポジトリを作った

Elasticsearch

本記事は情報検索・検索技術 Advent Calendar 2022の4日目の記事です。

こんにちは、pakioです。

先日のElasticON Tokyoに参加した際、とても興味深いセッションがありました。

The search for relevance with Vector Search

内容としては以下のブログと同じかと思います。 www.elastic.co

ざっくり説明するとElasticsearch + Ingest Pipelineを使えば自前でMLモデルから特徴量を抽出するようなサービスを立ち上げる必要なく、ドキュメントにembeddingを埋め込めるよと言った内容の講演でした。 かつ、Ingest Pipelineを利用することで、リアルタイム更新にも対応しているという優れものです。これは試してみるしかと思い、今回はその検証を行ったリポジトリを公開・及び主要なポイントとハマりどころを残します。

ElasticsearchのKNN/ANNに関してはhurutoriyaさんのこちらのブログ1やfujimotoshinjiさんのこちらのzennの記事2等類似した内容の記事も数多くありますが、MLノードを用いて、かつインデキシング時にリアルタイムで埋め込む事例、及びready-to-useなデモのコードに関してはあまり公開されていない為そのあたりを重点的に記載しています。

リポジトリ

github.com

デモ画面

上記リポジトリの手順通りセットアップすると、以下のようなデモ画面が表示されます。

デモ画面

データはwikimediaからenwikiの一部を利用しています。
wikimediaからのデータ投入に関しては、さっとさんのzenn3が大変参考になりました。元記事では日本語版であるjawikiでのデータ投入方法及びmapping設定についても言及されていますので、日本語で何かしらを試される場合のテストデータ作成にも大いに役立つ記事かと思います。

モデルはHugging Faceからsentence-transformers/msmarco-MiniLM-L-12-v3を利用し、出力は384次元のベクトルとなっています。

手順

1. Elasticsearchを起動 / モデルのアップロード

Elasticsearchを起動

/Esに移動し、Elasticsearch及びKibana(オプション)を起動します。

docker-compose up -d

Ingest Pipeline内でモデルを走らせる方法では、MLノードを有効化させる必要があるため、オプションとしてnode.roles=ml,ingestを指定した状態で起動しています。yaml

また、現時点でMLノードはBasic Licenseではサポートされていないため、検証用にTrial Licenseを有効化します。

curl -X POST 'http://127.0.0.1:9200/_license/start_trial?acknowledge=true'

text embedding生成モデルのアップロード

Eland4を用いて、起動したElasticsearchにモデルをアップロードします。 ElasdではHugging FaceのモデルIDをそのまま指定しアップロードが可能なため、ファインチューニングが必要ない場合には複雑な手順を踏むことなく検証が可能です。

eland_import_hub_model --url http://127.0.0.1:9200 --hub-model-id sentence-transformers/msmarco-MiniLM-L-12-v3 --task-type text_embedding

また、アップロードが完了した時点ではモデルは有効化されていないため、_startを叩いて有効化してあげます。

curl -X POST 'http://127.0.0.1:9200/_ml/trained_models/sentence-transformers__msmarco-minilm-l-12-v3/deployment/_start'

Index/Ingest Pipeline定義

次に、Index及びIngest Pipelineを定義します。

ここでまず1つ目のハマりポイントですが、Ingest Pipelineから出力されたベクトルは、指定したフィールドそのものではなく、一つ下の.predicted_valueに格納されます。Ingest Pipelineで指定したtarget_fieldではないためご注意ください。

2つ目のハマりポイントは、Ingest Pipelineの入力に使われるフィールドの指定についてです。テキストからembeddingを生成するため当然その入力となるフィールドが必要となりますが、デフォルトのフィールド名はtext_fieldとなっています。その値自体は変更不可なため、任意のフィールドを入力に利用したい場合、フィールド名とのマッピングを指定します。

{
  "settings": {
    "number_of_shards": 1,
    "number_of_replicas": 0
  },
  "mappings": {
    "properties": {
      "text": {
        "type": "text"
      },
      "url": {
        "type": "keyword"
      },
      "id": {
        "type": "keyword"
      },
      "title": {
        "type": "text"
      },
      "_vector.predicted_value": {  // ポイント①: .predicted_value必須 
        "type": "dense_vector",
        "dims": 384,
        "index": true,
        "similarity": "cosine"
      }
    }
  }
}

{
  "processors": [
    {
      "inference": {
        "model_id": "sentence-transformers__msmarco-minilm-l-12-v3",
        "target_field": "_vector",
        // ポイント②: "マッピングしたいフィールド名": "text_field"を指定する
        "field_map": {
          "title": "text_field"
        }
      }
    }
  ]
}

以上の手順でElasticsearch本体側の設定は完了です。次は実際にデータを投入します。

データの投入

1. wikimediaからデータを取得、加工

次に、/indexerに移動し、wikimediaからデータをダウンロード、インデキシング可能な形式に加工します。

加工手順の詳細については、以下の記事と同様の手順ですので、こちらを御覧ください。 zenn.dev

wget https://dumps.wikimedia.org/enwiki/20221201/enwiki-20221201-pages-articles-multistream1.xml-p1p41242.bz2

python wikiextractor/WikiExtractor.py -o output -b 10M enwiki-20221201-pages-articles-multistream1.xml-p1p41242.bz2 --json

ls ./output/AA/* -d | xargs -L 1 -P 10 bash -c './reformat_to_ndjson.py $0'

2. _bulkでデータを投入

加工後のデータを_bulk APIを用いて投入していきます。

ls ./output/AA/*_new.ndjson -d | xargs -L 1 bash -c 'echo $0 ; cat $0 | curl -s -X POST -H '\''Content-Type: application/x-ndjson'\'' '\''http://127.0.0.1:9200/wiki/_bulk?pipeline=embedding-pipeline&pretty'\'' --data-binary @-;'

ここで3つ目のハマりポイントですが、リクエストのパラメータにpipeline=embedding-pipelineを指定します。このパラメータを指定する事により、投入されたデータが即座にキューに追加され、随時処理・先程定義したパイプラインによってリアルタイムにembeddingが付与されます。

また、このコマンドを実行する際に429のステータスが返ってきた場合、処理が追いついていないことを示します。1.で WikiExtractor.pyの引数に指定した-b 10Mを小さくしてバッチサイズを減らすか、またはデータ投入時に適当なwaitを挟んで上げることで対処可能です。

ここまでの手順でデータの投入は完了です。

遊んで見る

1. 検証ツールの起動

最後に、/evalに移動し、検証ツールを起動します。ツールにはstreamlitを使用しているため、インストールがまだの方はpip install streamlitでインストールしてください。

streamlit run main.py

query部分に文字を入力すると先程投入したデータに対しての検索処理が開始されます。
また、画面上部のスライドバーを調整することで、クエリウェイトの調整も可能です。

最後に

今回のブログではElasticsearchのIngest Pipelineを用いて、MLモデルの学習/ホスティングなしにHybrid Searchを試してみました。 デモのgifにもあるように、口語文でも比較的関連性の高い結果を検索できる為大変体験がよく、今後のベクトル検索に対する個人的な期待がかなり高まる結果となりました。

  1. Elasticsearchの近似近傍探索を使って、ドラえもんのひみつ道具検索エンジンを作ってみた↩
  2. Elasticsearch の Machine Learning 機能を使ったベクトル検索↩
  3. Docker+ElasticsearchにWikipediaの情報をインデクシングする↩
  4. https://github.com/elastic/elandhttps://github.com/elastic/eland↩