エムスリーテックブログ

エムスリー(m3)のエンジニア・開発メンバーによる技術ブログです

日本語埋め込みモデルRuriを使ったBM42 on Elasticsearchと形態素解析器Sudachiによるトークン矯正

Elasticsearch ML 情報検索 python

こちらはエムスリー Advent Calendar 2024 1日目の記事です。

エムスリーエンジニアリンググループ AI・機械学習チームでソフトウェアエンジニアをしている中村(po3rin) です。

今回はQdrantが開発した新しいスコアリングアルゴリズムであるBM42を簡単に紹介し、それをElasticsearch上で構築する方法とその所感をお話しします。さらに形態素解析器のSudachiを使って類似語展開やトークン修正を行ない、BM42の精度を矯正する方法を試したのでその紹介をします。

BM42の紹介に関してはQdrantの記事が最も詳しいですが、このブログでも導入として簡単に紹介します。

qdrant.tech

BM25の弱点

BM25は検索においてクエリに関連する結果をスコアリングするために使用されます。セマンティック検索が活躍する現在でもBM25と組み合わせたハイブリッド検索は主力の方法です。

BM25の式をおさらいしておきましょう。ターム q_1 ... q_nを含むクエリ Qが与えられたときBM25のスコアは次のようになります。

\mathrm{score}(q, D) = \sum_i \mathrm{idf}(q_i)\times\frac{(k_1+1)f(q_i, D)}{f(q_i, D)+k_1(1-b+b\frac{|D|}{\mathrm{avg}(dl)})}

2024/12/01現在、環境によってはHantena BlogのTex数式の表示が乱れてしまっているようです。下記の記事で修正の仕方を解説して下さっています。

syleir.hatenablog.com

ここで idf(q)は qの逆ドキュメント頻度、 f(q, D)はドキュメント Dにおける qのターム頻度です。詳細な導入などは書籍「情報検索 :検索エンジンの実装と評価」が詳しいのでそちらをお勧めします。

情報検索 :検索エンジンの実装と評価

情報検索 :検索エンジンの実装と評価

  • 森北出版
Amazon

BM25は計算式にターム頻度を持つため、ドキュメントが十分に長く、重要語が複数回出現する場合にうまく機能します。しかし、RAGなどチャンク分割に対する検索や、短い記事タイトルに対して検索するような場合は、重要語であってもターム頻度が1回である場合が多く、BM25は上手く機能しません。特にチャンク分割の検索の場合は、チャンクの長さに差が発生しないので、BM25の式で機能するのが idf(q)だけになってしまう場合が多いのが現状です。

よって短いドキュメントやチャンクの検索ではドキュメント内の用語の重要度を再検討する必要が出てきます。そこで開発されたのがBM42です。

BM42とは

BM42はIDFとTransformerのAttentionを利用します。

Transformerモデルでは、ドキュメント内のmaskされたトークンを予測するようにトレーニングされているため、Attention行列はそのmaskしたトークンを予測するのに各トークンがどれだけ寄与するか(Attention weight)を表します。

Attention行列(Qdrantの記事「https://qdrant.tech/articles/bm42/」からの引用)

トークンの中でもTransformerモデルの [CLS]トークン は入力されたテキスト全体の要約や分類として利用できるように学習するため、[CLS]トークンはドキュメント文脈全体を代表するトークンとなっています。そのため、Attention行列の[CLS]トークン列を確認するとドキュメントの文脈全体に対する各トークンの重要性を取得できます。このトークンの重要度をBM25のターム頻度の代わりに利用するのがBM42です。

\mathrm{score}(Q, D) = \sum_i \mathrm{idf}(q_i) \times\ \mathrm{Attention}(CLS, q_i)

例として名古屋大学の塚越さんが公開している言語モデルであるRuriを使って文脈全体に対するトークンの重要度を確認してみます(retokenizationの部分はこちらの記事を大いに参考にしました。ありがとうございます)。

huggingface.co

import torch
from transformers import AutoTokenizer, AutoModel

model_name = "cl-nagoya/ruri-large"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)


