LlamaIndex v0.10のQueryEngineを使ってハイブリッド検索をやってみる

Aya*


2024/05/14

こんにちは。今回は，前回使用したQueryEngineを使ってハイブリッド検索実装にトライします。

今回はこちらのドキュメントをベースに紹介していきます。

はじめに

ここでは，ハイブリッド検索の非常に単純なバージョンを定義する方法を紹介します．

「AND」や「OR」条件を使用して，キーワードルックアップ検索とベクトル検索を組み合わせます．

今回は，日本語に対応させるために同じドキュメントに対して２つのノードを作成します．

設定

LlamaIndexのインストール

!pip install llama-index

sudachiのインストール
今回は日本語トークナイザーにsudachiを使用します．mecabも有名です．

!pip install -q sudachipy
!pip install -q sudachidict_core
import os

os.environ["OPENAI_API_KEY"] = "sk-"

データのダウンロード

トヨタのパンフレットPDFを読み込みます．

import requests
import os

def download_file(url, directory, filename):
    # ディレクトリが存在しない場合は作成
    if not os.path.exists(directory):
        os.makedirs(directory)

    # ファイルのパスを構築
    file_path = os.path.join(directory, filename)

    # ファイルをダウンロード
    response = requests.get(url)
    response.raise_for_status()  # ステータスコードが200以外の場合はエラーを発生させる

    # ファイルを保存
    with open(file_path, 'wb') as file:
        file.write(response.content)

    print(f"ファイルがダウンロードされ、{file_path}に保存されました。")

# 使用例
url = 'https://toyota.jp/pages/contents/corollasport/001_p_001/pdf/spec/corollasport_spec_201806.pdf'  # ダウンロードしたいファイルのURL
directory = 'data'  # 保存したいディレクトリのパス
filename = 'corollasport_spec_old.pdf'  # 保存するファイルの名前
download_file(url, directory, filename)

url = 'https://toyota.jp/pages/contents/corollasport/001_p_001/5.0/pdf/spec/corollasport_spec_202404.pdf'
filename = 'corollasport_spec_new.pdf'
download_file(url, directory, filename)

url = 'https://toyota.jp/pages/contents/corollatouring/001_p_001/5.0/pdf/spec/corollatouring_spec_202404.pdf'
filename = 'corollatouring_spec_new.pdf'
download_file(url, directory, filename)

データの読み込み

ドキュメントを読み込みます

# LlamaIndexのインポート
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
# dataフォルダ内の学習データを使い、インデックスを生成する
documents = SimpleDirectoryReader('data').load_data()

ノードの作成

from llama_index.core import Settings
nodes = Settings.node_parser.get_nodes_from_documents(documents)

ストレージコンテキストを作成します

from llama_index.core import StorageContext

# initialize storage context (by default it's in-memory)
storage_context = StorageContext.from_defaults()
storage_context.docstore.add_documents(nodes)

ドキュメントのトークナイズ

ここで，キーワード検索に使用するためにノードを処理します．

# sudachi

from sudachipy import tokenizer
from sudachipy import dictionary

# トークナイザーの設定
tokenizer_obj = dictionary.Dictionary().create()

# トークナイズの関数
def tokenize_text(text):
    # モード設定: A - 短い形、B - 中間形、C - 長い形
    mode = tokenizer.Tokenizer.SplitMode.C
    # トークン化
    tokens = tokenizer_obj.tokenize(text, mode)
    # 表層形を取得してスペースで結合
    tokenized_text = ' '.join([token.surface() for token in tokens])
    return tokenized_text

整形に使用する関数の定義

import re
# 日本語の間にあるスペースを削除する
def rm_space(text):
    # Regex pattern to match spaces where both sides are Japanese characters using specific Unicode ranges
    # Hiragana (U+3040-U+309F), Katakana (U+30A0-U+30FF), and common Kanji (U+4E00-U+9FAF)
    return re.sub(r'(?<=[\u3040-\u309F\u30A0-\u30FF\u4E00-\u9FAF])\s+(?=[\u3040-\u309F\u30A0-\u30FF\u4E00-\u9FAF])', '', text)

