Merge pull request #1 from Shakshi3104/feature

.gitignore

@@ -5,6 +5,8 @@
 # Develop
 .venv/
 logs/
+data/
+models/
 
 # Default
 # Byte-compiled / optimized / DLL files

cli_example.py

@@ -0,0 +1,14 @@
+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/data/notion_db.py

@@ -0,0 +1,216 @@
+import os
+from pathlib import Path
+import abc
+
+import pandas as pd
+from dotenv import load_dotenv
+
+import notion_client as nt
+from notion2md.exporter.block import StringExporter
+
+from loguru import logger
+
+
+class BaseNotionDatabase:
+    """
+    Notion DBからページのコンテンツを取り出すベースのクラス
+    """
+    def __init__(self):
+        load_dotenv()
+        self.notion_database_id = os.getenv("NOTION_DATABASE_ID")
+        self.integration_token = os.getenv("INTEGRATION_TOKEN")
+
+        # notion2mdの環境変数
+        os.environ["NOTION_TOKEN"] = os.getenv("INTEGRATION_TOKEN")
+
+        self.notion_client = nt.Client(auth=self.integration_token)
+
+    def load_database(self) -> list[dict]:
+        """
+        Notion DBのページ一覧を取得
+
+        Returns:
+
+        """
+        results = []
+        has_more = True
+        start_cursor = None
+
+        while has_more:
+            db = self.notion_client.databases.query(
+                **{
+                    "database_id": self.notion_database_id,
+                    "start_cursor": start_cursor
+                }
+            )
+            # 100件までしか1回に取得できない
+            # 100件以上ある場合 has_more = True
+            has_more = db["has_more"]
+            # 次のカーソル
+            start_cursor = db["next_cursor"]
+
+            # 取得結果
+            results += db["results"]
+
+        return results
+
+    @abc.abstractmethod
+    def load_content(self) -> list[dict]:
+        """
+        Notion DBのページの中身をdictで返す
+        Returns:
+
+        """
+        raise NotImplementedError
+
+
+class SakurapDB(BaseNotionDatabase):
+    def load_database(self) -> list[dict]:
+        """
+        Notion DBのページ一覧を取得
+
+        Returns:
+            results:
+                list[dict]
+
+        """
+        results = []
+        has_more = True
+        start_cursor = None
+
+        while has_more:
+            # "Rap詞 : 櫻井翔"がTrueのもののみ取得
+            db = self.notion_client.databases.query(
+                **{
+                    "database_id": self.notion_database_id,
+                    "filter": {
+                        "property": "Rap詞 : 櫻井翔",
+                        "checkbox": {
+                            "equals": True
+                        }
+                    },
+                    "start_cursor": start_cursor
+                }
+            )
+            # 100件までしか1回に取得できない
+            # 100件以上ある場合 has_more = True
+            has_more = db["has_more"]
+            # 次のカーソル
+            start_cursor = db["next_cursor"]
+
+            # 取得結果
+            results += db["results"]
+
+        return results
+
+    def __load_blocks(self, block_id: str) -> str:
+        """
+        Notionのページをプレーンテキストで取得する (Notion Official API)
+
+        Parameters
+        ----------
+        block_id:
+            str, Block ID
+
+        Returns
+        -------
+        texts:
+            str
+        """
+        block = self.notion_client.blocks.children.list(
+            **{
+                "block_id": block_id
+            }
+        )
+
+        # プレーンテキストを繋げる
+        def join_plain_texts():
+            text = [blck["paragraph"]["rich_text"][0]["plain_text"] if len(blck["paragraph"]["rich_text"])
+                    else "\n" for blck in block["results"]]
+
+            texts = "\n".join(text)
+            return texts
+
+        return join_plain_texts()
+
+    def load_content(self) -> list[dict]:
+        """
+        Notion DBのページの中身をdictで返す
+
+        Returns:
+            lyrics:
+                list[dict]
+        """
+
+        # DBのページ一覧を取得
+        db_results = self.load_database()
+        logger.info("🚦 [Notion] load database...")
+
+        # コンテンツ一覧
+        lyrics = []
+
+        logger.info("🚦 [Notion] start to load each page content ...")
+        # 各ページの処理
+        for result in db_results:
+            block_id = result["id"]
+            # rap_lyric = self.__load_blocks(block_id)
+
+            # Markdown形式でページを取得
+            rap_lyric = StringExporter(block_id=block_id).export()
+            # Markdownの修飾子を削除
+            rap_lyric = rap_lyric.replace("\n\n", "\n").replace("<br/>", "\n").replace("*", "")
+
+            lyrics.append(
+                {
+                    "title": result["properties"]["名前"]["title"][0]["plain_text"],
+                    "content": rap_lyric
+                }
+            )
+
+        logger.info("🚦 [Notion] Finish to load.")
+
+        return lyrics
+
+
+def fetch_sakurap_corpus(filepath: str, refetch=False) -> pd.DataFrame:
+    """
+    サクラップのコーパスを取得する
+    CSVファイルが存在しないときにNotionから取得する
+
+    Parameters
+    ----------
+    filepath:
+        str
+    refetch:
+        bool
+
+    Returns
+    -------
+
+    """
+    filepath = Path(filepath)
+
+    if not filepath.exists() or refetch:
+        # CSVファイルを保存するディレクトリが存在しなかったら作成する
+        if not filepath.parent.exists():
+            logger.info(f"🚦 [Notion] mkdir {str(filepath.parent)} ...")
+            filepath.parent.mkdir(parents=True, exist_ok=True)
+
+        logger.info("🚦 [Notion] fetch from Notion DB ...")
+        # dictを取得
+        rap_db = SakurapDB()
+        lyrics = rap_db.load_content()
+
+        lyrics_df = pd.DataFrame(lyrics)
+        lyrics_df.to_csv(filepath, index=False)
+    else:
+        logger.info("🚦 [Notion] load CSV file.")
+
+        lyrics_df = pd.read_csv(filepath)
+
+    return lyrics_df
+
+
+if __name__ == "__main__":
+    sakurap_db = SakurapDB()
+    lyrics = sakurap_db.load_content()

