概要

機械学習による検索ランキング改善ガイドという検索ランキング改善に関する優れた書籍が存在します。

この本ではSolrとElasticsearchについて詳しく解説されていますが、Vespaで機械学習モデルを使用して検索改善を行う具体的な方法については触れられていません。そこで、この記事では機械学習による検索ランキング改善ガイドの内容をVespaで実施しました。記事で作成したコードは別途公開しているリポジトリにて確認できます。

環境構築

環境構築の詳細は長くなるので、興味のある方は以下を開いてご確認ください。

環境構築の詳細

書籍ではElasticsearchをDockerで動かしていました。同様に、VespaもDockerで動かしています
- 今回は、検索ランキングの改善のみを試すので、マルチノード構成は取らず、1ノードのみで立てています
- 開発環境はDevContainerで構築します

開発環境とVespaの用意

使用したdocker-comopose.yamlは以下のとおりです
- 3つのserviceを定義しています
- workspaceは、作業用のコンテナでPythonやvespa-cliが入っています。後ほど中身を説明します
  - workspaceはDev Containerとして使用しています
- vespaは、vespaサーバを動かすコンテナです
  - healthcheckを記述しているのは、Vespaの起動後にDev Containerを起動するようにしたいからです
    - Dev Container側では、postCreateCommand.shでVespaの設定を反映する vespa deploy コマンドを実行しているため、先にvespaが起動している必要があります
- vispanaは、vespaのGUIツールです。ElasticsearchでのKibanaのようなのものです。Vespaの公式が出しているわけではないですが、Vespaの状態や設定の確認、クエリの実行などができます
  - https://github.com/vispana/vispana

docker-compose.yml

version: '3'

services:
  workspace:
    build:
      context: .
      dockerfile: .devcontainer/Dockerfile
    init: true
    environment:
      - TZ=Asia/Tokyo
    command: sleep infinity
    volumes:
      - .:/workspace:cached
      - workspace_venv:/workspace/.venv
      - workspace_bin:/workspace/.bin
    depends_on:
      vespa:
        condition: service_healthy
    networks:
      - vespa_network
  
  vespa:
    image: vespaengine/vespa
    container_name: vespa
    ports:
      - "8080:8080"
      - "19071:19071"
      - "19092:19092"
    volumes:
      - ./vespa-config:/vespa-config
      - vespa_var_storage:/opt/vespa/var
      - vespa_log_storage:/opt/vespa/logs
    networks:
      - vespa_network
    healthcheck:
      test: curl http://localhost:19071/state/v1/health
      timeout: 10s
      retries: 3
      start_period: 40s      

  vispana:
    image: vispana/vispana:latest
    container_name: vispana
    ports:
      - 4000:4000
    networks:
      - vespa_network

networks:
  vespa_network:
    driver: bridge

volumes:
  workspace_venv:
  workspace_bin:
  vespa_var_storage:
  vespa_log_storage:

workspace用のDockerfileは以下のようになっています
- Pythonの設定とvespa-cliのインストールを行っています

Dockerfile

ARG VARIANT=3.12-bookworm
ARG VESPA_CLI_VERSION=8.367.14
FROM mcr.microsoft.com/devcontainers/python:${VARIANT}

ENV PYTHONUNBUFFERED 1
ENV TZ Asia/Tokyo

# vespa-cliのインストール
RUN curl -L -o .bin/vespa-cli.tar.gz https://github.com/vespa-engine/vespa/releases/download/v${VESPA_CLI_VERSION}/vespa-cli_${VESPA_CLI_VERSION}_linux_amd64.tar.gz \
    tar -xvf .bin/vespa-cli.tar.gz -C .bin \
    mv .bin/vespa-cli_${VESPA_CLI_VERSION}_linux_amd64/bin/vespa .bin/vespa \
    rm .bin/vespa-cli.tar.gz \
    rm -rf .bin/vespa-cli_${VESPA_CLI_VERSION}_linux_amd64

Dev Containerの設定は以下のようになっています
- 先ほど説明したdocker-compose.ymlファイルを指定しています。また、今回poetryを使用してPythonのパッケージ管理をしているので、Dev Container Featuresでインストールしています

devcontaier.json