def get_token_attentions(text) -> dict[str, float]:
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
    with torch.no_grad():
        outputs = model(**inputs, output_attentions=True)

    attentions = outputs.attentions[-1][0, :, 0].mean(dim=0)
    #                                â–²  â–²     â–²                                 
    #                                │  │     └─── [CLS] token is the first one
    #                                │  └─────── First item of the batch         
    #                                └────────── Last transformer layer  
    tokens = tokenizer.convert_ids_to_tokens(inputs['input_ids'][0])

    token_attentions = {}
    current_word = ""
    current_weight = 0

    # retokenization
    for token, weight in zip(tokens[1:-1], attentions[1:-1]):
        if token.startswith("##"):
            current_word += token[2:]
            current_weight += weight
            continue

        if current_word:
            token_attentions[current_word] = current_weight
        current_word = token
        current_weight = weight.item()

    if current_word:
        token_attentions[current_word] = float(current_weight)

    return token_attentions

result = get_token_attentions("qdrantが開発した新しいランキングアルゴリズムであるBM42を試します。")
for k, v in result.items():
    print(f"{k}: {v:.4f}")

retokenizationのフェーズではQdrantの記事で紹介されていたのと同じように、分割されてしまったサブワードを単語に再びマージしています。その際のAttention weightはサブワードのAttention weightを合計となります。Attention weightはドキュメントの全単語の合計が1になるように正規化されてるので、そのドキュメントにおけるその単語の重要度という意味で足し合わせは可能です。

結果は次のように取得できます。

ランキング: 0.2108
qdrant: 0.1261
試し: 0.1148
アルゴリズム: 0.1063
42: 0.0750
ます: 0.0523
ã‚’: 0.0492
。: 0.0436
BM: 0.0301
ある: 0.0272
新しい: 0.0171
開発: 0.0149
で: 0.0058
た: 0.0056
が: 0.0054
し: 0.0051

この値を Attention(CLS, q)としてBM42を計算します。

Qdrantの記事ではSPLADEとの比較やメリデメなども説明されているので、もし興味があればそちらを参照ください。

BM42をElassticsearchで動かす

当然ながらElasticsearchにはBM42は実装されていないので、これをElasticsearchに組み込むにはどうすれば良いかを検討します。Elasticsearchではスパースベクトル検索がサポートされているのでそちらを使います。しかし後で説明する通り、ElasticsearchはQdrantと違いスパースベクトル検索のスコアとIDFの積をとる直接的な方法がないので、少し工夫する必要はあります。

ますはスパースベクトル検索用のmappingを用意します。

{
  "mappings": {
    "properties": {
      "title": {
        "type": "text",
        "analyzer": "whitespace"
      },
      "joined_tokens": {
        "type": "text",
        "analyzer": "whitespace"
      },
      "tokens": {
        "type": "sparse_vector"
      }
    }
  }
}

フィールドには今回はシンプルにタイトルそのままの文字列を格納するtitle、タイトルをtokenizeした結果をスペースでjoinしたjoined_tokens。スパースベクトルtokensの3つ用意します。JSONからElasticsearchのインデックスを作成した後に、実際にデータを入れていきます。

client = Elasticsearch("http://localhost:9200/")

texts = [
    "qdrantが開発した新しいランキングアルゴリズムであるBM42を試します。",
    "検索ランキングで使われるBM25とは？"
]

for i, t in enumerate(texts):
    tokens = get_token_attentions(t)
    joined_text = ' '.join(tokens.keys())
    doc = {
        "title": t,
        "joined_tokens": joined_text,
        "tokens": tokens
    }
    resp = client.index(index="test-index", id=i+1, document=doc)

今回は後にスコア計算の確認をするためにドキュメントを2つだけ用意しています。これで検索する準備が整いました。タームが1つの場合は次のクエリでBM42の検索ができます。

