[delete] delete to make vector search only

app.py

@@ -1,17 +1,17 @@
 import gradio as gr
 import pandas as pd
 
-from model.search.hybrid import HybridSearchClient
+from model.search.vector import RuriVoyagerSearchClient
 from model.data.notion_db import fetch_sakurap_corpus
 
 
-def search(search_client: HybridSearchClient):
+def search(search_client: RuriVoyagerSearchClient):
     def _search(query: str) -> pd.DataFrame:
         results = search_client.search_top_n(query)
         result = results[0]
         result["rank"] = result["rank"] + 1
-        result = result[["rank", "title", "content", "rank_sparse", "rank_dense"]]
-        result.columns = ["rank", "title", "rap lyric", "rank: surface", "rank: vector"]
+        result = result[["rank", "title", "content", "score"]]
+        result.columns = ["rank", "title", "rap lyric", "distance"]
         return result
 
     return _search
@@ -21,13 +21,13 @@ if __name__ == "__main__":
     # Load dataset
     sakurap_df = fetch_sakurap_corpus("./data/sakurap_corpus.csv")
     # Initialize search client
-    search_client = HybridSearchClient.from_dataframe(sakurap_df, "content")
+    search_client = RuriVoyagerSearchClient.from_dataframe(sakurap_df, "content")
 
     with gr.Blocks() as search_interface:
         gr.Markdown("""
-        # 💎 Cobalt
-        Demo app for hybrid search with vector and surface search using [Ruri](https://huggingface.co/cl-nagoya/ruri-large), [BM25](https://github.com/dorianbrown/rank_bm25) and [Voyager](https://spotify.github.io/voyager/).
-        
+        # 💎 Cobalt DuckDB 🦆
+        Demo app for vector search using [Ruri](https://huggingface.co/cl-nagoya/ruri-large) and DuckDB.
+
         You can search ARASHI's songs with rap lyrics by Sho Sakurai.
         """)
         # Input query
@@ -49,4 +49,4 @@ if __name__ == "__main__":
 
     # App launch
     search_interface.queue()
-    search_interface.launch(server_name="0.0.0.0")
+    search_interface.launch(server_name="0.0.0.0")

example.py

@@ -1,14 +0,0 @@
-import pandas as pd
-
-from model.search.hybrid import HybridSearchClient
-from model.data.notion_db import fetch_sakurap_corpus
-
-
-if __name__ == "__main__":
-    # Load dataset
-    sakurap_df = fetch_sakurap_corpus("./data/sakurap_corpus.csv")
-    # sakurap_df = pd.read_csv("./data/sakurap_corpus.csv")
-
-    # hybrid search
-    search_client = HybridSearchClient.from_dataframe(sakurap_df, "content")
-    results = search_client.search_top_n("嵐 5人の歴史")

model/search/hybrid.py

@@ -1,146 +0,0 @@
-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

model/search/surface.py

@@ -1,147 +0,0 @@
-from copy import deepcopy
-from typing import List, Union
-
-import pandas as pd
-import numpy as np
-
-from loguru import logger
-from tqdm import tqdm
-
-from rank_bm25 import BM25Okapi
-
-from model.search.base import BaseSearchClient
-from model.utils.tokenizer import MeCabTokenizer
-from model.utils.timer import stop_watch
-
-
-class BM25Wrapper(BM25Okapi):
-    def __init__(self, dataset: pd.DataFrame, target, tokenizer=None, k1=1.5, b=0.75, epsilon=0.25):
-        self.k1 = k1
-        self.b = b
-        self.epsilon = epsilon
-        self.dataset = dataset
-        corpus = dataset[target].values.tolist()
-        super().__init__(corpus, tokenizer)
-
-    def get_top_n(self, query, documents, n=5):
-        assert self.corpus_size == len(documents), "The documents given don't match the index corpus!"
-
-        scores = self.get_scores(query)
-        top_n = np.argsort(scores)[::-1][:n]
-
-        result = deepcopy(self.dataset.iloc[top_n])
-        result["score"] = scores[top_n]
-        return result
-
-
-class BM25SearchClient(BaseSearchClient):
-    def __init__(self, _model: BM25Okapi, _corpus: List[List[str]]):
-        """
-
-        Parameters
-        ----------
-        _model:
-            BM25Okapi
-        _corpus:
-            List[List[str]], 検索対象の分かち書き後のフィールド
-        """
-        self.model = _model
-        self.corpus = _corpus
-
-    @staticmethod
-    def tokenize_ja(_text: List[str]):
-        """MeCab日本語分かち書きによるコーパス作成
-
-        Args:
-            _text (List[str]): コーパス文のリスト
-
-        Returns:
-            List[List[str]]: 分かち書きされたテキストのリスト
-        """
-
-        # MeCabで分かち書き
-        parser = MeCabTokenizer.from_tagger("-Owakati")
-
-        corpus = []
-        with tqdm(_text) as pbar:
-            for i, t in enumerate(pbar):
-                try:
-                    # 分かち書きをする
-                    corpus.append(parser.parse(t).split())
-                except TypeError as e:
-                    if not isinstance(t, str):
-                        logger.info(f"🚦 [BM25SearchClient] Corpus index of {i} is not instance of String.")
-                        corpus.append(["[UNKNOWN]"])
-                    else:
-                        raise e
-        return corpus
-
-    @classmethod
-    def from_dataframe(cls, _data: pd.DataFrame, _target: str):
-        """
-        検索ドキュメントのpd.DataFrameから初期化する
-
-        Parameters
-        ----------
-        _data:
-            pd.DataFrame, 検索対象のDataFrame
-
-        _target:
-            str, 検索対象のカラム名
-
-        Returns
-        -------
-
-        """
-
-        logger.info("🚦 [BM25SearchClient] Initialize from DataFrame")
-
-        search_field = _data[_target]
-        corpus = search_field.values.tolist()
-
-        # 分かち書きをする
-        corpus_tokenized = cls.tokenize_ja(corpus)
-        _data["tokenized"] = corpus_tokenized
-
-        bm25 = BM25Wrapper(_data, "tokenized")
-        return cls(bm25, corpus_tokenized)
-
-    @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"🚦 [BM25SearchClient] Search top {n} | {_query}")
-
-        # 型チェック
-        if isinstance(_query, str):
-            _query = [_query]
-
-        # クエリを分かち書き
-        query_tokens = self.tokenize_ja(_query)
-
-        # ランキングtop-nをクエリ毎に取得
-        result = []
-        for query in tqdm(query_tokens):
-            df_res = self.model.get_top_n(query, self.corpus, n)
-            # ランク
-            df_res["rank"] = deepcopy(df_res.reset_index()).index
-            df_res = df_res.drop(columns=["tokenized"])
-            result.append(df_res)
-
-        logger.success(f"🚦 [BM25SearchClient] Executed")
-
-        return result