fujimotoshinji 2022/05/24

前回書いた knn Search の記事ではテキストのベクトル化を Elasticsearch 外部で実施しました。

Elasticsearch の Machine Learning 機能でベクトル化も含めて Elasticsearch でできるみたいで公式ブログに手順がまとめられていたので試してみました。なお、Machine Learning 機能は Platinum サブスクリプション以上に含まれる機能となりますのでご注意ください。

手順としては以下のとおりとなります。

環境準備
テキスト埋め込みモデルのデプロイ
テストデータの準備
Ingest パイプラインの作成
Reindex によるテストデータのベクトル化フィールド生成
ベクトル検索

fujimotoshinji 2022/05/24に更新

概要

このブログではテキスト埋め込みモデルを利用してテキストからベクトル表現を生成し、インデックスされたベクトル表現から検索キーワードのベクトル表現に類似するドキュメントを抽出します。
Elasticsearch Machine Learning の機能を利用することで Ingest パイプライン内でテキストからベクトル表現を生成し、ベクトル表現をフィールドに追加した状態でインデキシングできます。

サンプルのデータセットには 2019 TREC Deep Learning Track のテストステージのデータセットを用います。
テキスト埋め込みモデルには Sentence-Transformers を利用します。

fujimotoshinji 2022/05/24に更新

1. 環境準備

動作確認

環境 : macOS
Elastic Stack バージョン : 8.2.0

eland クライアント

Elasticsearch Machine Learning の操作に eland を利用します。
macOS でインストールしようと思ったら、ライブラリ関連でつまずいたので Docker 環境を用意しました。
Docker イメージは GitHub から Clone して docker build するだけでした。