GET /test-index/_search
{
    "query": {
        "script_score": {
            "query": {
                "bool": {
                    "filter": {
                        "match": {
                            "joined_tokens": "BM"
                        }
                    },
                    "should": [
                        {
                            "term": {
                                "tokens": {
                                    "value": "BM"
                                }
                            }
                        }
                    ]
                }
            },
            "script": {
                "source": "return _score / _termStats.docFreq().getSum() "
            }
        }
    }
}

あまり美しくないですが、クエリごとにスパースベクトル検索のスコアとIDFの積をとるのがElasticsearchだと少し難しいためにこのようなクエリになっています。

script_scoreではquery内で使った複数タームの統計値(IDFの平均やTFの合計など)しか取れないので、filterを挟んでそこでタームを1つに絞っています。実際にfilterクエリを外すと_termStats.docFreq().getSum()の部分が0になります。

script_scoreではreturn _score / _termStats.docFreq().getSum()というスクリプトでタームごとのスコアを計算しています。_scoreはスパースベクトル検索の結果のスコア(BM42ではAttention weight)、_termStats.docFreq().getSum()はfilterクエリで使ったタームのドキュメント頻度の合計値です。filter内ではタームを1つに絞っているので、_termStats.docFreq().getSum()はそのタームのドキュメント頻度そのものを表現しています。

ここで、 idf(q_i)の計算に必要な全ドキュメント数を使用していないのは、script_score内で参照できないためです。 idf(q_i)の定義とは違いますが、検索用途でスコアを並び替えるだけであれば、全アイテムに共通の値である全ドキュメント数は無視できます。

スコアの詳細を確認するために、同じbodyを_explainエンドポイントに投げてみます。

GET test-index/_explain/1

そうするとスパースベクトルのスコアとドキュメント頻度の逆数の積をとっていることが分かります。

{
    "_index": "test-index",
    "_id": "1",
    "matched": true,
    "explanation": {
        "value": 0.015045166,
        "description": "script score function, computed with script:\"Script{type=inline, lang='painless', idOrCode='return _score / _termStats.docFreq().getSum()', options={}, params={}}\"",
        "details": [
            {
                "value": 0.030090332,
                "description": "_score: ",
                "details": [
                    {
                        "value": 0.030090332,
                        "description": "sum of:",
                        "details": [
                            {
                                "value": 0.030090332,
                                "description": "Linear function on the tokens field for the BM feature, computed as w * S from:",
                                "details": [
                                    {
                                        "value": 1,
                                        "description": "w, weight of this function",
                                        "details": []
                                    },
                                    {
                                        "value": 0.030090332,
                                        "description": "S, feature value",
                                        "details": []
                                    }
                                ]
                            },
                            {
                                "value": 0,
                                "description": "match on required clause, product of:",
                                "details": [
                                    {
                                        "value": 0,
                                        "description": "# clause",
                                        "details": []
                                    },
                                    {
                                        "value": 1,
                                        "description": "joined_tokens:BM",
                                        "details": []
                                    }
                                ]
                            }
                        ]
                    }
                ]
            }
        ]
    }
}

複数タームの場合は、次のようにshouldで繋げてsumを取ります。

GET /test-index/_search

{
    "query": {
        "bool": {
            "should": [
                {
                    "script_score": {
                        "query": {
                            "bool": {
                                "filter": {
                                    "match": {
                                        "joined_tokens": "BM"
                                    }
                                },
                                "should": [
                                    {
                                        "term": {
                                            "tokens": {
                                                "value": "BM"
                                            }
                                        }
                                    }
                                ]
                            }
                        },
                        "script": {
                            "source": "return _score / _termStats.docFreq().getSum() "
                        }
                    }
                },
                {
                    "script_score": {
                        "query": {
                            "bool": {
                                "filter": {
                                    "match": {
                                        "joined_tokens": "検索"
                                    }
                                },
                                "should": [
                                    {
                                        "term": {
                                            "tokens": {
                                                "value": "検索"
                                            }
                                        }
                                    }
                                ]
                            }
                        },
                        "script": {
                            "source": "return _score / _termStats.docFreq().getSum() "
                        }
                    }
                }
            ]
        }
    }
}

