flashrag-dev 0.1.1__py3-none-any.whl

flashrag/refiner/refiner.py

@@ -0,0 +1,252 @@
+from typing import List
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from flashrag.retriever.encoder import Encoder
+from tqdm import tqdm
+import re
+import torch
+import numpy as np
+
+class BaseRefiner:
+    r"""Base object of Refiner method"""
+
+    def __init__(self, config):
+        self.config = config
+        self.name = config["refiner_name"]
+        self.model_path = config["refiner_model_path"]
+        self.device = config["device"]
+        self.input_prompt_flag = config["refiner_input_prompt_flag"] if "refiner_input_prompt_flag" in config else False
+
+    def run(self, item) -> str:
+        r"""Get refining result.
+
+        Args:
+            item: dataset item, contains question, retrieval result...
+
+        Returns:
+            str: refining result of this item
+        """
+        pass
+
+    def batch_run(self, dataset, batch_size=None) -> List[str]:
+        return [self.run(item) for item in dataset]
+
+
+class LLMLinguaRefiner(BaseRefiner):
+    """Implementation for (Long)LLMLingua."""
+
+    def __init__(self, config):
+        super().__init__(config)
+        default_config = {
+            "rate": 0.55,
+            "condition_in_question": "after_condition",
+            "reorder_context": "sort",
+            "dynamic_context_compression_ratio": 0.3,
+            "condition_compare": True,
+            "context_budget": "+100",
+            "rank_method": "longllmlingua",
+        }
+        if "llmlingua_config" in config and config["llmlingua_config"] is not None:
+            self.compress_config = config["llmlingua_config"]
+        else:
+            self.compress_config = default_config
+
+        from flashrag.refiner.llmlingua_compressor import PromptCompressor
+
+        self.refiner = PromptCompressor(model_name=self.model_path)
+
+    def format_reference(self, retrieval_result):
+        format_reference = ""
+        for idx, doc_item in enumerate(retrieval_result):
+            content = doc_item["contents"]
+            title = content.split("\n")[0]
+            text = "\n".join(content.split("\n")[1:])
+            format_reference += f"Doc {idx+1}(Title: {title}) {text}\n"
+
+        return format_reference
+
+    def batch_run(self, dataset):
+        output = []
+        for item in tqdm(dataset, desc="Refining process: "):
+            question = item.question
+            retrieval_result = item.retrieval_result
+            # TODO: suit more cases
+            if self.input_prompt_flag:
+                input_prompt = item.prompt
+                prompt_split = input_prompt.split("\n\n")
+                # need fixed format prompt: instr + demon(retrieval results) + question
+                instruction, question = prompt_split[0], prompt_split[-1]
+                demonstration = "\n".join(prompt_split[1:-1])
+                item_output = self.refiner.compress_prompt(
+                    [i for i in demonstration.split("\n") if i != ""],
+                    instruction=instruction,
+                    question=question,
+                    **self.compress_config,
+                )
+            else:
+                docs = self.format_reference(retrieval_result).split("\n")
+                docs = [i for i in docs if i != ""]
+                item_output = self.refiner.compress_prompt(
+                    docs, instruction="", question=question, **self.compress_config
+                )
+            output.append(item_output["compressed_prompt"])
+        return output
+
+
+class SelectiveContextRefiner(BaseRefiner):
+    """Implementation for Selective Context"""
+
+    def __init__(self, config):
+        super().__init__(config)
+        from flashrag.refiner.selective_context_compressor import SelectiveContext
+
+        default_config = {"reduce_ratio": 0.5}
+
+        self.refiner = SelectiveContext(model_type="gpt2", model_path=self.model_path, lang="en")
+        if "sc_config" in config and config["sc_config"] is not None:
+            self.compress_config = config["sc_config"]
+        else:
+            self.compress_config = default_config
+
+    def format_reference(self, retrieval_result):
+        format_reference = ""
+        for idx, doc_item in enumerate(retrieval_result):
+            content = doc_item["contents"]
+            title = content.split("\n")[0]
+            text = "\n".join(content.split("\n")[1:])
+            format_reference += f"Doc {idx+1}(Title: {title}) {text}\n"
+
+        return format_reference
+
+    def batch_run(self, dataset):
+        # only use text
+        all_inputs = []
+        for item in dataset:
+            retrieval_result = item.retrieval_result
+            all_inputs.append(self.format_reference(retrieval_result))
+
+        output = []
+        for text in tqdm(all_inputs, desc="Refining process: "):
+            compress_text, _ = self.refiner(text, **self.compress_config)
+            output.append(compress_text)
+        return output
+
+
+class ExtractiveRefiner(BaseRefiner):
+    """Implementation for Extractive compressor.
+    Using retrieval method to select sentences or other granularity data.
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+        # number of keeping sentences
+        self.topk = config["refiner_topk"]
+        self.pooling_method = config["refiner_pooling_method"]
+        self.encode_max_length = config["refiner_encode_max_length"]
+        self.mini_batch_size = config['refiner_mini_batch_size'] if 'refiner_mini_batch_size' in config else 256
+        # load model
+        self.encoder = Encoder(
+            model_name=self.name, 
+            model_path=self.model_path, 
+            pooling_method=self.pooling_method, 
+            max_length=self.encode_max_length, 
+            use_fp16=True
+        )
+
+    def batch_run(self, dataset, batch_size=16):
+        questions = dataset.question
+        # only use text
+        retrieval_results = dataset.retrieval_result
+        retrieval_results = [
+            ["\n".join(doc_item["contents"].split("\n")[1:]) for doc_item in item_result]
+            for item_result in retrieval_results
+        ]
+
+        # split into sentences: [[sent1, sent2,...], [...]]
+        sent_lists = [
+            [i.strip() for i in re.split(r"(?<=[.!?])\s+", " ".join(res)) if len(i.strip()) > 5]
+            for res in retrieval_results
+        ]
+        score_lists = []  # matching scores, size == sent_lists
+        for idx in tqdm(range(0, len(questions), batch_size), desc="Refining process: "):
+            batch_questions = questions[idx : idx + batch_size]
+            batch_sents = sent_lists[idx : idx + batch_size]
+            question_embs = self.encoder.encode(batch_questions, is_query=True)
+            
+            flatten_batch_sents = sum(batch_sents, [])
+            sent_embs = []
+            for s_index in tqdm(range(0, len(flatten_batch_sents), self.mini_batchsize), desc='Sentence encoding..,'):
+                mini_batch_sents = flatten_batch_sents[s_index:s_index+self.mini_batchsize]
+                mini_sent_embs = self.encoder.encode(mini_batch_sents, is_query=False)
+                sent_embs.append(mini_sent_embs)
+            sent_embs = np.concatenate(sent_embs, axis=0)
+
+            scores = question_embs @ sent_embs.T
+            start_idx = 0
+            for row_score, single_list in zip(scores, batch_sents):
+                row_score = row_score.tolist()
+                score_lists.append(row_score[start_idx : start_idx + len(single_list)])
+                start_idx += len(single_list)
+
+        # select topk sents
+        retain_lists = []
+        for sent_scores, sent_list in zip(score_lists, sent_lists):
+            assert len(sent_scores) == len(sent_list)
+            if len(sent_scores) < self.topk:
+                retain_lists.append(sent_list)
+                continue
+
+            topk_idxs = torch.topk(torch.Tensor(sent_scores), min(self.topk, len(sent_scores))).indices.tolist()
+            retain_lists.append([sent_list[idx] for idx in sorted(topk_idxs) if idx < len(sent_list)])
+
+        return [" ".join(sents) for sents in retain_lists]
+
+
+class AbstractiveRecompRefiner(BaseRefiner):
+    """Implementation for Abstractive RECOMP compressor:
+    RECOMP: Improving Retrieval-Augmented LMs with Compression and Selective Augmentation.
+    """
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.max_input_length = config["refiner_max_input_length"]
+        self.max_output_length = config["refiner_max_output_length"]
+
+        # load model
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model_path)
+        self.model.cuda()
+        self.model.eval()
+
+    def batch_run(self, dataset, batch_size=2):
+        # only use text
+        retrieval_results = dataset.retrieval_result
+        retrieval_results = [
+            ["\n".join(doc_item["contents"].split("\n")[1:]) for doc_item in item_result]
+            for item_result in retrieval_results
+        ]
+
+        # input processing in recomp training format
+        format_inputs = [
+            "Question: {question}\n Document: {document}\n Summary: ".format(
+                question=item.question, document="\n".join(docs)
+            )
+            for item, docs in zip(dataset, retrieval_results)
+        ]
+
+        results = []
+        for idx in tqdm(range(0, len(format_inputs), batch_size), desc="Refining process: "):
+            batch_inputs = format_inputs[idx : idx + batch_size]
+            batch_inputs = self.tokenizer(
+                batch_inputs, return_tensors="pt", padding=True, truncation=True, max_length=self.max_input_length
+            ).to(self.device)
+
+            batch_outputs = self.model.generate(**batch_inputs, max_length=self.max_output_length)
+
+            batch_outputs = self.tokenizer.batch_decode(
+                batch_outputs, skip_special_tokens=True, clean_up_tokenization_spaces=False
+            )
+
+            results.extend(batch_outputs)
+
+        return results

flashrag/refiner/selective_context_compressor.py

@@ -0,0 +1,294 @@
+# Implementation of Selective-Context, modified from official repo: https://github.com/liyucheng09/Selective_Context
+# Licensed under The MIT License
+
+
+import re
+from typing import List, Tuple
+import spacy
+import numpy as np
+import os
+import sys
+
+sys.path.append("..")
+from dataclasses import dataclass
+from nltk.tokenize import word_tokenize
+import time
+import torch
+from transformers import GPT2Tokenizer, GPT2LMHeadModel, BertTokenizer
+
+
+@dataclass
+class LexicalUnits:
+    unit_type: str
+    text: List[str]
+    self_info: List[float] = None
+
+    def __add__(self, other):
+        assert self.unit_type == other.unit_type, "Cannot add two different unit types"
+        return LexicalUnits(self.unit_type, self.text + other.text, self.self_info + other.self_info)
+
+    def __radd__(self, other):
+        if other == 0:
+            return self
+        return NotImplementedError()
+
+    def add_to_head(self, token, self_info):
+        return LexicalUnits(self.unit_type, [token] + self.text, [self_info] + self.self_info)
+
+    def add_to_tail(self, token, self_info):
+        return LexicalUnits(self.unit_type, self.text + [token], self.self_info + [self_info])
+
+
+class SelectiveContext:
+
+    def __init__(self, model_type="gpt2", model_path="openai-community/gpt2", lang="en"):
+
+        self.model_type = model_type
+        self.model_path = model_path
+        self.lang = lang
+
+        # this means we calculate self-information sentence by sentence
+        self.sent_level_self_info = True
+
+        self._prepare_phrase_tokenizer()
+        self.sent_tokenize_pattern = r"(?<!\w\.\w.)(?<![A-Z][a-z]\.)(?<=\.|\?)\s"
+        self.phrase_mask_token = ""
+        self.sent_mask_token = "<...some content omitted.>"
+
+        self._prepare_model()
+
+    def _prepare_phrase_tokenizer(self):
+        # we use space to tokenize sentence into phrases
+        # for English, we should use `spacy.load("en_core_web_sm").add_pipe('merge_noun_chunks')`
+        # for Chinese, use `nlp = spacy.load('zh_core_web_sm')`` directly
+        lang = self.lang
+        if lang == "en":
+            self.nlp = spacy.load("en_core_web_lg", disable=["ner"])
+            self.nlp.add_pipe("merge_noun_chunks")
+        elif lang == "zh":
+            self.nlp = spacy.load("zh_core_web_sm", disable=["ner"])
+
+    def _prepare_model(self):
+        # Load tokenizer
+        if self.lang == "zh":
+            self.tokenizer = BertTokenizer.from_pretrained("uer/gpt2-chinese-cluecorpussmall")
+        elif self.lang == "en":
+            self.tokenizer = GPT2Tokenizer.from_pretrained(self.model_path)
+        else:
+            raise NotImplementedError()
+
+        if self.model_type == "gpt2":
+            if self.lang == "zh":
+                self.model = GPT2LMHeadModel.from_pretrained("uer/gpt2-chinese-cluecorpussmall")
+            else:
+                self.model = GPT2LMHeadModel.from_pretrained(self.model_path)
+            self.model.to("cuda")
+            self.model.eval()
+
+            print("model loaded")
+
+            self.max_token_length = self.model.config.n_positions
+            self.get_self_information = self._get_self_info_via_gpt2
+
+        elif self.model_type == "curie":
+            global openai
+            import openai
+
+            self.max_token_length = 2048
+
+            self.get_self_information = self._get_self_info_via_curie
+
+    def get_self_information(self, text: str) -> Tuple[List[str], List[float]]:
+        # it takes text as input, and return a list of words and a list of self-information scores
+        raise NotImplementedError
+
+    def _get_self_info_via_gpt2(self, text: str) -> Tuple[List[str], List[float]]:
+        if self.lang == "en":
+            text = f"<|endoftext|>{text}"
+        elif self.lang == "zh":
+            text = f"[CLS]{text}"
+        with torch.inference_mode(mode=True):
+            encoding = self.tokenizer(
+                text, max_length=1024, truncation=True, add_special_tokens=False, return_tensors="pt"
+            )
+            encoding = encoding.to(self.model.device)
+            outputs = self.model(**encoding)
+            logits = outputs.logits
+            probs = torch.softmax(logits, dim=-1)
+            self_info = -torch.log(probs)
+
+        input_ids = encoding["input_ids"]
+        input_ids_expaned = input_ids[:, 1:].unsqueeze(-1)
+
+        tokens = [self.tokenizer.decode(token_) for token_ in input_ids.squeeze().tolist()[1:]]
+        return tokens, self_info[:, :-1].gather(-1, input_ids_expaned).squeeze(-1).squeeze(0).tolist()
+
+    def _get_self_info_via_curie(self, text: str) -> Tuple[List[str], List[float]]:
+        num_retry = 3
+        openai.api_key = os.environ["OPENAI_API_KEY"]
+
+        for _ in range(num_retry):
+            try:
+                r = openai.Completion.create(
+                    model="curie",
+                    prompt=f"<|endoftext|>{text}",
+                    max_tokens=0,
+                    temperature=0,
+                    echo=True,
+                    logprobs=0,
+                )
+                break
+            except Exception as e:
+                print(e)
+                time.sleep(1)
+
+        result = r["choices"][0]
+        tokens, logprobs = result["logprobs"]["tokens"][1:], result["logprobs"]["token_logprobs"][1:]
+
+        assert len(tokens) == len(logprobs), f"Expected {len(tokens)} logprobs, got {len(logprobs)}"
+
+        self_info = [-logprob for logprob in logprobs]
+        return tokens, self_info
+
+    def _lexical_unit(self, sents):
+
+        if self.sent_level_self_info:
+            sent_self_info = []
+            all_noun_phrases = []
+            all_noun_phrases_info = []
+            all_tokens = []
+            all_token_self_info = []
+
+            for sent in sents:
+                # print(sent)
+                tokens, self_info = self.get_self_information(sent)
+                sent_self_info.append(np.mean(self_info))
+
+                all_tokens.extend(tokens)
+                all_token_self_info.extend(self_info)
+
+                noun_phrases, noun_phrases_info = self._calculate_lexical_unit(tokens, self_info)
+
+                # We need to add a space before the first noun phrase for every sentence except the first one
+                if len(all_noun_phrases) != 0:
+                    noun_phrases[0] = f" {noun_phrases[0]}"
+                all_noun_phrases.extend(noun_phrases)
+                all_noun_phrases_info.extend(noun_phrases_info)
+
+            return [
+                LexicalUnits("sent", text=sents, self_info=sent_self_info),
+                LexicalUnits("phrase", text=all_noun_phrases, self_info=all_noun_phrases_info),
+                LexicalUnits("token", text=all_tokens, self_info=all_token_self_info),
+            ]
+
+    def _calculate_lexical_unit(self, tokens, self_info):
+        def _unit_info(tokens, self_info, units):
+            current_unit_idx = 0
+            current_position = 0
+            unit_self_info = [[] for _ in range(len(units))]
+
+            for idx, (token, info) in enumerate(zip(tokens, self_info)):
+                current_position += len(token)
+                if current_position == len(units[current_unit_idx]):
+                    unit_self_info[current_unit_idx].append(info)
+                    current_position = current_position - len(units[current_unit_idx])
+                    current_unit_idx += 1
+                elif current_position > len(units[current_unit_idx]):
+                    counter_ = 1
+                    current_position = current_position - len(units[current_unit_idx])
+                    current_unit_idx += 1
+                    while current_position >= len(units[current_unit_idx]):
+                        counter_ += 1
+                        current_position = current_position - len(units[current_unit_idx])
+                        current_unit_idx += 1
+                        if current_unit_idx >= len(units):
+                            break
+                    partial_info = info / counter_
+                    for _ in range(counter_):
+                        unit_self_info[(current_unit_idx - 1) - _].append(partial_info)
+                else:
+                    if token == " ":
+                        continue
+                    unit_self_info[current_unit_idx].append(info)
+
+            unit_self_info_ = [np.mean(info) for info in unit_self_info]
+            return unit_self_info_
+
+        def _noun_phrases(sent):
+            noun_phrases = []
+            doc = self.nlp(sent)
+            for index, chunk in enumerate(doc):
+                if index == 0:
+                    noun_phrases.append(chunk.text)
+                else:
+                    noun_phrases.append(doc[index - 1].whitespace_ + chunk.text)
+            return noun_phrases
+
+        if self.sent_level_self_info:
+            # in this case, the self_info is for each sentence
+            # we only need to calculate the self_info for each phrase
+
+            sent = "".join(tokens)
+            # noun_phrases = [chunk.text for chunk in self.nlp(sent).noun_chunks]
+            noun_phrases = _noun_phrases(sent)
+            # noun_phrases[-1] = noun_phrases[-1] + ' '
+            noun_phrases_info = _unit_info(tokens, self_info, noun_phrases)
+
+            return noun_phrases, noun_phrases_info
+
+    def beautify_context(self, context: str) -> str:
+        context = re.sub(r"\s+", " ", context)
+        return context
+
+    def self_info_mask(self, sents: List[str], self_info: List[float], mask_level):
+        # mask_level: mask sentences, phrases, or tokens
+        sents_after_mask = []
+        masked_sents = []
+
+        self.ppl_threshold = np.nanpercentile(self_info, self.mask_ratio * 100)
+
+        for sent, info in zip(sents, self_info):
+            if info < self.ppl_threshold:
+                masked_sents.append(sent)
+                sents_after_mask.append(self.mask_a_sent(sent, mask_level))
+            else:
+                sents_after_mask.append(sent)
+        masked_context = " ".join(sents_after_mask) if mask_level == "sent" else "".join(sents_after_mask)
+
+        return masked_context, masked_sents
+
+    def mask_a_sent(self, sent, level):
+        if level == "phrase":
+            return self.phrase_mask_token
+        elif level == "sent":
+            if self.keep_leading_word:
+                leading_few_words = " ".join(word_tokenize(sent)[: self.num_lead_words]) + " "
+            else:
+                leading_few_words = ""
+            return leading_few_words + self.mask_token
+        elif level == "token":
+            return ""
+
+    def __call__(self, text: str, reduce_ratio: float = 0.5, reduce_level: str = "phrase") -> List[str]:
+        context = self.beautify_context(text)
+
+        self.mask_ratio = reduce_ratio
+
+        sents = re.split(self.sent_tokenize_pattern, context)
+        sents = [sent.strip() for sent in sents if sent.strip()]
+
+        # You want the reduce happen at sentence level, phrase level, or token level?
+        assert reduce_level in [
+            "sent",
+            "phrase",
+            "token",
+        ], f"reduce_level should be one of ['sent', 'phrase', 'token'], got {reduce_level}"
+        sent_lus, phrase_lus, token_lus = self._lexical_unit(sents)
+        # print(phrase_lus, '^^^^')
+        lexical_level = {"sent": sent_lus, "phrase": phrase_lus, "token": token_lus}
+
+        # context is the reduced context, masked_sents denotes what context has been filtered out
+        context, masked_sents = self.self_info_mask(
+            lexical_level[reduce_level].text, lexical_level[reduce_level].self_info, reduce_level
+        )
+        return context, masked_sents

flashrag/retriever/__init__.py

@@ -0,0 +1,3 @@
+from flashrag.retriever.retriever import *
+from flashrag.retriever.reranker import *
+from flashrag.retriever.utils import *

flashrag/retriever/__main__.py

@@ -0,0 +1,4 @@
+from . import index_builder
+
+if __name__ == "__main__":
+    index_builder.main()

flashrag/retriever/encoder.py

@@ -0,0 +1,120 @@
+from typing import List
+import torch
+import numpy as np
+from flashrag.retriever.utils import load_model, pooling
+
+
+def parse_query(model_name, query_list, is_query=True):
+    """
+    processing query for different encoders
+    """
+
+    def is_zh(str):
+        import unicodedata
+
+        zh_char = 0
+        for c in str:
+            try:
+                if "CJK" in unicodedata.name(c):
+                    zh_char += 1
+            except:
+                continue
+        if zh_char / len(str) > 0.2:
+            return True
+        else:
+            return False
+
+    if isinstance(query_list, str):
+        query_list = [query_list]
+
+    if "e5" in model_name.lower():
+        if is_query:
+            query_list = [f"query: {query}" for query in query_list]
+        else:
+            query_list = [f"passage: {query}" for query in query_list]
+
+    if "bge" in model_name.lower():
+        if is_query:
+            if is_zh(query_list[0]):
+                query_list = [f"为这个句子生成表示以用于检索相关文章：{query}" for query in query_list]
+            else:
+                query_list = [
+                    f"Represent this sentence for searching relevant passages: {query}" for query in query_list
+                ]
+
+    return query_list
+
+
+class Encoder:
+    def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16):
+        self.model_name = model_name
+        self.model_path = model_path
+        self.pooling_method = pooling_method
+        self.max_length = max_length
+        self.use_fp16 = use_fp16
+
+        self.model, self.tokenizer = load_model(model_path=model_path, use_fp16=use_fp16)
+
+    @torch.inference_mode(mode=True)
+    def encode(self, query_list: List[str], is_query=True) -> np.ndarray:
+        query_list = parse_query(self.model_name, query_list, is_query)
+
+        inputs = self.tokenizer(
+            query_list, max_length=self.max_length, padding=True, truncation=True, return_tensors="pt"
+        )
+        inputs = {k: v.cuda() for k, v in inputs.items()}
+
+        if "T5" in type(self.model).__name__:
+            # T5-based retrieval model
+            decoder_input_ids = torch.zeros((inputs["input_ids"].shape[0], 1), dtype=torch.long).to(
+                inputs["input_ids"].device
+            )
+            output = self.model(**inputs, decoder_input_ids=decoder_input_ids, return_dict=True)
+            query_emb = output.last_hidden_state[:, 0, :]
+
+        else:
+            output = self.model(**inputs, return_dict=True)
+            query_emb = pooling(
+                output.pooler_output, output.last_hidden_state, inputs["attention_mask"], self.pooling_method
+            )
+
+        query_emb = query_emb.detach().cpu().numpy()
+        query_emb = query_emb.astype(np.float32, order="C")
+        return query_emb
+
+
+class STEncoder:
+    def __init__(self, model_name, model_path, max_length, use_fp16):
+        import torch
+        from sentence_transformers import SentenceTransformer
+
+        self.model_name = model_name
+        self.model_path = model_path
+        self.max_length = max_length
+        self.use_fp16 = use_fp16
+
+        self.model = SentenceTransformer(
+            model_path, model_kwargs={"torch_dtype": torch.float16 if use_fp16 else torch.float}
+        )
+
+    @torch.inference_mode(mode=True)
+    def encode(self, query_list: List[str], batch_size=64, is_query=True) -> np.ndarray:
+        query_list = parse_query(self.model_name, query_list, is_query)
+        query_emb = self.model.encode(
+            query_list, batch_size=batch_size, convert_to_numpy=True, normalize_embeddings=True
+        )
+        query_emb = query_emb.astype(np.float32, order="C")
+
+        return query_emb
+
+    @torch.inference_mode(mode=True)
+    def multi_gpu_encode(self, query_list: List[str], is_query=True, batch_size=None) -> np.ndarray:
+        query_list = parse_query(self.model_name, query_list, is_query)
+        pool = self.model.start_multi_process_pool()
+        query_emb = self.model.encode_multi_process(
+            query_list, pool, convert_to_numpy=True, normalize_embeddings=True, batch_size=batch_size
+        )
+        self.model.stop_multi_process_pool(pool)
+        query_emb.astype(np.float32, order="C")
+
+        return query_emb