🎃

ローカルでOpenSearchのLTRプラグインを動かしてみる

に公開

はじめに

Learning To Rank(LTR)は、機械学習を活用して検索結果のランキングを最適化する技術で、ユーザーの行動データや文書の特徴量を学習してより関連性の高い結果を上位表示できます。オープンソースの検索エンジンであるOpenSearchには、このLTR機能を利用できるプラグインが用意されています。今回はドキュメントを参考にして、LTRプラグインを実際にローカル環境で動かしてみました。その備忘録として、具体的な手順をまとめたいと思います。

環境構築

Dockerを使ってOpenSearchコンテナを起動します。公式ドキュメントLTRプラグインのリポジトリを参考に、docker-compose.yamlとDockerfileを用意します。Dockerfileでは、OpenSearch公式のDockerイメージをベースに、LTRプラグインをインストールします。LTRプラグインに合わせて、OpenSearchのバージョンは2.17.1にしています。

Dockerfile
FROM opensearchproject/opensearch:2.17.1

ADD --chown=opensearch:opensearch https://github.com/opensearch-project/opensearch-learning-to-rank-base/releases/download/2.17.1/ltr-2.17.1-os2.17.1.zip ltr-2.17.1-os2.17.1.zip

RUN /usr/share/opensearch/bin/opensearch-plugin install --batch file:ltr-2.17.1-os2.17.1.zip

1ノード構成でコンテナを定義します。今回はLTRプラグインの動作確認が目的なので、セキュリティ設定を無効化しています。

docker-compose.yaml
version: '3'
services:
  opensearch-node:
    build: .
    container_name: opensearch-node
    environment:
      - cluster.name=opensearch-cluster
      - node.name=opensearch-node
      - discovery.seed_hosts=opensearch-node
      - cluster.initial_cluster_manager_nodes=opensearch-node
      - bootstrap.memory_lock=true
      - OPENSEARCH_JAVA_OPTS=-Xms512m -Xmx512m
      - OPENSEARCH_INITIAL_ADMIN_PASSWORD=${OPENSEARCH_INITIAL_ADMIN_PASSWORD}
      - "DISABLE_INSTALL_DEMO_CONFIG=true" # 動作確認用なのでdisableにする
      - "DISABLE_SECURITY_PLUGIN=true" # 動作確認用なのでdisableにする
    ulimits:
      memlock:
        soft: -1
        hard: -1
      nofile:
        soft: 65536
        hard: 65536
    ports:
      - 9200:9200
    networks:
      - opensearch-net

networks:
  opensearch-net:

環境変数にパスワードを設定し、コンテナを起動します。

$ export OPENSEARCH_INITIAL_ADMIN_PASSWORD=<password>
$ docker-compose up -d

コンテナが正常に起動したか確認します。

$ curl http://localhost:9200
{
  "name" : "opensearch-node",
  "cluster_name" : "opensearch-cluster",
  "cluster_uuid" : "4l50cH0eQd-8SG1HjdC59w",
  "version" : {
    "distribution" : "opensearch",
    "number" : "2.17.1",
    "build_type" : "tar",
    "build_hash" : "1893d20797e30110e5877170e44d42275ce5951e",
    "build_date" : "2024-09-26T21:59:32.078798875Z",
    "build_snapshot" : false,
    "lucene_version" : "9.11.1",
    "minimum_wire_compatibility_version" : "7.10.0",
    "minimum_index_compatibility_version" : "7.0.0"
  },
  "tagline" : "The OpenSearch Project: https://opensearch.org/"
}

特徴量セットをデプロイする

LTRプラグインでは、feature storeというインデックスに特徴量セットやモデルを保存します。
まずは、特徴量セットをデプロイするために、feature storeインデックスをPUTリクエストで作成します。

$ curl -X PUT "http://localhost:9200/_ltr"

次に、モデルが使用する特徴量セットを用意します。featuresetオブジェクトの中に、featuresという配列を定義し、その中に個々の特徴量をオブジェクトとして記述していきます。今回は、titleoverviewという2つのフィールドを持つドキュメントを想定して、tmbd_titletmdb_multiという特徴量セットを定義しています。

featureset.json
{
    "featureset": {
        "features": [
            {
                "name": "tmdb_title",
                "params": [
                    "keywords"
                ],
                "template": {
                    "match": {
                        "title": "{{keywords}}"
                    }
                }
            },
            {
                "name": "tmdb_multi",
                "params": [
                    "keywords"
                ],
                "template": {
                    "multi_match": {
                        "query": "{{keywords}}",
                        "fields": ["title", "overview"]
                    }
                }
            }
        ]
    }
}

more_moview_featuresという名前で特徴量セットをデプロイします。

$ curl -XPOST "http://localhost:9200/_ltr/_featureset/more_movie_features" -H "Content-Type: application/json" -d "@featureset.json"

機械学習モデルをOpenSearchにデプロイする

LTRプラグインでは、モデルを学習する機能は提供されていません。そのため、pythonでモデルを学習させて、モデルの定義ファイルを作成します。OpenSearchでサポートされているモデルはXGboostとRankLibです。今回はXGboostを使用します。

学習には、リポジトリに用意されているデモファイルを使用します。

  • xgboost.txt: libsvm形式でサンプルのトレーニングデータが記載されています
  • featmap.txt: 使用する特徴量が記載されたconfigファイルです

