File size: 5,123 Bytes
3e4b2ef
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b58c0e3
3e4b2ef
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b58c0e3
3e4b2ef
02c2acd
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3e4b2ef
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b58c0e3
3e4b2ef
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
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