model/search/base.py

@@ -0,0 +1,20 @@
+import abc
+from typing import List, Union
+
+import pandas as pd
+
+
+class BaseSearchClient:
+    """
+    検査インタフェースクラス
+    """
+    corpus: pd.DataFrame | list | None = None
+
+    @classmethod
+    @abc.abstractmethod
+    def from_dataframe(cls, _data: pd.DataFrame, _target: str):
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def search_top_n(self, _query: Union[List[str], str], n: int=10) -> List[pd.DataFrame]:
+        raise NotImplementedError()

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

model/search/surface.py

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

model/search/vector.py

@@ -0,0 +1,182 @@
+import os
+
+from typing import List, Union, Optional
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+
+from copy import deepcopy
+from dotenv import load_dotenv
+from loguru import logger
+from tqdm import tqdm
+
+import sentence_transformers as st
+
+import voyager
+
+from model.search.base import BaseSearchClient
+from model.utils.timer import stop_watch
+
+
+def array_to_string(array: np.ndarray) -> str:
+    """
+    np.ndarrayを文字列に変換する
+
+    Parameters
+    ----------
+    array:
+        np.ndarray
+
+    Returns
+    -------
+    array_string:
+        str
+    """
+    array_string = f"{array.tolist()}"
+    return array_string
+
+
+class RuriEmbedder:
+    def __init__(self, model: Optional[st.SentenceTransformer] = None):
+
+        load_dotenv()
+
+        # モデルの保存先
+        self.model_dir = Path("models/ruri")
+        model_filepath = self.model_dir / "ruri-large"
+
+        # モデル
+        if model is None:
+            if model_filepath.exists():
+                logger.info(f"🚦 [RuriEmbedder] load ruri-large from local path: {model_filepath}")
+                self.model = st.SentenceTransformer(str(model_filepath))
+            else:
+                logger.info(f"🚦 [RuriEmbedder] load ruri-large from HuggingFace🤗")
+                token = os.getenv("HF_TOKEN")
+                self.model = st.SentenceTransformer("cl-nagoya/ruri-large", token=token)
+                # モデルを保存する
+                logger.info(f"🚦 [RuriEmbedder] save model ...")
+                self.model.save(str(model_filepath))
+        else:
+            self.model = model
+
+    def embed(self, text: Union[str, list[str]]) -> np.ndarray:
+        """
+
+        Parameters
+        ----------
+        text:
+            Union[str, list[str]], ベクトル化する文字列
+
+        Returns
+        -------
+        embedding:
+             np.ndarray, 埋め込み表現. トークンサイズ 1024
+        """
+        embedding = self.model.encode(text, convert_to_numpy=True)
+        return embedding
+
+
+class RuriVoyagerSearchClient(BaseSearchClient):
+    def __init__(self, dataset: pd.DataFrame, target: str,
+                 index: voyager.Index,
+                 model: RuriEmbedder):
+        load_dotenv()
+        # オリジナルのコーパス
+        self.dataset = dataset
+        self.corpus = dataset[target].values.tolist()
+
+        # 埋め込みモデル
+        self.embedder = model
+
+        # Voyagerインデックス
+        self.index = index
+
+    @classmethod
+    @stop_watch
+    def from_dataframe(cls, _data: pd.DataFrame, _target: str):
+        """
+        検索ドキュメントのpd.DataFrameから初期化する
+
+        Parameters
+        ----------
+        _data:
+            pd.DataFrame, 検索対象のDataFrame
+
+        _target:
+            str, 検索対象のカラム名
+
+        Returns
+        -------
+
+        """
+        logger.info("🚦 [RuriVoyagerSearchClient] Initialize from DataFrame")
+
+        search_field = _data[_target]
+        corpus = search_field.values.tolist()
+
+        # 埋め込みモデルの初期化
+        embedder = RuriEmbedder()
+
+        # Ruriの前処理
+        corpus = [f"文章: {c}" for c in corpus]
+
+        # ベクトル化する
+        embeddings = embedder.embed(corpus)
+
+        # 埋め込みベクトルの次元
+        num_dim = embeddings.shape[1]
+        logger.debug(f"🚦⚓️ [RuriVoyagerSearchClient] Number of dimensions of Embedding vector is {num_dim}")
+
+        # Voyagerのインデックスを初期化
+        index = voyager.Index(voyager.Space.Cosine, num_dimensions=num_dim)
+        # indexにベクトルを追加
+        _ = index.add_items(embeddings)
+
+        return cls(_data, _target, index, embedder)
+
+    @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"🚦 [RuriVoyagerSearchClient] Search top {n} | {_query}")
+
+        # 型チェック
+        if isinstance(_query, str):
+            _query = [_query]
+
+        # Ruriの前処理
+        _query = [f"クエリ: {q}" for q in _query]
+
+        # ベクトル化
+        embeddings_queries = self.embedder.embed(_query)
+
+        # ランキングtop-nをクエリ毎に取得
+        result = []
+        for embeddings_query in tqdm(embeddings_queries):
+            # Voyagerのインデックスを探索
+            neighbors_indices, distances = self.index.query(embeddings_query, k=n)
+            # 類似度スコア
+            df_res = deepcopy(self.dataset.iloc[neighbors_indices])
+            df_res["score"] = distances
+            # ランク
+            df_res["rank"] = deepcopy(df_res.reset_index()).index
+
+            result.append(df_res)
+
+        return result