shouldクエリは各クエリのスコアの合計を取るので、これでBM42のスコア計算ができます。

このように少し周りくどいクエリじゃないと現状BM42をElasticsearch上で再現できないのが問題です。また、TopKクエリ処理最適化ができないので、普通のスコア計算よりもパフォーマンスが悪いです。TopKクエリ処理最適化に関しては次の記事が詳しいです。

engineering.mercari.com

最適解はカスタムPluginを実装することだと思いますが、今回は一番シンプルな実装方法を紹介しました。(時間があれば実装にチャレンジしてみようと思います)

Sudachiによる矯正

ここまででBM42をElasticsearch上で動かす方法を紹介しましたが、まだ次の気になる点があります。

  • モデルによっては意図しないトークンが生成される問題
  • 表記揺れ、シノニムが吸収できない問題

モデルによっては意図しないトークンが生成される問題

日本語埋め込みモデルだと、特にエムスリーがメインで扱う医療用語は細かく分割されるので、意図しないヒットを生成する可能性があります。特に医療検索で顕著なのは漢方の名前です。例えば「半夏厚朴湯と柴胡加竜骨牡蛎湯の併用」というドキュメントは漢方名がトークンに分解されます。

example_text = "半夏厚朴湯と柴胡加竜骨牡蛎湯の併用"
result = get_token_attentions(example_text)
for k, v in sorted(result.items(), key = lambda fruit : fruit[1], reverse=True):
    print(f"{k}: {v:.4f}")

結果は次のようになります。

併用: 0.1649
半夏: 0.1020
湯: 0.0594
厚朴: 0.0523
柴胡: 0.0436
牡蛎: 0.0323
竜骨: 0.0310
加: 0.0213
の: 0.0156
と: 0.0149

そのため、このような名詞は辞書でさらにトークンを結合してあげたくなります。qdrantでは##プレフィックスを持つトークンを結合していますが、##プレフィックスを持たないトークンの結合を考える場合は、分割したくない単語を辞書で持っておくと良さそうです。エムスリーではSudachiを使った形態素解析による検索が既に動いているので、その辞書による形態素解析の結果使ってトークン結合することを考えます。エムスリーでのSudachi利用に関しては過去の記事で紹介しています。

www.m3tech.blog

下記はSudachiでトークンをSudachiの形態素解析の結果を元に結合する例です。

def retokenize_with_sudachi(tokens, text):
    """
    Sudachiの形態素解析結果を元に、トークン化を結合する
    """

    tokenizer_obj = dictionary.Dictionary(config_path="./sudachi.json", dict_type="core").create()  
    sudachi_tokens = [m.surface() for m in tokenizer_obj.tokenize(text, mode)]

    result = {}
    for token in sudachi_tokens:
        result[token] = 0
        for t in tokens:
            if t in token:
                result[token] += float(tokens[t])

    return result

これでSudachiの形態素解析の結果を使ったトークン結合ができます。

example_text = "半夏厚朴湯と柴胡加竜骨牡蛎湯の併用"
result = get_token_attentions(example_text)
for k, v in sorted(result.items(), key = lambda item : item[1], reverse=True):
    print(f"{k}: {v:.4f}")

print("----------")
result = retokenize_with_sudachi(result, example_text)
for k, v in sorted(result.items(), key = lambda item : item[1], reverse=True):
    print(f"{k}: {v:.4f}")

結果は次のようになります。

併用: 0.1649
半夏: 0.1020
湯: 0.0594
厚朴: 0.0523
柴胡: 0.0436
牡蛎: 0.0323
竜骨: 0.0310
加: 0.0213
の: 0.0156
と: 0.0149
----------
半夏厚朴湯: 0.2138
柴胡加竜骨牡蛎湯: 0.1876
併用: 0.1649
の: 0.0156
と: 0.0149