$ git clone https://github.com/elastic/eland.git
$ cd eland
$ docker build -t eland .
[+] Building 2.4s (11/11) FINISHED
 => [internal] load build definition from Dockerfile                                          0.0s
 => => transferring dockerfile: 399B                                                          0.0s
 => [internal] load .dockerignore                                                             0.0s
 => => transferring context: 34B                                                              0.0s
 => [internal] load metadata for docker.io/library/debian:11.1                                2.3s
 => [auth] library/debian:pull token for registry-1.docker.io                                 0.0s
 => [internal] load build context                                                             0.0s
 => => transferring context: 9.04kB                                                           0.0s
 => [1/5] FROM docker.io/library/debian:11.1@sha256:45ee40a844048c2f6d0105899c1a17733530b56d  0.0s
 => CACHED [2/5] RUN apt-get update &&     apt-get install -y build-essential pkg-config cma  0.0s
 => CACHED [3/5] ADD . /eland                                                                 0.0s
 => CACHED [4/5] WORKDIR /eland                                                               0.0s
 => CACHED [5/5] RUN python3 -m pip install --no-cache-dir --disable-pip-version-check .[all  0.0s
 => exporting to image                                                                        0.0s
 => => exporting layers                                                                       0.0s
 => => writing image sha256:51fe587c344ecb56430aba5ea931ff07297650046ffddd29e51a9eb543b2485a  0.0s
 => => naming to docker.io/library/eland                                                      0.0s

$ docker run eland eland_import_hub_model
usage: eland_import_hub_model [-h] (--url URL | --cloud-id CLOUD_ID)
                              --hub-model-id HUB_MODEL_ID
                              [--es-model-id ES_MODEL_ID] [-u ES_USERNAME]
                              [-p ES_PASSWORD] [--es-api-key ES_API_KEY]
                              --task-type
                              {text_embedding,question_answering,zero_shot_classification,ner,fill_mask,text_classification}
                              [--quantize] [--start] [--clear-previous]
                              [--insecure] [--ca-certs CA_CERTS]
eland_import_hub_model: error: the following arguments are required: --hub-model-id, --task-type

Elasticsearch / Kibana の起動

ローカルに Elasticsearch/Kibana を起動します。ネットワーク設定を簡単にするために eland も docker-compose.yml に定義しておきます。

version: "3"

services:
  elasticsearch:
    image: docker.elastic.co/elasticsearch/elasticsearch:8.2.0
    ports:
      - 9200:9200
    environment:
      - discovery.type=single-node
      - xpack.security.enabled=false
    ulimits:
      memlock:
        soft: -1
        hard: -1

  kibana:
    image: docker.elastic.co/kibana/kibana:8.2.0
    environment:
      ELASTICSEARCH_HOSTS: http://elasticsearch:9200
    ports:
       - 5601:5601

  eland:
    depends_on:
      - elasticsearch
      - kibana
    image: eland
    command: "echo specify eland command"

docker-compose コマンドで起動する

$ docker-compose up -d

起動を待って、以下の URL にアクセスして Kibana にアクセスできることを確認します。

http://localhost:5601

トライアルライセンスの有効化

Machine Learning を利用するためにトライアルライセンスを有効化します。

Stack Management > License Management > Start trial で有効化できます。

fujimotoshinji 2022/05/24に更新

2. テキスト埋め込みモデルのデプロイ

Elasticsearch Machine Learning にテキスト埋め込みモデルをデプロイします。
利用するモデルは概要にも書いた通り、 Hugging Face にある msmarco-distilbert-base-tas-b を利用します。

eland クライアントを利用することで Hugging Face のモデル ID？（GitHub の user/reposiroty みたいなの？）を指定するだけでデプロイできます。

$ docker-compose run eland eland_import_hub_model \
    --url http://elasticsearch:9200/ \
    --hub-model-id sentence-transformers/msmarco-distilbert-base-tas-b \
    --task-type text_embedding \
    --start

2022-05-23 05:52:04,990 INFO : Establishing connection to Elasticsearch
2022-05-23 05:52:05,001 INFO : Connected to cluster named 'docker-cluster' (version: 8.2.0)
2022-05-23 05:52:05,002 INFO : Loading HuggingFace transformer tokenizer and model 'sentence-transformers/msmarco-distilbert-base-tas-b'
Downloading: 100%|█████████████████████████████████████████████████| 547/547 [00:00<00:00, 568kB/s]
Downloading: 100%|█████████████████████████████████████████████████| 548/548 [00:00<00:00, 417kB/s]
Downloading: 100%|███████████████████████████████████████████████| 226k/226k [00:00<00:00, 343kB/s]
Downloading: 100%|████████████████████████████████████████████████| 112/112 [00:00<00:00, 10.7kB/s]
Downloading: 100%|██████████████████████████████████████████████| 253M/253M [00:05<00:00, 45.8MB/s]
Downloading: 100%|█████████████████████████████████████████████████| 690/690 [00:00<00:00, 349kB/s]
Downloading: 100%|████████████████████████████████████████████████| 190/190 [00:00<00:00, 98.6kB/s]
Downloading: 100%|████████████████████████████████████████████| 3.95k/3.95k [00:00<00:00, 1.12MB/s]
Downloading: 100%|█████████████████████████████████████████████████| 548/548 [00:00<00:00, 279kB/s]
Downloading: 100%|████████████████████████████████████████████████| 122/122 [00:00<00:00, 86.6kB/s]
Downloading: 100%|████████████████████████████████████████████████| 229/229 [00:00<00:00, 98.3kB/s]
Downloading: 100%|██████████████████████████████████████████████| 265M/265M [00:05<00:00, 52.0MB/s]
Downloading: 100%|██████████████████████████████████████████████| 53.0/53.0 [00:00<00:00, 34.2kB/s]
Downloading: 100%|████████████████████████████████████████████████| 112/112 [00:00<00:00, 74.0kB/s]
Downloading: 100%|███████████████████████████████████████████████| 466k/466k [00:00<00:00, 694kB/s]
Downloading: 100%|█████████████████████████████████████████████████| 547/547 [00:00<00:00, 492kB/s]
Downloading: 100%|███████████████████████████████████████████████| 232k/232k [00:00<00:00, 342kB/s]
2022-05-23 05:52:48,649 INFO : Creating model with id 'sentence-transformers__msmarco-distilbert-base-tas-b'
2022-05-23 05:52:48,861 INFO : Uploading model definition
100%|██████████████████████████████████████████████████████████| 64/64 [00:19<00:00,  3.30 parts/s]
2022-05-23 05:53:08,239 INFO : Uploading model vocabulary
2022-05-23 05:53:08,316 INFO : Starting model deployment
2022-05-23 05:53:13,542 INFO : Model successfully imported with id 'sentence-transformers__msmarco-distilbert-base-tas-b'

モデルのインストールが完了したら Kibana から API を実行して推論してみます。

POST 
/_ml/trained_models/sentence-transformers__msmarco-distilbert-base-tas-b/deployment/_infer
{
  "docs":  {
      "text_field": "how is the weather in jamaica"
    }
}

結果として 768 次元のベクトルが返ってきます。

{
  "predicted_value" : [
      -0.0919460579752922,
      -0.4940606653690338,
      0.035987671464681625,
       …
  ]
}

fujimotoshinji 2022/05/24

3. テストデータの準備

2019 TREC Deep Learning Track のテストステージで使用されたデータセット msmarco-passagetest2019-top1000.tsv を利用します。

Kibana から TSV ファイルを簡単にインデックスに取り込むことができます。

Kibana の Home から Upload a file を選択します。

TSV ファイルをブラウザにドラッグ＆ドロップします。

2列の TSV ファイルであることを確認して、Import を選択します

インデックスに collection、フィールド名に id, text を指定して import を選択します

アップロード後、collectionという名前のインデックスに182,469件の文書が格納されていることが確認できます

GET _cat/indices
yellow open collection Lj6lDI6JTgSzsDz_dWQtsg 1 1 182469 0 65.2mb 65.2mb

fujimotoshinji 2022/05/24

4. Ingest パイプラインの作成

インデキシングした TSV ファイルはテキストであり、ベクトル表現がないのでベクトル表現のフィールドを追加します。
Reindex にモデルから推論した（ベクトル表現を生成した）フィールドを追加する Ingest パイプラインを作成します。
Ingest パイプラインは Machine Learning の推論を実行できます。

PUT _ingest/pipeline/text-embeddings
{
  "description": "Text embedding pipeline",
  "processors": [
    {
      "inference": {
      "model_id": "sentence-transformers__msmarco-distilbert-base-tas-b",
      "target_field": "text_embedding",
      "field_map": {
        "text": "text_field"
       }
    }
  }
  ],
  "on_failure": [
    {
      "set": {
        "description": "Index document to 'failed-<index>'",
        "field": "_index",
        "value": "failed-{{{_index}}}"
      }
    },
    {
      "set": {
        "description": "Set error message",
        "field": "ingest.failure",
        "value": "{{_ingest.on_failure_message}}"
      }
    }
  ]
}

{"acknowledged" : true}

fujimotoshinji 2022/05/24に更新

5. Reindex によるテストデータのベクトル化フィールド生成

Reindex することでベクトル表現を追加したフィールドを追加したドキュメントをインデキシングします。
Reindex に Ingest パイプラインを通すことでテキスト埋め込みモデルの推論で返ってきたベクトル表現のフィールドをドキュメントに追加できます。

新しいインデックスはテキスト埋め込みモデルに合わせた 768次元の dense_vector フィールドを定義します。

PUT collection-with-embeddings
{
  "mappings": {
    "properties": {
      "text_embedding.predicted_value": {
        "type": "dense_vector",
        "dims": 768,
        "index": true,
        "similarity": "cosine"
      },
      "text": {
        "type": "text"
      }
    }
  }
}

{
  "acknowledged" : true,
  "shards_acknowledged" : true,
  "index" : "collection-with-embeddings"
}

作成したインデックスに向けて Reindex します。

POST _reindex?wait_for_completion=false
{
  "source": {
    "index": "collection"
  },
  "dest": {
    "index": "collection-with-embeddings",
    "pipeline": "text-embeddings"
  }
}

{"task" : "88Y0c7UuSD2ZIsQe8FB-8A:62304"}

返ってきた Task ID の Task API や、Kibana の Machine Learning の画面から進捗状況を見ることができます。

GET _tasks/88Y0c7UuSD2ZIsQe8FB-8A:62304

{
  "completed" : true,
  "task" : {
    "node" : "88Y0c7UuSD2ZIsQe8FB-8A",
    "id" : 62304,
    "type" : "transport",
    "action" : "indices:data/write/reindex",
    "status" : {
      "total" : 182469,
      "updated" : 0,
      "created" : 31000,
      "deleted" : 0,
      "batches" : 31,
      "version_conflicts" : 0,
      "noops" : 0,
      "retries" : {
        "bulk" : 0,
        "search" : 0
      },
      "throttled_millis" : 0,
      "requests_per_second" : -1.0,
      "canceled" : "by user request",
      "throttled_until_millis" : 0
    },
    "description" : "reindex from [collection] to [collection-with-embeddings]",
    "start_time_in_millis" : 1653288203832,
    "running_time_in_nanos" : 3072732486316,
    "cancellable" : true,
    "cancelled" : false,
    "headers" : {
      "trace.id" : "f5479039e1e2e728a20a1979550b5d0b"
    }
  },
  "response" : {
    "took" : 3072714,
    "timed_out" : false,
    "total" : 182469,
    "updated" : 0,
    "created" : 31000,
    "deleted" : 0,
    "batches" : 31,
    "version_conflicts" : 0,
    "noops" : 0,
    "retries" : {
      "bulk" : 0,
      "search" : 0
    },
    "throttled" : "0s",
    "throttled_millis" : 0,
    "requests_per_second" : -1.0,
    "canceled" : "by user request",
    "throttled_until" : "0s",
    "throttled_until_millis" : 0,
    "failures" : [ ]
  }
}

1時間ぐらいで 31,000 ドキュメントと 6時間ぐらいかかりそうだったので全部完了するのは諦めました。

データを 1件取得して確認してみると、768次元の数値配列がフィールドに追加されています。

POST collection-with-embeddings/_search
{"size": 1}

{
  "took" : 4545,
  "timed_out" : false,
  "_shards" : {
    "total" : 1,
    "successful" : 1,
    "skipped" : 0,
    "failed" : 0
  },
  "hits" : {
    "total" : {
      "value" : 3000,
      "relation" : "eq"
    },
    "max_score" : 1.0,
    "hits" : [
      {
        "_index" : "collection-with-embeddings",
        "_id" : "iyyf74ABrVs-dJ6RerRR",
        "_score" : 1.0,
        "_source" : {
          "text_embedding" : {
            "predicted_value" : [
              0.057356320321559906,
              0.16028179228305817,
              -0.18122528493404388,
              ...
            ],
            "model_id" : "sentence-transformers__msmarco-distilbert-base-tas-b"
          },
          "id" : 7130104,
          "text" : "This is the definition of RNA along with examples of types of RNA molecules. This is the definition of RNA along with examples of types of RNA molecules. RNA Definition"
        }
      }
    ]
  }
}

fujimotoshinji 2022/05/24

6. ベクトル検索

準備が終わったのでベクトル検索してみます。
もちろん検索キーワードも 768次元のベクトルである必要があります。
検索キーワードを推論して生成したベクトル表現で knn 検索してみます。

POST /_ml/trained_models/sentence-transformers__msmarco-distilbert-base-tas-b/deployment/_infer
{
  "docs": {
    "text_field": "how is the weather in jamaica"
  }
}

{
  "predicted_value" : [
      -0.0919460579752922,
      -0.4940606653690338,
      0.035987671464681625,
       …
  ]
}

_knn_search API でベクトル検索します。


}