model/utils/timer.py

@@ -0,0 +1,20 @@
+from functools import wraps
+import time
+
+from loguru import logger
+
+
+# https://qiita.com/hisatoshi/items/7354c76a4412dffc4fd7
+def stop_watch(func):
+    """
+    処理にかかる時間計測をするデコレータ
+    """
+    @wraps(func)
+    def wrapper(*args, **kargs):
+        logger.debug(f"🚦 [@stop_watch] measure time to run `{func.__name__}`.")
+        start = time.time()
+        result = func(*args, **kargs)
+        elapsed_time = time.time() - start
+        logger.debug(f"🚦 [@stop_watch] take {elapsed_time:.3f} sec to run `{func.__name__}`.")
+        return result
+    return wrapper

model/utils/tokenizer.py

@@ -0,0 +1,63 @@
+import abc
+from typing import Optional
+
+import MeCab
+# from janome.tokenizer import Tokenizer
+
+
+class BaseTokenizer:
+    @abc.abstractmethod
+    def parse(self, _text: str) -> str:
+        """
+        分かち書きした結果を返す
+
+        Parameters
+        ----------
+        _text:
+            str, 入力文章
+
+        Returns
+        -------
+        parsed:
+            str, 分かち書き後の文章, スペース区切り
+        """
+        raise NotImplementedError
+
+
+class MeCabTokenizer(BaseTokenizer):
+    def __init__(self, _parser: MeCab.Tagger) -> None:
+        self.parser = _parser
+
+    @classmethod
+    def from_tagger(cls, _tagger: Optional[str]):
+        parser = MeCab.Tagger(_tagger)
+        return cls(parser)
+
+    def parse(self, _text: str):
+        return self.parser.parse(_text)
+
+
+# class JanomeTokenizer(BaseTokenizer):
+#     def __init__(self, _tokenizer: Tokenizer):
+#         self.tokenizer = _tokenizer
+#
+#     @classmethod
+#     def from_user_simple_dictionary(cls, _dict_filepath: Optional[str] = None):
+#         """
+#         簡易辞書フォーマットによるユーザー辞書によるイニシャライザー
+#
+#         https://mocobeta.github.io/janome/#v0-2-7
+#
+#         Parameters
+#         ----------
+#         _dict_filepath:
+#             str, 簡易辞書フォーマットで書かれたユーザー辞書 (CSVファイル)のファイルパス
+#         """
+#
+#         if _dict_filepath is None:
+#             return cls(Tokenizer())
+#         else:
+#             return cls(Tokenizer(udic=_dict_filepath, udic_type='simpledic'))
+#
+#     def parse(self, _text: str) -> str:
+#         return " ".join(list(self.tokenizer.tokenize(_text, wakati=True)))

requirements.txt

@@ -9,9 +9,7 @@ tqdm
 python-dotenv
 
 # Visualization
-streamlit>=1.24
-st-pages
-streamlit-webrtc
+gradio
 
 tqdm>=4.65
 matplotlib>=3.7
@@ -25,8 +23,6 @@ pandas>=2.0
 opencv-python>=4.8
 pillow>=9.5
 
-# LLM
-
 # Others
 python-magic==0.4.27
 emoji>=2.6.0
@@ -39,6 +35,7 @@ voyager
 mecab-python3
 unidic-lite
 fugashi
+sentencepiece
 sentence-transformers>=3.0
 
 # Notion