{
  "name": "building-search-app-w-ml-vespa",
  "dockerComposeFile": "../docker-compose.yml",
  "service": "workspace",
  "workspaceFolder": "/workspace",
  "features": {
    "ghcr.io/devcontainers-contrib/features/poetry:2": {}
  },
  "remoteEnv": {
    "PATH": "${containerEnv:PATH}:/workspace/.bin"
  },
  "postCreateCommand": "./.devcontainer/postCreateCommand.sh",
  "customizations": {
    // 長いので省略
  }
}

コンテナ作成後実行されるpostCreateCommand.shでは、poetry install と vespaの設定のデプロイを行っています

postCreateCommand.sh

#!/bin/sh
# postCreateCommand.sh

echo "START Install"

sudo chown -R vscode:vscode .

poetry config virtualenvs.in-project true
poetry install

echo "SETUP Vespa"

vespa config set target http://vespa:19071
vespa status
vespa deploy vespa-config --wait 30

echo "FINISH"

これらの設定により、Dev Containerを起動するだけで、vespaが設定の反映まで完了した状態で使用できるようになります
なお、Dev Containerの設定には、以下の記事を参考にさせていただきました

https://tech.dentsusoken.com/entry/2023/05/02/Dev_Container%E3%82%92%E4%BD%BF%E3%81%A3%E3%81%A6%E3%82%B9%E3%83%86%E3%83%83%E3%83%97%E3%83%90%E3%82%A4%E3%82%B9%E3%83%86%E3%83%83%E3%83%97%E3%81%A7%E4%BD%9C%E3%82%8BPython%E3%82%A2%E3%83%97%E3%83%AA%E3%82%B1

https://blog.johtani.info/blog/2023/07/21/multi-containers-with-dev-container/

機械学習モデルを用いたリランキングを行うまでのステップ

機械学習を用いたリランキングを行うには、以下のステップを踏みます

Vespaのスキーマ設定と特徴量取得のためのVespaのrankprofile設定
Vespaへのデータフィード
Vespaで検索を行い特徴量データを取得
モデルの学習
機械学習モデルを使うためのrankprofile設定とモデルのデプロイ
機械学習モデルを用いた検索の実施

1.Vespaのスキーマ設定と特徴量取得のためのVespaのrankprofile設定

まず、Vespaのフィードするドキュメントのスキーマについて説明します。スキーマ定義ファイルは.sdという拡張子のファイルです。スキーマ定義ファイルの中身は以下のようになっています。idとpageviewフィールドはattributeとして、titleとtextフィールドはindexとして定義しました。

simplewiki.sd

schema simplewiki {

    document simplewiki {

        field id type string {
            indexing: summary | attribute
            attribute: fast-search
        }

        field title type string {
            indexing: summary | index
            index: enable-bm25
        }

        field text type string {
            indexing: summary | index
            index: enable-bm25
        }

        field pageviews type int {
            indexing: summary | attribute
        }
    }

    fieldset default {
        fields: text
    }
}

スキーマ設定に加えて、特徴量を収集する際に使用するrankprofileをbase.profileという名前で作成します。
機械学習ランキングモデルを学習するには、学習のインプットである特徴量となるデータが必要です。Vespaではrankfeatureのsummary-featuresのリストに追加することで、検索時に特徴量を返すことができます。

base.profile

rank-profile base inherits default {

    summary-features {
        queryTermCount
        nativeRank
        attribute(pageviews)
    }
}

ここではサンプルで上記の３つの特徴量のみ指定していますが、Vespaではデフォルトで他にもたくさんの特徴量が用意されています。Vespaで扱うことができる特徴量は以下のページにまとまっています。

これらのVespaの設定は、決められた構造で一つのディレクトリにまとめる必要があります。今回は、vespa-configというディレクトリにまとめました。vespa-configディレクトリは以下のような構成になっています。

.
├── schemas
│   ├── simplewiki
│   │   └── base.profile
│   └── simplewiki.sd
└── services.xml

このファイル群をVespaにデプロイすることで、Vespaの設定が完了します。

Vespaへのデータフィード

スキーマ設定が完了したので、次にVespaへドキュメントをフィードしていきます。
Vespaへのデータフィードはフィード用のエンドポイントにhttpリクエストを送ることで行います。今回は、vespaへのリクエストにvespaのpythonクライアントであるpyvespaを使用しました。

#!usr/bin/env python
import bz2
from vespa.application import Vespa


