投稿者: shinsuke

Elasticsearchでknnクエリーを利用して、ベクトル検索を試すことができるの試してみる。まず、Elasticsearchを起動する。

git clone https://github.com/codelibs/docker-elasticsearch.git
cd docker-elasticsearch
docker compose -f compose.yaml up -d

という感じで、Dockerがあれば、http://localhost:9200でアクセスできるElasticsearchが起動できる。

ML機能はライセンスが必要なので、以下でトライアルを有効にする。

curl -XPOST -H "Content-Type:application/json" "http://localhost:9200/_license/start_trial?acknowledge=true"

ElasticsearchのML機能を利用して、テキストデータを渡されたら、ベクトルに変換してくれるようにingestを設定して、ingestでテキストからベクトルを推論するモデルを登録していく。

まずは、モデルを登録していくのが、モデルをインポートするためには、Pythonのelandを用いて、登録することになる。そのために、Pythonモジュールをインストールする。今回は、日本語を利用するので、以下のようなモジュールを入れておく。

pip install 'eland[pytorch]'
pip install fugashi
pip install unidic-lite

次はモデルのインポート。今回はtohoku/bert-base-japanese-v2を指定する。（今回は、簡単に試せるように、httpsなどのsecurityを無効な状態にしてあるので、ユーザー名などの指定が不要である）

eland_import_hub_model --url http://localhost:9200 --hub-model-id cl-tohoku/bert-base-japanese-v2 --task-type text_embedding

インポートできたら、モデルをデプロイする。

curl -XPOST -H "Content-Type:application/json" "http://localhost:9200/_ml/trained_models/cl-tohoku__bert-base-japanese-v2/deployment/_start"

ingestを登録して、テキストがベクトルに変換できるようにする。

curl -XPUT -H "Content-Type:application/json" "http://localhost:9200/_ingest/pipeline/text_embedding" -d '
{
  "description": "Text embedding pipeline",
  "processors": [
    {
      "inference": {
        "model_id": "cl-tohoku__bert-base-japanese-v2",
        "target_field": "content_vector",
        "field_map": {
          "content": "text_field"
        }
      }
    }
  ]
}
'

contentフィールドにテキスト情報を入れる。ベクトルは、content_vectorオブジェクトに入る。ベクトル自体はcontent_vector.predicted_valueに格納される。

ここまでくれば、下準備ができたので、次にデータを入れるインデックスを作る。

curl -XPUT -H "Content-Type:application/json" "http://localhost:9200/docs" -d '
{
  "settings": {
    "index": {
      "number_of_shards": 1,
      "number_of_replicas": 0
    }
  },
  "mappings": {
    "properties": {
      "content": {
        "type": "text",
        "analyzer": "standard"
      },
      "content_vector.model_id": {
        "type": "keyword"
      },
      "content_vector.predicted_value": {
        "type": "dense_vector",
        "dims": 768,
        "index": true,
        "similarity": "l2_norm"
      }
    }
  }
}
'

検証なので、contentにデータを入れるシンプルなインデックスです。このインデックスにデータを入れrます。テスト用のデータは、ChatGPTに10件生成してもらったので、それを入れます。

curl -XPOST -H "Content-Type: application/x-ndjson" "http://localhost:9200/_bulk?pipeline=text_embedding" -d '
{"index": {"_index": "docs"}}
{"content": "明日の天気は晴れの予報です。"}
{"index": {"_index": "docs"}}
{"content": "彼女が好きな色はピンクです。"}
{"index": {"_index": "docs"}}
{"content": "日本語を勉強しています。"}
{"index": {"_index": "docs"}}
{"content": "新しいレストランがオープンしました。"}
{"index": {"_index": "docs"}}
{"content": "今日は疲れたので早く寝ます。"}
{"index": {"_index": "docs"}}
{"content": "おいしいコーヒーを飲みました。"}
{"index": {"_index": "docs"}}
{"content": "新しいスマートフォンが発売されます。"}
{"index": {"_index": "docs"}}
{"content": "友達と映画を見に行きました。"}
{"index": {"_index": "docs"}}
{"content": "最近、本を読む習慣を始めました。"}
{"index": {"_index": "docs"}}
{"content": "子供たちが公園で遊んでいます。"}
'

シンプルなデータなので、面白みにかけますが、検索は以下のリクエストを投げると、結果が帰ってきます。

curl -XPOST -H "Content-Type: application/json" "http://localhost:9200/docs/_search" -d '
{
  "knn": {
    "field": "content_vector.predicted_value",
    "query_vector_builder": {
      "text_embedding": {
        "model_id": "cl-tohoku__bert-base-japanese-v2",
        "model_text": "彼女が好きな色は？"
      }
    },
    "k": 3,
    "num_candidates": 10
  },
  "_source": [
    "content"
  ]
}
'

「彼女が好きな色はピンクです。」が1件目で取得できると思いますが、2件目以降は類似度が高い順に返ってくるだけです。

という感じですが、詳しくはElasticsearchのドキュメントを参照するのが良いと思います。