検索エンジンPyTerrierを使った日本語検索パイプラインの実装

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

今回は社内でPyTerrierを採用して文書検索BatchをPythonで実装したので、PyTerrierの紹介とPyTerrierで日本語検索を実装する方法を紹介します(日本語でPyTerrierを扱う記事は多分初?)。

PyTerrierとは

Terrierのロゴ

PyTerrierは、Pythonでの情報検索実験のためのプラットフォームです。 JavaベースのTerrierを内部的に使用して、インデックス作成と検索操作を行うことができます。基本的なQuery RewritingやBM25などの各種スコアリングがすぐに使え、また学習済みモデルの組み込みや評価なども簡単にできるため、開発と評価を一気通貫で行うことが可能です。

ECIR2021ではLearning to rankの実験などPyTerrierで行うチュートリアルが公開されています。

github.com

パイプラインを演算子で構築できるのが特徴で、例えば、TF-IDFで100件取ってきて、BM25でリランキングするパイプラインは下記のように宣言的に実装できます。

tfidf = pt.BatchRetrieve(index, wmodel="TF_IDF")
bm25 = pt.BatchRetrieve(index, wmodel="BM25")
pipeline = (tfidf % 100) >> bm25

パイプラインの評価もすぐに行うことができます。例えば下記はTF-IDFとBM25の比較をmap(Mean Average Precision)メトリクスで行う例です。

pt.Experiment([tf_idf, bm25], topic, qrels, eval_metrics=["map"])

このようにPyTerrierでは情報検索の実験環境としても非常に優れたインターフェースを提供しています。

弊社でのPyTerrier利用

社内で、「数十万件のtermリストが記事に出現するかをオフラインで確認する」というタスクを実装することになり、その中でPyTerrierを使ってみることとしました。

アーキテクチャ

ちなみに弊社で日々使っているElasticsearchを使ってしまうというのも候補としてありましたが、Elasticsearchを利用するとコア処理のテストがミドルウェアに依存することになり、動作確認のたびにESを立てる、落とすなどの面倒な処理が必要なため、今回は見送りました。

他にも、termが長い時に別の手法を使って高速化していたりもするのですが、それに関しては別の記事で詳細を説明したいと思います。

PyTerrierで日本語検索

PyTerrierで日本語検索をする際には少しコツが必要です。PyTerrierで用意しているTokenizerに日本語の形態素解析はないので、自前で用意してあげる必要があります。

PyTerrierで英語以外の検索例が公開されているので、これも参考にしてください。

colab.research.google.com

今回はSudachiで形態素解析して、PyTerrierで検索する方法を紹介します。Sudachiの紹介やSudachiをElasticsearchに導入した記事を弊社から公開しているので、Sudachiに興味のある方は是非そちらもご覧ください。

www.m3tech.blog

早速PyTerrierで日本語検索をする方法を紹介していきます。モジュールは下記を用意します。また、PyTerrierのcoreはJavaで実装されているので、Javaの環境も用意しておきましょう。

import os

import pyterrier as pt
import pandas as pd
from sudachipy import dictionary, tokenizer

PyTerrierを初期化します。

if not pt.started():
  pt.init()

今回検索する対象のドキュメントを用意しておきます。

df = pd.DataFrame([
        ["d1", "検索方法の検討"]
    ], columns=["docno", "text"])

PyTerrierはPandasのDataFrameをそのままインデックスするインターフェースが用意されているので便利です。 ドキュメントとクエリの両方を形態素解析するので、それぞれTokenizerを用意してあげます。ドキュメントは品詞でインデックスするタームを絞ります。

class DocTokenizer():
    tokenizer_obj = dictionary.Dictionary().create()
    mode = tokenizer.Tokenizer.SplitMode.C

    def tokenize(self, txt: str) -> list[str]:
        return [
            m.dictionary_form() for m in self.tokenizer_obj.tokenize(txt, self.mode)
            if len(set(['名詞', '動詞', '形容詞', '副詞', '形状詞']) & set(m.part_of_speech())) != 0
        ]

class TokenizeDoc():
    tokenizer = DocTokenizer()

    def tokenize(self, df: pd.DataFrame):
        df['tokens'] = df['text'].apply(lambda x: ' '.join(self.tokenizer.tokenize(x)))
        return df

これで事前にドキュメントをタームに分割する用意ができました。ドキュメントのDataFrameをTokenizeします。

doc_tokenizer = TokenizeDoc()
phrase_query_converter = PhraseQueryConverter()

df = doc_tokenizer.tokenize(df=df)
df

#  docno   text          tokens
#   d1     検索方法の検討   検索 方法 検討

これでドキュメントの準備ができたので、実際にIndex処理を行います。日本語の場合はスペースで区切られるUTFTokeniserを利用します。事前にドキュメントをタームのスペース区切りにしてあるので、そのまま渡してあげればインデックス完了です。