{
  "took" : 13,
  "timed_out" : false,
  "_shards" : {
    "total" : 1,
    "successful" : 1,
    "skipped" : 0,
    "failed" : 0
  },
  "hits" : {
    "total" : {
      "value" : 100,
      "relation" : "eq"
    },
    "max_score" : 0.965587,
    "hits" : [
      {
        "_index" : "collection-with-embeddings",
        "_id" : "7Cyf74ABrVs-dJ6RgeZJ",
        "_score" : 0.965587,
        "_source" : {
          "id" : 6112586,
          "text" : "The coolest times in Jamaica are the winter months of December to February. The average temperature at this time is 75 °F. Winter temperatures vary with the summer temperatures by 10 °F to 23 °F.Remember our climate is cool and not cold.or the most part, the average temperature in Jamaica is between 80 °F and 90 °F (27 °FCelsius-29 °Celsius). Luckily, the weather in Jamaica is always vacation friendly. You will hardly experience long periods of rain fall, and you will become accustomed to weeks upon weeks of sunny weather."
        }
      },
      {
        "_index" : "collection-with-embeddings",
        "_id" : "Ziyf74ABrVs-dJ6RgedJ",
        "_score" : 0.9655287,
        "_source" : {
          "id" : 6467518,
          "text" : "Weather in Jamaica. The weather patterns you'll encounter in Jamaica can vary dramatically around the island. Photo credit: © Dennis Sabo | Dreamstime.com. Regardless of when you visit, the tropical climate and warm temperatures of Jamaica essentially guarantee beautiful weather during your vacation. Average temperatures in Jamaica range between 80 degrees Fahrenheit and 90 degrees Fahrenheit, with July and August being the hottest months and February the coolest."
        }
      },
      {
        "_index" : "collection-with-embeddings",
        "_id" : "jiyf74ABrVs-dJ6Rfcqw",
        "_score" : 0.96532303,
        "_source" : {
          "id" : 823272,
          "text" : "Before I tell you when is the best time to visit Jamaica based on weather conditions, let us take a look at the monthly climate of Jamaica. Here is a brief summary of the weather in Jamaica for each month. January. January has an average high temperature of 82.3 °F and a low of 74 °F. You can expect to experience rain about ten days of the month."
        }
      },
      {
        "_index" : "collection-with-embeddings",
        "_id" : "_S2f74ABrVs-dJ6RhQ7A",
        "_score" : 0.963459,
        "_source" : {
          "id" : 3391439,
          "text" : "The coolest times in Jamaica are the winter months of December to February. The average temperature at this time is 75 °F. Winter temperatures vary with the summer temperatures by 10 °F to 23 °F. Remember our climate is cool and not cold. That’s because Jamaica does not have a “cold” weather; in fact, our winter is relatively warm when compared to New York, Miami, England and Canada."
        }
      },
      {
        "_index" : "collection-with-embeddings",
        "_id" : "Pyyf74ABrVs-dJ6RgvmN",
        "_score" : 0.9628865,
        "_source" : {
          "id" : 3074599,
          "text" : "Regardless of when you visit, the tropical climate and warm temperatures of Jamaica essentially guarantee beautiful weather during your vacation. Average temperatures in Jamaica range between 80 degrees Fahrenheit and 90 degrees Fahrenheit, with July and August being the hottest months and February the coolest. Temperatures in Jamaica generally vary approximately 10 degrees from summer to winter."
        }
      },
      {
        "_index" : "collection-with-embeddings",
        "_id" : "MCyf74ABrVs-dJ6Re8Tm",
        "_score" : 0.9595746,
        "_source" : {
          "id" : 8160224,
          "text" : "Guide to Jamaica weather in January The average maximum daytime temperature in Jamaica in January is a warm 29°C (84°F) with moderate heat & humidity. The average night-time temperature is usually a comfortable 20°C (68°F)."
        }
      },
      {
        "_index" : "collection-with-embeddings",
        "_id" : "Siyf74ABrVs-dJ6RertV",
        "_score" : 0.9590724,
        "_source" : {
          "id" : 7218012,
          "text" : "Thanks to its position in the warm Caribbean Sea, Jamaica enjoys generally pleasant weather year round. While October may not be the peak time to visit this tiny yet world-renowned tourist destination, any month is a fine month to travel to Jamaica. While you'll see some rain showers, you'll still enjoy plenty of warm and beach-friendly weather."
        }
      },
      {
        "_index" : "collection-with-embeddings",
        "_id" : "eSyf74ABrVs-dJ6RgeZJ",
        "_score" : 0.95751655,
        "_source" : {
          "id" : 6103289,
          "text" : "Jamaica Weather in September: Still Hot & Still Humid. Jamaica weather in September is very much like weather in August with one notable exception... ...the possibility of more rain during this month. However, the days are still hot, still humid, and still filled with long hours of daylight."
        }
      },
      {
        "_index" : "collection-with-embeddings",
        "_id" : "aSyf74ABrVs-dJ6RgedJ",
        "_score" : 0.957516,
        "_source" : {
          "id" : 6467524,
          "text" : "Find the best month to visit Jamaica based on the weather. Jan Feb Mar Apr May Jun Jul Aug Sep Oct Nov Dec Expect 29°C daytime maximum temperatures in the shade with on average 8 hours of sunshine per day in Jamaica in February."
        }
      },
      {
        "_index" : "collection-with-embeddings",
        "_id" : "gSyf74ABrVs-dJ6RerdU",
        "_score" : 0.9567431,
        "_source" : {
          "id" : 8033648,
          "text" : "Typical Weather in March. March temperatures in Jamaica range from an average low of 73.0 degrees to an average high of 84.0 degrees. On a typical March day, during the early morning hours you will probably find the temperature will be around 74.1 degrees."
        }
      }
    ]
  }
}

10件とも Jamaica の天気に関する結果が返ってきた。
今回の 18万件のテストデータの内容をわかっていないので詳細な評価はできないが、少なくとも全く関係ないものを返しているということはなさそう。
文字列から生成した数値同士で類似したものを返しているようです。

fujimotoshinji 2022/05/24

まとめ

全文検索の弱み・辛みをベクトル検索で解決できることがある
- 全文検索の上位互換ではない、ベクトル検索にはベクトル検索の弱み・辛みがある
- パッと思いつくのはチューニングの難しさ、性能、不透明さ、技術としての難しさ
Elasticsearch の Machine Learning を利用することでベクトル検索する上で Elasticsearch の外での処理をなくすことができる
- eland などのツールが整っているのでとにかく簡単（あくまで手順として）にインデキシング、ベクトル検索ができる
推論の性能が遅い
- 設定次第で早くなるのかな
- Docker コンテナのリソース割り当ての問題かとそこそこ大きいリソースの EC2（m6g.2xlarge）で試したら 10倍ぐらい遅くてはてなになった
Machine Learning 機能は有償機能（Platinum サブスクリプション以上）

このスクラップは2022/05/24にクローズされました