上記の実装では同じトークンが二回出てくることを考慮していないなどの改良の余地はありますが、これで絶対に分解したくないトークンを結合することができました。

表記揺れ、シノニムが吸収できない問題

2点目については、類義語をシノニム辞書に定義して管理しているチームもあると思います。しかし、BM42ではモデルのTokenizerを利用するのでシノニムは無視されます。SPLADEであれば文脈に基づく関連性の高い単語もベクトルに含めることができますが、意図しないトークンも追加され、検索時におかしなヒットをしてしまう可能性があります。

例えばYuichi Tatenoさんが公開している日本語のSPLADEモデルを使って「Qdrantが開発した新しいランキングアルゴリズムであるBM42を試します」の出力を見ると次のようになります(デモ画面のスクリーンショットを掲載します)。

huggingface.co

日本語SPLADEによるスパースベクトル変換

「試す」を拡張して「挑戦」というトークンが含まれたりなど、上手く機能しているように見えますが、「製品」や「秋山」など無駄なトークンも含まれてしまっています。

そこで既にあるシノニム辞書でトークンを同じスコアで拡張することを考えます。

ダブル配列でシノニムを格納し、もしトークンが検索でヒットすればシノニムとしてトークンを同じスコアで追加します。今回はkey-valueのようにシンプルに使えるダブル配列のPython実装のPydatrieを利用しました。

github.com

from pydatrie import DoubleArrayTrie


synonyms = DoubleArrayTrie(
    {
        "ばね指": "弾発指",
        "弾発指": "ばね指"
    }
)

def token_expantion(tokens) -> dict[str, float]:
    """
    トークンの類義語を追加する
    """

    result = {}
    for k, v in tokens.items():
        result[k] = v
        syn = synonyms.get(k)
        if syn is not None:
            result[syn] = v
    return result

example_text = "ばね指の症状について"
result = get_token_attentions(example_text)
for k, v in sorted(result.items(), key = lambda item : item[1], reverse=True):
    print(f"{k}: {v:.4f}")

print("----------")
result = retokenize_with_sudachi(result, example_text)
for k, v in sorted(result.items(), key = lambda item : item[1], reverse=True):
    print(f"{k}: {v:.4f}")

print("----------")
result = token_expantion(result)
for k, v in sorted(result.items(), key = lambda item : item[1], reverse=True):
    print(f"{k}: {v:.4f}")

結果は次のようになります。

ばね指: 0.5203
症状: 0.1462
の: 0.0684
に: 0.0675
つい: 0.0506
----------
ばね指: 0.5203
症状: 0.1462
の: 0.0684
に: 0.0675
つい: 0.0506
て: 0.0000
----------
ばね指: 0.5203
弾発指: 0.5203
症状: 0.1462
の: 0.0684
に: 0.0675
つい: 0.0506
て: 0.0000

こちらで用意したシノニム辞書によるトークン拡張ができました。他にもstopwordの削除や、トークンのnormalizeなどもの改良点はありそうです。

まとめ

今回はElasticsearch上でBM42を構築する方法に加え、Sudachiを使ってトークン結合をする方法や、シノニムによるトークン展開をする方法を紹介しました。それっぽくいい感じにBM42が利用できるところまで確認できましたが、ユーザーのクエリを受けた際に、Transformerの推論に加え、形態素解析を挟んだり、シノニムの検索をする必要があるので、パフォーマンスの面でまだ課題がありそうです。またElasticsearchのTopKクエリ処理最適化もできない状況なので、Elasticsearchのパフォーマンスにも影響を与えるため、実際に実務で利用を考える場合はまだまだ課題がありそうです。

当然、BM42はハイブリッド検索と合わせて使うことが理想なので、実務利用の場合はスパースベクトル検索の結果と合わせて、密ベクトル検索やBM25での結果をRRFで混ぜ合わせると良いでしょう。

We are hiring !!

エムスリーでは検索や推薦が大好きなエンジニアを募集しています！少しでも興味がある方は、次のURLからカジュアル面談をご応募ください！

jobs.m3.com