indexer = pt.DFIndexer('./askd-terrier', overwrite=True, blocks=True)
indexer.setProperty('tokeniser', 'UTFTokeniser')
indexer.setProperty('termpipelines', '')
index_ref = indexer.index(df['tokens'], docno=df['docno'])
index = pt.IndexFactory.of(index_ref)

後はクエリの処理です。PyTerrierではクエリ言語をサポートしており、And検索やPhrase検索が可能です。例えばAnd検索は+term1 +term2のように記述でき、Phrase検索は"term1 term2"のように記述できます。その他の記述方法はドキュメントをご覧ください。

http://terrier.org/docs/v5.1/querylanguage.html

今回はPhrase検索を使ってみます。形態素解析したクエリをフレーズクエリ言語に展開する実装です。

class QueryTokenizer():
    tokenizer_obj = dictionary.Dictionary().create()
    mode = tokenizer.Tokenizer.SplitMode.C

    def tokenize(self, txt: str) -> list[str]:
        return [m.surface() for m in self.tokenizer_obj.tokenize(txt, self.mode)]

class PhraseQueryConverter():
    query_tokenizer = QueryTokenizer()

    def convert(self, text: str) -> str:
        tokens = [t for t in self.query_tokenizer.tokenize(text)]
        if len(tokens) <= 1:
            return text
        joined = ' '.join(tokens)
        return f'"{joined}"'

クエリを処理する準備ができたので、実際に検索パイプラインを実装します。今回はクエリをフレーズクエリに変換して、BM25でスコアリングして上位100件を取得するパイプラインを用意しました。

pipe = (pt.apply.query(lambda row: phrase_query_converter.convert(row.query)) >> \ 
        (pt.BatchRetrieve(index, wmodel='BM25') % 100).compile())

compile()は検索パイプラインのDAGを書き換えて最適化してくれます。例えばcompile無しだとクエリにヒットするドキュメントを全件とってきて、BM25でスコアリングして上位100件を取得します。一方でcompile()を行うとLuceneでも採用されているBlock Max WANDなどの動的プルーニング手法に書き換えられ、検索がより高速になります。compileによる最適化についてはこちらの論文が詳しいです。

これで検索パイプラインを実装する準備ができました。

res = pipe.search('検索方法')
res

#  qid docid   docno   rank    score   query_0 query
# 0    1   0   d1  0   -1.584963   検索方法    "検索 方法"

ヒットしたドキュメントのIDとともにrankã‚„scoreが返ってきます。また、query_0には元のクエリ、queryには実際に検索が走ったクエリが結果に記載されます。もちろんフレーズクエリに書き換えているので、検索検討などのクエリにはヒットしません。

res = pipe.search('検索検討')
res

# empty...

Phrase Queryの注意点

現在Issueにあげているのですが、フレーズ検索のタームがインデックスされていないものだと、そのタームを無視して検索をする挙動を発見しました。

github.com

具体的には、今回の例で言うと、下記のようなクエリでもフレーズクエリでヒットしてしまいます。

res = pipe.search('検索専門')
res

#  qid docid   docno   rank    score   query_0 query
# 0    1   0   d1  0   -1.584963   検索専門    "検索 専門"

直近のできる対応としては、インデックスされているタームをチェックして、もし存在しないなら、そのままのクエリを投げることでヒットを防ぐなどの対応が考えられます。

def convert(self, text: str, lexicon) -> str:
    tokens = [t for t in self.query_tokenizer.tokenize(text)]

    if len(tokens) <= 1:
            return text

    # indexed tokens inculde query term (bug?: phrase query ignore non indexed term)
    for t in tokens:
        if lexicon.getLexiconEntry(t) is None:
            return text

    joined = ' '.join(tokens)
    return f'"{joined}"'


lex = index.getLexicon()

pipe = (pt.apply.query(lambda row: phrase_query_converter.convert(row.query, lex)) >> pt.BatchRetrieve(index, wmodel='BM25').compile())

弊社ではフレーズクエリが必要だったので、一旦この方法で対応しています。根本の原因は現在調査中です。

まとめ

PyTerrierの紹介と、PyTerrierで日本語検索する方法を簡単に紹介しました。Pythonでサクッと検索したい時には便利です。一方で、PyTerrierは今回のようなLexical Searchにとどまらず、情報検索モデルの適用や、実験の評価などでも活躍するので、興味のある方は是非触ってみてください。個人的にはECIR2021のチュートリアルが非常に良い入門になりました。

https://github.com/terrier-org/ecir2021tutorial

We're hiring !!!

エムスリーでは検索&推薦基盤の開発&改善を通して医療を前進させるエンジニアを募集しています！社内では日々検索や推薦についての議論が活発に行われています。

「ちょっと話を聞いてみたいかも」という人はこちらから！ jobs.m3.com