RAG処理の改善: langchainでハイブリッド検索を実装してみる(勉強メモ)

こちらの記事を読んでいて、ローカル環境だけでセマンティックハイブリッド検索を自分のRAG処理に実装したいなーと思ったのが発端です。

ハイブリッド検索とは

定義が微妙にいろいろあるようですが、キーワード検索とセマンティック検索のような異なる検索手法を組み合わせた検索処理のことを指します。
※ 正確にはDense Retriever(密ベクトルでの検索)とSparse Retriever(疎ベクトルでの検索)の組み合わせらしい。。。？

イメージとしては下記が理解しやすかったです。

ハイブリッド検索をサポートしている著名なクラウドサービスで言えば、Azure Cognitive SearchやPinecone、Weaviateがあります。

今回は上記のクラウドサービスを使わず、Databricks環境下でのハイブリッド検索をlangchainで実装してみます。

私自身NLPにさほど詳しくないため、特にキーワード検索まわりの実装は問題だらけかもしれません。参考にする場合は注意ください。

参考記事

Step1. 環境・データの準備

いつものようにDatabricks上で実装・検証しています。

ノートブックを作成し、必要なモジュールをインストール。(いくらか余計なものが入っています）

%pip install -U -qq transformers accelerate ctranslate2 langchain faiss-cpu sentencepiece rank_bm25 mecab-python3 unidic-lite

dbutils.library.restartPython()

検索に用いるデータは、以前の記事で利用・作成したdolly-15k-jaにチャンクデータ・埋め込みデータを追加したデータセットを使います。

こんな感じで、chunkとembeddingのペアを一列に保持したデータです。

Step2. セマンティック検索用Retrieverの作成

これも以前の記事中で行ったように、FAISSでVectorstoreを作成して、そこからRetrieverを取得します。
embeddingsは事前にインスタンス化したlangchainのEmbeddingsオブジェクトが入ります。
埋め込みモデルはintfloat/multilingual_e5_largeを使っています。

def create_faiss_vectorstore(
    df: DataFrame,
    embeddings: Embeddings,
    index_col: str = "index",
    embeddings_pairs_col: str = "embedding_pairs",
) -> VectorStore:
    """指定されたDataframeからFAISSでVectorStoreを構築"""

    df = df.select(
        index_col, F.explode(embeddings_pairs_col).alias(embeddings_pairs_col)
    )
    pdf = df.toPandas()

    metadata = pdf[[index_col]].to_dict(orient="records")
    text_embedding_pairs = list(
        pdf[embeddings_pairs_col].map(lambda x: (x["chunks"], x["embeddings"]))
    )

    faiss = FAISS.from_embeddings(
        text_embeddings=text_embedding_pairs,
        metadatas=metadata,
        embedding=embeddings,
    )

    return faiss

faiss_retriever = create_faiss_vectorstore(df, embeddings=embeddings).as_retriever(search_kwargs={"k": 5})

Step2. キーワード検索用Retrieverの作成

langchainにはBM25RetrieverというBM25アルゴリズムでの検索を行うRetrieverが提供されています。
(内部的にrank_bm25モジュールを使って実現しています）

※BM25とは↓

とはいえ、大事なのはどちらかと言えば前処理の領域です。
前処理はMecabを使った形態素解析やストップワードの除去など、簡単なものだけ行っています。

なお、今回は前処理をRetriever作成処理内で実施しましたが、事前実施＆保管しておいたほうが処理時間的にはよいです。

import MeCab
import requests

# 日本語ストップワード辞書
url = "https://raw.githubusercontent.com/stopwords-iso/stopwords-ja/master/stopwords-ja.txt"
stopwords_jp = requests.get(url).text.split("\n")

def remove_stopwords(tokens: List[str]) -> List[str]:
    """指定した単語リストからストップワードを除去した結果を返す"""
    tokens = [token for token in tokens if token not in stopwords_jp]
    return tokens


def preprocess_jp(text: str) -> List[str]:
    """日本語テキストを前処理してトークンリストを返す"""

    # 改行コードの除去
    text = text.replace("\n", "")

    # Mecabで特定の品詞のみのトークンリストを作成
    pos_list = ["名詞", "動詞", "形容詞"]
    tagger = MeCab.Tagger()
    node = tagger.parseToNode(text)

    word_list = []
    while node:
        pos = node.feature.split(",")[0]
        if pos in pos_list:
            word = node.surface
            word_list.append(word)
        node = node.next
    return remove_stopwords(word_list)


def create_bm25_retriever(
    df: DataFrame,
    index_col: str = "index",
    embeddings_pairs_col: str = "embedding_pairs",
    bm25_params: Optional[Dict[str, Any]] = None,
    **kwargs: Any,
) -> BaseRetriever:

    # BM25Retrieverのfrom_textsの改変
    # Sparkで前処理した結果を使ってBM25Retrieverを生成する

    try:
        from rank_bm25 import BM25Okapi
    except ImportError:
        raise ImportError(
            "Could not import rank_bm25, please install with `pip install "
            "rank_bm25`."
        )

    # Spark上で自然言語前処理を実施
    @pandas_udf("array<string>")
    def preprocess_chunk(texts: pd.Series) -> pd.Series:
        return texts.apply(preprocess_jp)

    df = (
        df.select(index_col, F.col(f"{embeddings_pairs_col}.chunks"))
        .withColumn("chunks", F.explode("chunks"))
        .withColumn("tokens", preprocess_chunk("chunks"))
    )
    pdf = df.toPandas()

    metadatas = pdf[[index_col]].to_dict(orient="records")
    texts = pdf["chunks"]
    texts_processed = pdf["tokens"]

    bm25_params = bm25_params or {}
    vectorizer = BM25Okapi(texts_processed, **bm25_params)
    metadatas = metadatas or ({} for _ in texts)
    docs = [Document(page_content=t, metadata=m) for t, m in zip(texts, metadatas)]
    return BM25Retriever(
        vectorizer=vectorizer, docs=docs, preprocess_func=preprocess_jp, **kwargs
    )

# BM25を用いたRetriever
bm25_retriever = create_bm25_retriever(df, k=5)

Step3. ハイブリッド検索用Retrieverの作成

langchainはEnsembleRetrieverという複数のRetrieverからの検索結果をReciprocal Rank Fusion(RRF)を用いてマージするRetrieverが提供されています。
そのままこれを使います。

ensemble_retriever = EnsembleRetriever(retrievers=[bm25_retriever, faiss_retriever], weights=[0.5, 0.5])

weightsは半々で設定してみました。

以上でハイブリッド検索用のRetriverが作成できました。

Step4. 検証してみる

単純な検索ワードの結果を比較してみます。

1. セマンティック検索のみ

faiss_retriever.get_relevant_documents("ブドウの木の育て方")

結果(整形しています)

[
    Document(
        page_content="ブドウは湿気の少ない砂地を好むので、水やりの後は必ず地面を乾燥させる。ブドウは、持続的に生産できる量よりもはるかに多くの果実を育てることがよくあります。春先、まだ寒いうちに樹液が出始める前に、積極的に剪定してください。剪定の方法としては、昨年伸びた芽を10～15個ほど残したまま2～4本のつるを残し、今年の芽が伸びるようにジョイントを残しておきます",
        metadata={"index": "1022"},
    ),
    Document(page_content="新しい木のお手入れ方法について教えてください。", metadata={"index": "562"}),
    Document(page_content="ぶどうのお手入れはどうしたらいいのでしょうか？", metadata={"index": "1022"}),
    Document(
        page_content="1.必要な分だけ水を少なめに与え、水を与え過ぎない\n2.1日2時間以上、間接照明に当てておくこと",
        metadata={"index": "3231"},
    ),
    Document(
        page_content="。このような樹木は、一度植えるだけで、定期的に水や肥料で手入れをする必要があります。",
        metadata={"index": "9499"},
    ),
]

セマンティック検索単体でも意味合い的に近いものが取れていますね。

2. キーワード検索のみ

bm25_retriever.get_relevant_documents("ブドウの木の育て方")

結果(整形しています)

[
    Document(
        page_content="。木や蔓から落ちたブドウ、ナツメヤシ、イチジクなどは、炎天下で乾燥させることができた。木や蔓から落ちたブドウやナツメヤシ、イチジクなどが炎天下で乾燥し、食べられるようになるのを狩猟採集民は観察し、その安定性と凝縮された甘さを評価した。",
        metadata={"index": "3590"},
    ),
    Document(
        page_content="。木や蔓から落ちたブドウ、ナツメヤシ、イチジクなどは、炎天下で乾燥させることができた。木や蔓から落ちたブドウやナツメヤシ、イチジクなどが炎天下で乾燥し、食べられるようになるのを狩猟採集民は観察し、その安定性と凝縮された甘さを評価した。",
        metadata={"index": "5547"},
    ),
    Document(
        page_content="1680年、スペインのイエズス会宣教師が、宗教行事で使うワインを作るためのブドウを育てるため、地中海のブドウの木を植え、カリフォルニアのワイン造りの歴史を築きました。",
        metadata={"index": "4194"},
    ),
    Document(
        page_content="。木や蔓から落ちたブドウやナツメヤシ、イチジクなどが炎天下で乾燥し、食べられるようになるのを狩猟採集民は観察し、その安定性と凝縮された甘さを評価した。",
        metadata={"index": "4403"},
    ),
    Document(
        page_content="。木や蔓から落ちたブドウやナツメヤシ、イチジクなどが炎天下で乾燥し、食べられるようになるのを狩猟採集民は観察し、その安定性と凝縮された甘さを評価した。",
        metadata={"index": "7806"},
    ),
]

ブドウというキーワードで引っ掛かっていますね。
※　同じ文言が複数出ているのは不具合ではなく、そういうデータセットのためです。

3. ハイブリッド検索

ensemble_retriever.get_relevant_documents("ブドウの木の育て方")

結果(整形しています)

[
    Document(
        page_content="。木や蔓から落ちたブドウ、ナツメヤシ、イチジクなどは、炎天下で乾燥させることができた。木や蔓から落ちたブドウやナツメヤシ、イチジクなどが炎天下で乾燥し、食べられるようになるのを狩猟採集民は観察し、その安定性と凝縮された甘さを評価した。",
        metadata={"index": "5547"},
    ),
    Document(
        page_content="。木や蔓から落ちたブドウやナツメヤシ、イチジクなどが炎天下で乾燥し、食べられるようになるのを狩猟採集民は観察し、その安定性と凝縮された甘さを評価した。",
        metadata={"index": "7806"},
    ),
    Document(
        page_content="ブドウは湿気の少ない砂地を好むので、水やりの後は必ず地面を乾燥させる。ブドウは、持続的に生産できる量よりもはるかに多くの果実を育てることがよくあります。春先、まだ寒いうちに樹液が出始める前に、積極的に剪定してください。剪定の方法としては、昨年伸びた芽を10～15個ほど残したまま2～4本のつるを残し、今年の芽が伸びるようにジョイントを残しておきます",
        metadata={"index": "1022"},
    ),
    Document(page_content="新しい木のお手入れ方法について教えてください。", metadata={"index": "562"}),
    Document(
        page_content="1680年、スペインのイエズス会宣教師が、宗教行事で使うワインを作るためのブドウを育てるため、地中海のブドウの木を植え、カリフォルニアのワイン造りの歴史を築きました。",
        metadata={"index": "4194"},
    ),
    Document(page_content="ぶどうのお手入れはどうしたらいいのでしょうか？", metadata={"index": "1022"}),
    Document(
        page_content="1.必要な分だけ水を少なめに与え、水を与え過ぎない\n2.1日2時間以上、間接照明に当てておくこと",
        metadata={"index": "3231"},
    ),
    Document(
        page_content="。このような樹木は、一度植えるだけで、定期的に水や肥料で手入れをする必要があります。",
        metadata={"index": "9499"},
    ),
]

セマンティック検索とキーワード検索の結果が混ざったような結果が得られました。
※ 件数増えてしまっていますね。

ハイブリッド検索を取り入れることで、RAGにおける回答精度を上げられる（かもしれない）こと期待できると思います。

今後の改善をするとしたら

検索結果をRe-rankingさせることで、より質問に関係するコンテキストのみに絞るやり方が考えられます。
(Azure Cognitive Searchのセマンティックハイブリッド検索でも用いられる考え方）

Re-rankingはざっくり言えば、得られた各検索結果が当初の質問に答えられる内容かどうかを言語モデルに判断・スコアリングさせ、そのスコアが高いもののみを抜き出す手法です。
langchainでもLLMを使った手法がカバーされていますし、コード自体は難しくないのですが、処理時間がかかる＋モデルによって結果がイマイチ安定しないため、もう少し勉強が必要そうです。

まとめ

いろいろ詰めるところが多い（特に前処理周り）ですが、ハイブリッドぽい検索は実現できたかなと思います。
RAGのパフォーマンス改善で言えば、チャンク分割の仕方とか他にも工夫するべきポイントが多数ありますが、ハイブリッド検索も取り入れてやってみるといいかなと思いました。

しかし、やがてLakehouse AIがリリースされて、Databricksの利用においては簡単にできるようになるんだろうなあ。。。