以下のpythonスクリプトで、モデルの定義ファイルを作成します。

xgb.py
import xgboost as xgb
import json

dtrain = xgb.DMatrix('xgboost.txt?format=libsvm')
param = {'max_depth':2, 'eta':1, 'silent':1, 'objective':'reg:linear'}
num_round = 2

bst = xgb.train(param, dtrain, num_round)

model = bst.get_dump(fmap='featmap.txt', dump_format='json')

with open('xgb-model.json', 'w') as output:
    json_model = {
        "model": {
            "name": "xgb-model",
            "model": {
                "type": "model/xgboost+json",
                "definition": "[" + ",".join(list(model)) + "]",
            },
        }
    }
    json.dump(json_model, output, indent=2)
    output.write("\n")

必要なライブラリをインストールしてスクリプトを実行すると、xgb-model.jsonが出力されます。

$ pip install xgboost
$ python xgb.py
xgb-model.json
{
  "model": {
    "name": "xgb-model",
    "model": {
      "type": "model/xgboost+json",
      "definition": "[  { \"nodeid\": 0, \"depth\": 0, \"split\": \"tmdb_multi\", \"split_condition\": 14.3700256, \"yes\": 1, \"no\": 2, \"missing\": 2 , \"children\": [\n    { \"nodeid\": 1, \"leaf\": -0.943529487 }, \n    { \"nodeid\": 2, \"leaf\": 1.60399997 }\n  ]},  { \"nodeid\": 0, \"depth\": 0, \"split\": \"tmdb_title\", \"split_condition\": 8.13691807, \"yes\": 1, \"no\": 2, \"missing\": 2 , \"children\": [\n    { \"nodeid\": 1, \"depth\": 1, \"split\": \"tmdb_title\", \"split_condition\": 4.41261435, \"yes\": 3, \"no\": 4, \"missing\": 4 , \"children\": [\n      { \"nodeid\": 3, \"leaf\": -0.0288857929 }, \n      { \"nodeid\": 4, \"leaf\": -0.0950589329 }\n    ]}, \n    { \"nodeid\": 2, \"leaf\": 0.637333214 }\n  ]}]"
    }
  }
}

POSTリクエストでデプロイ完了です。

$ curl -XPOST "http://localhost:9200/_ltr/_featureset/more_movie_features/_createmodel" -H "Content-Type: application/json" -d "@xgb-model.json"

機械学習モデルを使って検索する

モデルをデプロイしたので、OpeanSearchにサンプルデータを投入して、実際に検索をしてみます。
サンプルデータとして、sample_data.jsonを用意しました。動作確認が目的なので、簡単なデータにしています。

sample_data.json
{"index": {"_index": "more_movie_features", "_id": "1"}}
{"title": "sample movie 1", "overview": "This is a sample overview for movie 1."}
{"index": {"_index": "more_movie_features", "_id": "2"}}
{"title": "sample movie 2", "overview": "This is a sample overview for movie 2."}
{"index": {"_index": "more_movie_features", "_id": "3"}}
{"title": "sample movie 3", "overview": "This is a sample overview for movie 3."}

POSTリクエストでデータを投入します。

$ curl -XPOST "http://localhost:9200/_bulk" -H "Content-Type: application/x-ndjson" --data-binary "@sample_data.json"

デプロイしたモデルを使って検索する場合、通常の検索クエリとは異なる特殊な形式を使用します。LTRプラグインでは、rescoreという機能を利用して、最初の検索結果の順位をモデルで再評価します。

search_query.json
{
  "query": {
    "query_string": {
      "query": "sample",
      "fields": ["title"]
    }
  },
  "rescore": {
    "window_size": 10, // 再スコアリングの対象となるドキュメント数
    "query": {
      "rescore_query": { // 再スコアリングで使用するクエリ
        "sltr": {
          "params": {
            "keywords": "sample" // 特徴量に渡すパラメータ
          },
          "model": "xgb-model", // 使用するモデル名
          "featureset": "more_movie_features" // 使用する特徴量セット名
        }
      },
      "query_weight": 0, // ベースクエリのスコアの重み(0に設定するとモデルのスコアのみが使われます)
      "rescore_query_weight": 1 // LTRモデルのスコアの重み
    }
  }
}

GETリクエストで検索します。

curl -XGET "http://localhost:9200/more_movie_features/_search" -H "Content-Type: application/json" -d "@search_query.json"

以下のような結果が返ってくれば成功です。

{"took":1,"timed_out":false,"_shards":{"total":1,"successful":1,"skipped":0,"failed":0},"hits":{"total":{"value":3,"relation":"eq"},"max_score":-0.97241527,"hits":[{"_index":"more_movie_features","_id":"4","_score":-0.97241527,"_source":{"title": "sample movie 1", "overview": "This is a sample overview for movie 1."}},{"_index":"more_movie_features","_id":"5","_score":-0.97241527,"_source":{"title": "sample movie 2", "overview": "This is a sample overview for movie 2."}},{"_index":"more_movie_features","_id":"6","_score":-0.97241527,"_source":{"title": "sample movie 3", "overview": "This is a sample overview for movie 3."}}]}}

終わりに

今回は、ローカル環境でOpenSearchのLTRプラグインを試してみました。今後としては、セマンティック検索やハイブリッド検索など試していきたいです。

nextbeat Tech Blog

Discussion