概要

Elastic 8.0 から NLP関連の機械学習ジョブを Elastic の Machine Learning で実行できるようになりました。(NLP のリファレンス) NLPで対応できるジョブにもいくつか種類はありますが、今回はテキストエンべディングを活用した類似文章の検索を試してみたいと思います。

テキストエンべディングは、かなり雑にまとめると、文章をベクトル化し機械的に比較できるようにするための手法です。

吾輩は猫である ==(MLでテキストエンべディング)==> [[0.12618473172187805, -0.11623416095972061, -0.08502635359764099, -0.2107191085815429, ...]

Elastic に置き換えて話すと、機械学習ジョブで取り込む文章をベクトル化（数字の配列）し、ドキュメントの dense_vector フィールドにベクトルを書き込みます。検索したい文章を同じジョブでベクトル化し、ベクトル検索で類似するベクトルを探します。ベクトルの類似度が高ければ高いほど意味合いの似た文章という事になります。

*8.5 時点ではNLPジョブは betaです。

前提

Elastic の機械学習機能を使うには、Platinum 以上の有償サブスクリプションが必要になりますので、ご注意ください。
上記前提のもと、Self managed の場合、ノードに機械学習のロールを持った物が必要となり、Elastic Cloud の場合、機械学習ノードが必要です。

NLPモデルのアップロードにPython と eland というライブラリも使います。詳細はこちら。
# 私自身あまり Python は強くないので、それ関連の細かい話は割愛しています。Python の勉強がてらやっています。。。

実装

実際に試すためにやった事ですが、おおまかに3点

利用する機械学習モデルのアップロード
検索対象とする文章データの取り込み
文章を検索する

コードを全部 github にあげたましたので、もし気になる方いらっしゃいましたら。(Github Repo )

1) 利用する機械学習モデルのアップロード

アップロードのスクリプト

huggingface から探してきたモデルを eland を使って Elasticsearch のクラスターにアップロードしています。eland を pip でインストールすると eland_import_hub_model がコマンドラインから実行できるので、あとは必要な引数を与えるだけです。

run.sh

eland_import_hub_model --es-api-key --url $ESS_URL --hub-model-id $HUB_MODEL_ID --task-type $TASK_TYPE --start

2) 検索対象とする文章データの取り込み

取り込みスクリプト
# コードの最後の2行は、試行錯誤中に全部回らないように、1000件程度で止めるようにしてました。全件入れると 17,000件ほどあると思います。それなりに処理の負荷がかかったのでタイムアウトに悩まされたりもしました。

11行目〜17行目あたりに変更が必要な変数があります。

load_aoz.py

...
repo_path = "/path/to/aozorabunko-master" # DL/clone repo from https://github.com/aozorabunko/aozorabunko

es = Elasticsearch(
        "https://xxxxyyyyzzzz.es.asia-northeast1.gcp.cloud.es.io:9243",
        api_key=("REPLACE_WITH_BASE64_API_KEY"),
        request_timeout=60
        )
...

Github に上がっている青空文庫のサイト（Github Repo）をダウンロードし、青空文庫に登録されている小説をデータソースとして使ってみました。

取り込むために書いたスクリプトの簡単な流れの紹介ですが、

最初にインデックスの削除、インジェストパイプラインの作成、インデックスの作成（マッピングの定義）を行っています。

ここで重要なのは、マッピングで dense_vector フィールドの指定です。MLモデルが返すベクトル長に合わせて dims の値を設定します。別のモデルを使う場合はここの確認が必須です。

"text_embedding.predicted_value": {
  "type": "dense_vector",
  "dims": 512,
  "index": True,
  "similarity": "cosine"
},

次に、特定のファイルの命名規則（実際の小説の中身の入ったHTMLファイル）にそってファイルの一覧を取得し、ひとつずつ処理していきます。特定のHTMLタグや、classを引っ張ってきて、フィールドに落とし込んでいます。HTMLの解釈自体は Beautiful Soup というライブラリにまかせています。

load_aoz.py

...
path = repo_path + "/cards/**/files/*_*.html"

aoz_bulk = []
    for ind, fn in enumerate(glob.glob(path, recursive=True)): 
...