def generate_bulk_buffer():
    buf = []
    with bz2.open("dataset/simplewiki-202109-pages-with-pageviews-20211001.bz2", "rt") as bz2f:
        for line in bz2f:
            id, title, text, pageviews = line.rstrip().split("\t")
            buf.append(
                {
                    "id": id,
                    "fields": {
                        "title": title,
                        "text": text,
                        "pageviews": pageviews
                    }
                }
            )
            if 500 <= len(buf):
                yield buf
                buf.clear()
    if buf:
        yield buf
        

client = Vespa(url = 'http://vespa', port = 8080)

def callback(response, id):
    if not response.is_successful():
        print(
            f"Failed to feed document {id} with status code {response.status_code}: Reason {response.get_json()}"
        )
    
for buf in generate_bulk_buffer():
    client.feed_iterable(buf, schema = "simplewiki", callback=callback)

上記がフィードのためのコードです。書籍のElasticsearchへのフィードを行うコードとほぼ同じ形で実装できています。

Vespaで検索を行い特徴量データを取得

検索結果の取得も、Elasticsearch版とほとんど同じコードになっています。
違いとしては、vespaに送るクエリです。重複する部分については割愛していますが、以下がvespaに送るクエリを生成するコードです。

def generate_query_to_collect_features(keywords, size=10):
    return {
        "yql": "select title, summaryfeatures from simplewiki where userInput(@userinput)",
        "userinput": keywords,
        "ranking": {"profile": "base"},
        "hits": size,
        "presentation.timing": True,
    }

"ranking": {"profile": "base"} で、先ほど作成したrankprofileプロファイルを指定して検索を行っています。また、select title, summaryfeaturesでレスポンスにドキュメントのタイトルと、summaryfeatureを含めるように指示しています。

実際にクエリを発行した結果、ドキュメントごとに以下のようにsummaryfeatureの値を取得できます。

      {
        "id": "index:simplewiki/0/c6cbfacce17abb72eca191bf",
        "relevance": 0.35871986310761605,
        "source": "simplewiki",
        "fields": {
          "title": "Movie rights",
          "summaryfeatures": {
            "attribute(pageviews)": 2,
            "nativeRank": 0.35871986310761605,
            "queryTermCount": 1,
            "vespa.summaryFeatures.cached": 0
          },
          "relevance": 0.35871986310761605
        }
      },

このsummaryfeaturesの値をCSVファイルなどに落とすことで、機械学習モデルのトレーニングデータとします。

モデルの学習

モデルの学習については、検索エンジンの種類によらず同じであるため、割愛します。
書籍では、モデルとしてXGBoostモデルが利用されていました。Vespaではプラグインなどを追加することなくデフォルトでXGBoostモデルが利用できるので、同じ方法でモデルを作成しています。

機械学習モデルを使うためのrankprofile設定とモデルのデプロイ

モデルが作成できたら、モデルとモデルを使うためのrankprofileをVespaにデプロイします。

先ほどのbase.profileとは別に機械学習モデルを使用してランキングを行う際に指定するrerank.profileを用意します。

profile

rank-profile xgboost inherits base {

    second-phase {
        expression: xgboost("hands_on_model.json")
    }
}

rerank.profileではbase.profileを継承し、second-pahseでリランキングの設定をしています。
second-phaseでは、作成したモデルの名前をxgboost関数に渡します。

この時点でvespa-configディレクトリの構成は以下のようになっています。

.
├── models
│   └── hands_on_model.json
├── schemas
│   ├── simplewiki
│   │   ├── base.profile
│   │   └── rerank.profile
│   └── simplewiki.sd
└── services.xml

このモデルを含めた設定をVespaにデプロイすることで、検索時にリランキングが可能になります。

機械学習モデルを用いた検索の実施

最後に、検索クエリで機械学習モデルを使用したリランキングを指定する方法を説明します。
といっても、非常に簡単で使用するランクプロファイル名を指定するだけです。


def generate_query_to_search_with_mlr(keywords, size=10, window_size=100):
    return {
        "yql": "select title from simplewiki where userInput(@userinput)",
        "userinput": keywords,
        "ranking": {
            "profile": "xgboost",
            "rerankCount": window_size,
        },
        "hits": size,
        "presentation.timing": True,
    }

まとめ

Vespaで機械学習による検索ランキングの改善を行いました
わずかな設定だけで、機械学習モデルによるリランキングが実現できました
機械学習を用いたリランキングについてはVespaがデフォルトでサポートしていることもあり、プラグインなどを使用することなく簡単に実装することができました