[add] Implement hybrid search

model/search/hybrid.py

@@ -0,0 +1,146 @@
+from typing import Union, List
+
+import pandas as pd
+from copy import deepcopy
+
+from dotenv import load_dotenv
+from loguru import logger
+from tqdm import tqdm
+
+from model.search.base import BaseSearchClient
+from model.search.surface import BM25SearchClient
+from model.search.vector import RuriVoyagerSearchClient
+
+from model.utils.timer import stop_watch
+
+
+def reciprocal_rank_fusion(sparse: pd.DataFrame, dense: pd.DataFrame, k=60) -> pd.DataFrame:
+    """
+    Reciprocal Rank Fusionを計算する
+
+    Notes
+    ----------
+    RRFの計算は以下の式
+
+    .. math:: RRF = \sum_{i=1}^n \frac{1}{k+r_i}
+
+    Parameters
+    ----------
+    sparse:
+        pd.DataFrame, 表層検索の検索結果
+    dense:
+        pd.DataFrame, ベクトル検索の結果
+    k:
+        int,
+
+    Returns
+    -------
+    rank_results:
+        pd.DataFrame, RRFによるリランク結果
+
+    """
+    # カラム名を変更
+    sparse = sparse.rename(columns={"rank": "rank_sparse"})
+    dense = dense.rename(columns={"rank": "rank_dense"})
+    # denseはランク以外を落として結合する
+    dense_ = dense["rank_dense"]
+
+    # 順位を1からスタートするようにする
+    sparse["rank_sparse"] += 1
+    dense_ += 1
+
+    # 文書のインデックスをキーに結合する
+    rank_results = pd.merge(sparse, dense_, how="left", left_index=True, right_index=True)
+
+    # RRFスコアの計算
+    rank_results["rrf_score"] = 1 / (rank_results["rank_dense"] + k) + 1 / (rank_results["rank_sparse"] + k)
+
+    # RRFスコアのスコアが大きい順にソート
+    rank_results = rank_results.sort_values(["rrf_score"], ascending=False)
+    rank_results["rank"] = deepcopy(rank_results.reset_index()).index
+
+    return rank_results
+
+
+class HybridSearchClient(BaseSearchClient):
+    def __init__(self, dense_model: BaseSearchClient, sparse_model: BaseSearchClient):
+        self.dense_model = dense_model
+        self.sparse_model = sparse_model
+
+    @classmethod
+    @stop_watch
+    def from_dataframe(cls, _data: pd.DataFrame, _target: str):
+        """
+        検索ドキュメントのpd.DataFrameから初期化する
+
+        Parameters
+        ----------
+        _data:
+            pd.DataFrame, 検索対象のDataFrame
+
+        _target:
+            str, 検索対象のカラム名
+
+        Returns
+        -------
+
+        """
+        # 表層検索の初期化
+        dense_model = BM25SearchClient.from_dataframe(_data, _target)
+        # ベクトル検索の初期化
+        sparse_model = RuriVoyagerSearchClient.from_dataframe(_data, _target)
+
+        return cls(dense_model, sparse_model)
+
+    @stop_watch
+    def search_top_n(self, _query: Union[List[str], str], n: int = 10) -> List[pd.DataFrame]:
+        """
+        クエリに対する検索結果をtop-n個取得する
+
+        Parameters
+        ----------
+        _query:
+            Union[List[str], str], 検索クエリ
+        n:
+            int, top-nの個数. デフォルト 10.
+
+        Returns
+        -------
+        results:
+            List[pd.DataFrame], ランキング結果
+        """
+
+        logger.info(f"🚦 [HybridSearchClient] Search top {n} | {_query}")
+
+        # 型チェック
+        if isinstance(_query, str):
+            _query = [_query]
+
+        # ランキングtop-nをクエリ毎に取得
+        result = []
+        for query in tqdm(_query):
+            assert len(self.sparse_model.corpus) == len(
+                self.dense_model.corpus), "The document counts do not match between sparse and dense!"
+
+            # ドキュメント数
+            doc_num = len(self.sparse_model.corpus)
+
+            # 表層検索
+            logger.info(f"🚦 [HybridSearchClient] run surface search ...")
+            sparse_res = self.sparse_model.search_top_n(query, n=doc_num)
+            # ベクトル検索
+            logger.info(f"🚦 [HybridSearchClient] run vector search ...")
+            dense_res = self.dense_model.search_top_n(query, n=doc_num)
+
+            # RRFスコアの計算
+            logger.info(f"🚦 [HybridSearchClient] compute RRF scores ...")
+            rrf_res = reciprocal_rank_fusion(sparse_res[0], dense_res[0])
+
+            # 結果をtop Nに絞る
+            top_num = 10
+            rrf_res = rrf_res.head(top_num)
+            logger.info(f"🚦 [HybridSearchClient] return {top_num} results")
+
+            result.append(rrf_res)
+
+        return result