# 連続するスペースを削除する
def rm_multi_space(text):
    return re.sub(r'\s{2,}', ' ', text)

ノードをコピーしてキーワード検索ノードを作成

tnodes = nodes.copy()
for node in tnodes:
    node.text = rm_multi_space(tokenize_text(rm_space(node.text)))
tnodes[0].text

同じデータに対するベクトル インデックスとキーワード テーブル インデックスを定義する

from llama_index.core import SimpleKeywordTableIndex, VectorStoreIndex

vector_index = VectorStoreIndex(nodes, storage_context=storage_context)
keyword_index = SimpleKeywordTableIndex(tnodes, storage_context=storage_context)

カスタムレトリーバーの定義

# import QueryBundle
from llama_index.core import QueryBundle

# import NodeWithScore
from llama_index.core.schema import NodeWithScore

# Retrievers
from llama_index.core.retrievers import (
    BaseRetriever,
    VectorIndexRetriever,
    KeywordTableSimpleRetriever,
)

from typing import List

カスタム検索器の定義

class CustomRetriever(BaseRetriever):
    """Custom retriever that performs both semantic search and hybrid search."""

    def __init__(
        self,
        vector_retriever: VectorIndexRetriever,
        keyword_retriever: KeywordTableSimpleRetriever,
        mode: str = "AND",
    ) -> None:
        """Init params."""

        self._vector_retriever = vector_retriever
        self._keyword_retriever = keyword_retriever
        if mode not in ("AND", "OR"):
            raise ValueError("Invalid mode.")
        self._mode = mode
        super().__init__()

    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
        """Retrieve nodes given query."""

        vector_nodes = self._vector_retriever.retrieve(query_bundle)
        keyword_nodes = self._keyword_retriever.retrieve(query_bundle)

        vector_ids = {n.node.node_id for n in vector_nodes}
        keyword_ids = {n.node.node_id for n in keyword_nodes}

        combined_dict = {n.node.node_id: n for n in vector_nodes}
        combined_dict.update({n.node.node_id: n for n in keyword_nodes})

        if self._mode == "AND":
            retrieve_ids = vector_ids.intersection(keyword_ids)
        else:
            retrieve_ids = vector_ids.union(keyword_ids)

        retrieve_nodes = [combined_dict[rid] for rid in retrieve_ids]
        return retrieve_nodes

検索器をプラグインに入れるクエリエンジン

from llama_index.core import get_response_synthesizer
from llama_index.core.query_engine import RetrieverQueryEngine

# define custom retriever
vector_retriever = VectorIndexRetriever(index=vector_index, similarity_top_k=2)
keyword_retriever = KeywordTableSimpleRetriever(index=keyword_index)
custom_retriever = CustomRetriever(vector_retriever, keyword_retriever)

# define response synthesizer
response_synthesizer = get_response_synthesizer()

# assemble query engine
custom_query_engine = RetrieverQueryEngine(
    retriever=custom_retriever,
    response_synthesizer=response_synthesizer,
)

# vector query engine
vector_query_engine = RetrieverQueryEngine(
    retriever=vector_retriever,
    response_synthesizer=response_synthesizer,
)
# keyword query engine
keyword_query_engine = RetrieverQueryEngine(
    retriever=keyword_retriever,
    response_synthesizer=response_synthesizer,
)

実行

query_text = "カローラスポーツはどんな車ですか？"
query = tokenize_text(query_text)
response = custom_query_engine.query(query)
print(response)

Empty Response

今回はクエリを入力しても空の回答が返ってきました．

これらのコードはこちらで実行できます．

さらに，キーワード検索をBM25検索器を用いた実装も行いました．


まとめ

  • CustomRetrieverクラスを作成することで任意のハイブリッド検索器を定義できる

それでは、ご覧いただきありがとうございました。

この記事が気に入ったらサポートをしてみませんか？