ある程度まとめて処理した方が効率がいいので、_bulk API を使ってインデックスしています。

load_aoz.py

...
if ind % 100 == 0 and len(aoz_bulk) > 0:
    print("bulk at " + str(ind))
    # print(aoz_bulk)
    resp = es.bulk(index=index_name, body=aoz_bulk)
    print(resp)
    aoz_bulk = []
...

たまにサイトのフォーマットがバラつくことがあったり、命名規則までちゃんと見てないため、たまにパースエラーがあったりしますが、今回はあまり気にせずすっ飛ばしています。

結果的に、17,000件弱ほど取り込めました。

GET aoz/_count
=>
{
  "count": 16763,
  "_shards": {
    "total": 1,
    "successful": 1,
    "skipped": 0,
    "failed": 0
  }
}

3) 文章を検索する

検索スクリプト

検索の前に、一度 Infer trained model API を使って検索したい文章を text embedding のジョブを通してベクトル化しています。文章の中の文字列ではなく、意図を検索するには文章の特徴を抽出できるモデルを使ってベクトル化し、ベクトル同士の比較が必要になります。

search.py

...
resp = es.ml.infer_trained_model(model_id=model_name, docs=infer_docs)
input_embed = resp['inference_results'][0]['predicted_value']
...

検索は一般的な _search エンドポイントでの検索で Query DSL を書いています。8.4から検索クエリーに knn オプションが追加されたため、ベクトル検索を従来の match クエリー等と一緒に使えるようになりました。先程ベクトル化した文章を、knn の検索対象ベクトルに埋め込んでいます。boost を使うと重み付けによる検索結果のチューニングもできるようになります。

search.py

...
knn_query = {
    "field": "text_embedding.predicted_value",
    "query_vector": input_embed,
    "k": 10,
    "num_candidates": 100,
    "boost": 0.3
}

search_query = {
    "match": {
        "text_field": {
            "query": sys.argv[1],
            "boost": 0.7
        }
    }
}

search_fields = ["title", "author", "file_path"]

search_result = es.search(knn=knn_query, query=search_query, fields=search_fields, source=False)
...

あとは検索を実行し、結果を表示するだけです。

最終的な動き

先程作ったインデックスに対して「吾輩は猫である」を検索してみました。
path: の値に、https://www.aozora.gr.jp をつけると実際の文章が確認できますので興味のある方はどの程度の検索精度だったか確認していただけます。

# 検索
$ python3 search.py 吾輩は猫である
=>
夏目漱石 - 吾輩は猫である	score: 8.415148	path: /cards/000148/files/789_14547.html
夢野久作 - 超人鬚野博士	score: 8.358631	path: /cards/000096/files/2099_22206.html
夏目漱石 - 『吾輩は猫である』上篇自序	score: 8.080666	path: /cards/000148/files/47148_32217.html
伊藤左千夫 - 『悲しき玩具』を読む	score: 8.07749	path: /cards/000058/files/59823_72664.html
夏目漱石 - 猫の広告文	score: 7.901422	path: /cards/000148/files/4683_9476.html
石田孫太郎 - 猫と色の嗜好	score: 7.778905	path: /cards/000229/files/1153_16151.html
久米正雄 - 「私」小説と「心境」小説	score: 7.5832276	path: /cards/001151/files/59165_74429.html
尾崎放哉 - 俺の記	score: 7.439347	path: /cards/000195/files/43775_25707.html
和辻哲郎 - 漱石の人物	score: 7.38118	path: /cards/001395/files/49898_46721.html
宮本百合子 - 漱石の「行人」について	score: 7.2935934	path: /cards/000311/files/2823_8912.html

最後に

とりあえず、NLPを使う流れをまとめようと書いてみました。あまり精度に関して等は踏み込んでいませんが、こんな記事があったりしますので、踏み込んだ実装の深堀りに関してはこちらを読んでいただくのが良いかと思います。

補足で、実は、 _semantic_search API (ドキュメント, GitHub PR)というのがあって、検索時の煩わしい変換部分をワンステップで検索できるようなもの開発しているようです。これが正式にリリースされると、よりNLPの利用がスムーズでわかりやすくなるのではないかなと思っています。