flashrag-dev 0.1.1__py3-none-any.whl

flashrag/retriever/retriever.py

@@ -0,0 +1,346 @@
+import json
+import os
+
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+import warnings
+from typing import List, Dict
+import functools
+from tqdm import tqdm
+import faiss
+
+from flashrag.utils import get_reranker
+from flashrag.retriever.utils import load_corpus, load_docs
+from flashrag.retriever.encoder import Encoder, STEncoder
+
+
+def cache_manager(func):
+    """
+    Decorator used for retrieving document cache.
+    With the decorator, The retriever can store each retrieved document as a file and reuse it.
+    """
+
+    @functools.wraps(func)
+    def wrapper(self, query_list, num=None, return_score=False):
+        if num is None:
+            num = self.topk
+        if self.use_cache:
+            if isinstance(query_list, str):
+                new_query_list = [query_list]
+            else:
+                new_query_list = query_list
+
+            no_cache_query = []
+            cache_results = []
+            for query in new_query_list:
+                if query in self.cache:
+                    cache_res = self.cache[query]
+                    if len(cache_res) < num:
+                        warnings.warn(f"The number of cached retrieval results is less than topk ({num})")
+                    cache_res = cache_res[:num]
+                    # separate the doc score
+                    doc_scores = [item.pop("score") for item in cache_res]
+                    cache_results.append((cache_res, doc_scores))
+                else:
+                    cache_results.append(None)
+                    no_cache_query.append(query)
+
+            if no_cache_query != []:
+                # use batch search without decorator
+                no_cache_results, no_cache_scores = self._batch_search_with_rerank(no_cache_query, num, True)
+                no_cache_idx = 0
+                for idx, res in enumerate(cache_results):
+                    if res is None:
+                        assert new_query_list[idx] == no_cache_query[no_cache_idx]
+                        cache_results = (
+                            no_cache_results[no_cache_idx],
+                            no_cache_scores[no_cache_scores],
+                        )
+                        no_cache_idx += 1
+
+            results, scores = (
+                [t[0] for t in cache_results],
+                [t[1] for t in cache_results],
+            )
+
+        else:
+            results, scores = func(self, query_list, num, True)
+
+        if self.save_cache:
+            # merge result and score
+            save_results = results.copy()
+            save_scores = scores.copy()
+            if isinstance(query_list, str):
+                query_list = [query_list]
+                if "batch" not in func.__name__:
+                    save_results = [save_results]
+                    save_scores = [save_scores]
+            for query, doc_items, doc_scores in zip(query_list, save_results, save_scores):
+                for item, score in zip(doc_items, doc_scores):
+                    item["score"] = score
+                self.cache[query] = doc_items
+
+        if return_score:
+            return results, scores
+        else:
+            return results
+
+    return wrapper
+
+
+def rerank_manager(func):
+    """
+    Decorator used for reranking retrieved documents.
+    """
+
+    @functools.wraps(func)
+    def wrapper(self, query_list, num=None, return_score=False):
+        results, scores = func(self, query_list, num, True)
+        if self.use_reranker:
+            results, scores = self.reranker.rerank(query_list, results)
+            if "batch" not in func.__name__:
+                results = results[0]
+                scores = scores[0]
+        if return_score:
+            return results, scores
+        else:
+            return results
+
+    return wrapper
+
+
+class BaseRetriever:
+    """Base object for all retrievers."""
+
+    def __init__(self, config):
+        self.config = config
+        self.retrieval_method = config["retrieval_method"]
+        self.topk = config["retrieval_topk"]
+
+        self.index_path = config["index_path"]
+        self.corpus_path = config["corpus_path"]
+
+        self.save_cache = config["save_retrieval_cache"]
+        self.use_cache = config["use_retrieval_cache"]
+        self.cache_path = config["retrieval_cache_path"]
+
+        self.use_reranker = config["use_reranker"]
+        if self.use_reranker:
+            self.reranker = get_reranker(config)
+
+        if self.save_cache:
+            self.cache_save_path = os.path.join(config["save_dir"], "retrieval_cache.json")
+            self.cache = {}
+        if self.use_cache:
+            assert self.cache_path is not None
+            with open(self.cache_path, "r") as f:
+                self.cache = json.load(f)
+
+    def _save_cache(self):
+        with open(self.cache_save_path, "w") as f:
+            json.dump(self.cache, f, indent=4)
+
+    def _search(self, query: str, num: int, return_score: bool) -> List[Dict[str, str]]:
+        r"""Retrieve topk relevant documents in corpus.
+
+        Return:
+            list: contains information related to the document, including:
+                contents: used for building index
+                title: (if provided)
+                text: (if provided)
+
+        """
+
+        pass
+
+    def _batch_search(self, query_list, num, return_score):
+        pass
+
+    @cache_manager
+    @rerank_manager
+    def search(self, *args, **kwargs):
+        return self._search(*args, **kwargs)
+
+    @cache_manager
+    @rerank_manager
+    def batch_search(self, *args, **kwargs):
+        return self._batch_search(*args, **kwargs)
+
+    @rerank_manager
+    def _batch_search_with_rerank(self, *args, **kwargs):
+        return self._batch_search(*args, **kwargs)
+
+    @rerank_manager
+    def _search_with_rerank(self, *args, **kwargs):
+        return self._search(*args, **kwargs)
+
+
+class BM25Retriever(BaseRetriever):
+    r"""BM25 retriever based on pre-built pyserini index."""
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.backend = config['bm25_backend']
+
+        if self.backend == 'pyserini':
+            # Warning: the method based on pyserini will be deprecated
+            from pyserini.search.lucene import LuceneSearcher
+
+            self.searcher = LuceneSearcher(self.index_path)
+            self.contain_doc = self._check_contain_doc()
+            if not self.contain_doc:
+                self.corpus = load_corpus(self.corpus_path)
+            self.max_process_num = 8
+        elif self.backend == 'bm25s':
+            import Stemmer
+            import bm25s
+
+            stemmer = Stemmer.Stemmer('english')
+            self.searcher = bm25s.BM25.load(self.index_path, mmap=True, load_corpus=True)
+            self.searcher.backend = 'numba'
+            self.tokenizer = bm25s.tokenization.Tokenizer(stemmer=stemmer)
+        else:
+            assert False, 'Invalid bm25 backend!'
+
+    def _check_contain_doc(self):
+        r"""Check if the index contains document content"""
+        return self.searcher.doc(0).raw() is not None
+
+    def _search(self, query: str, num: int = None, return_score=False) -> List[Dict[str, str]]:
+        if num is None:
+            num = self.topk
+        if self.backend == 'pyserini': 
+            hits = self.searcher.search(query, num)
+            if len(hits) < 1:
+                if return_score:
+                    return [], []
+                else:
+                    return []
+
+            scores = [hit.score for hit in hits]
+            if len(hits) < num:
+                warnings.warn("Not enough documents retrieved!")
+            else:
+                hits = hits[:num]
+
+            if self.contain_doc:
+                all_contents = [json.loads(self.searcher.doc(hit.docid).raw())["contents"] for hit in hits]
+                results = [
+                    {
+                        "title": content.split("\n")[0].strip('"'),
+                        "text": "\n".join(content.split("\n")[1:]),
+                        "contents": content,
+                    }
+                    for content in all_contents
+                ]
+            else:
+                results = load_docs(self.corpus, [hit.docid for hit in hits])
+        elif self.backend == 'bm25s':
+            query_tokens = self.tokenizer.tokenize([query])
+            results, scores = self.searcher.retrieve(query_tokens, k=num)
+            results = results[0]
+            scores = scores[0]
+        else:
+            assert False, 'Invalid bm25 backend!'
+
+        if return_score:
+            return results, scores
+        else:
+            return results
+
+    def _batch_search(self, query_list, num: int = None, return_score=False):
+        if self.backend == 'pyserini': 
+            # TODO: modify batch method
+            results = []
+            scores = []
+            for query in query_list:
+                item_result, item_score = self._search(query, num, True)
+                results.append(item_result)
+                scores.append(item_score)
+        elif self.backend == 'bm25s':
+            query_tokens = self.tokenizer.tokenize(query_list)
+            results, scores = self.searcher.retrieve(query_tokens, k=num)
+        else:
+            assert False, 'Invalid bm25 backend!'
+
+        if return_score:
+            return results, scores
+        else:
+            return results
+
+
+class DenseRetriever(BaseRetriever):
+    r"""Dense retriever based on pre-built faiss index."""
+
+    def __init__(self, config: dict):
+        super().__init__(config)
+        self.index = faiss.read_index(self.index_path)
+        if config["faiss_gpu"]:
+            co = faiss.GpuMultipleClonerOptions()
+            co.useFloat16 = True
+            co.shard = True
+            self.index = faiss.index_cpu_to_all_gpus(self.index, co=co)
+
+        self.corpus = load_corpus(self.corpus_path)
+
+        if config["use_sentence_transformer"]:
+            self.encoder = STEncoder(
+                model_name=self.retrieval_method,
+                model_path=config["retrieval_model_path"],
+                max_length=config["retrieval_query_max_length"],
+                use_fp16=config["retrieval_use_fp16"],
+            )
+        else:
+            self.encoder = Encoder(
+                model_name=self.retrieval_method,
+                model_path=config["retrieval_model_path"],
+                pooling_method=config["retrieval_pooling_method"],
+                max_length=config["retrieval_query_max_length"],
+                use_fp16=config["retrieval_use_fp16"],
+            )
+        self.topk = config["retrieval_topk"]
+        self.batch_size = self.config["retrieval_batch_size"]
+
+    def _search(self, query: str, num: int = None, return_score=False):
+        if num is None:
+            num = self.topk
+        query_emb = self.encoder.encode(query)
+        scores, idxs = self.index.search(query_emb, k=num)
+        scores = scores.tolist()
+        idxs = idxs[0]
+        scores = scores[0]
+
+        results = load_docs(self.corpus, idxs)
+        if return_score:
+            return results, scores
+        else:
+            return results
+
+    def _batch_search(self, query_list: List[str], num: int = None, return_score=False):
+        if isinstance(query_list, str):
+            query_list = [query_list]
+        if num is None:
+            num = self.topk
+
+        batch_size = self.batch_size
+
+        results = []
+        scores = []
+
+        for start_idx in tqdm(range(0, len(query_list), batch_size), desc="Retrieval process: "):
+            query_batch = query_list[start_idx : start_idx + batch_size]
+            batch_emb = self.encoder.encode(query_batch)
+            batch_scores, batch_idxs = self.index.search(batch_emb, k=num)
+            batch_scores = batch_scores.tolist()
+            batch_idxs = batch_idxs.tolist()
+
+            flat_idxs = sum(batch_idxs, [])
+            batch_results = load_docs(self.corpus, flat_idxs)
+            batch_results = [batch_results[i * num : (i + 1) * num] for i in range(len(batch_idxs))]
+
+            scores.extend(batch_scores)
+            results.extend(batch_results)
+
+        if return_score:
+            return results, scores
+        else:
+            return results

flashrag/retriever/utils.py

@@ -0,0 +1,49 @@
+import json
+import datasets
+from transformers import AutoTokenizer, AutoModel, AutoConfig
+
+
+def load_model(model_path: str, use_fp16: bool = False):
+    model_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+    model = AutoModel.from_pretrained(model_path, trust_remote_code=True)
+    model.eval()
+    model.cuda()
+    if use_fp16:
+        model = model.half()
+    tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True, trust_remote_code=True)
+
+    return model, tokenizer
+
+
+def pooling(pooler_output, last_hidden_state, attention_mask=None, pooling_method="mean"):
+    if pooling_method == "mean":
+        last_hidden = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
+        return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
+    elif pooling_method == "cls":
+        return last_hidden_state[:, 0]
+    elif pooling_method == "pooler":
+        return pooler_output
+    else:
+        raise NotImplementedError("Pooling method not implemented!")
+
+
+def load_corpus(corpus_path: str):
+    corpus = datasets.load_dataset("json", data_files=corpus_path, split="train")
+    return corpus
+
+
+def read_jsonl(file_path):
+    with open(file_path, "r") as f:
+        while True:
+            new_line = f.readline()
+            if not new_line:
+                return
+            new_item = json.loads(new_line)
+
+            yield new_item
+
+
+def load_docs(corpus, doc_idxs):
+    results = [corpus[int(idx)] for idx in doc_idxs]
+
+    return results

flashrag/utils/__init__.py

@@ -0,0 +1,2 @@
+from flashrag.utils.utils import *
+from flashrag.utils.pred_parse import *

flashrag/utils/constants.py

@@ -0,0 +1,51 @@
+OPENAI_MODEL_DICT = {
+    # chat
+    "gpt-4o": "o200k_base",
+    "gpt-4": "cl100k_base",
+    "gpt-3.5-turbo": "cl100k_base",
+    "gpt-3.5": "cl100k_base",  # Common shorthand
+    "gpt-35-turbo": "cl100k_base",  # Azure deployment name
+    # base
+    "davinci-002": "cl100k_base",
+    "babbage-002": "cl100k_base",
+    # embeddings
+    "text-embedding-ada-002": "cl100k_base",
+    "text-embedding-3-small": "cl100k_base",
+    "text-embedding-3-large": "cl100k_base",
+    # DEPRECATED MODELS
+    # text (DEPRECATED)
+    "text-davinci-003": "p50k_base",
+    "text-davinci-002": "p50k_base",
+    "text-davinci-001": "r50k_base",
+    "text-curie-001": "r50k_base",
+    "text-babbage-001": "r50k_base",
+    "text-ada-001": "r50k_base",
+    "davinci": "r50k_base",
+    "curie": "r50k_base",
+    "babbage": "r50k_base",
+    "ada": "r50k_base",
+    # code (DEPRECATED)
+    "code-davinci-002": "p50k_base",
+    "code-davinci-001": "p50k_base",
+    "code-cushman-002": "p50k_base",
+    "code-cushman-001": "p50k_base",
+    "davinci-codex": "p50k_base",
+    "cushman-codex": "p50k_base",
+    # edit (DEPRECATED)
+    "text-davinci-edit-001": "p50k_edit",
+    "code-davinci-edit-001": "p50k_edit",
+    # old embeddings (DEPRECATED)
+    "text-similarity-davinci-001": "r50k_base",
+    "text-similarity-curie-001": "r50k_base",
+    "text-similarity-babbage-001": "r50k_base",
+    "text-similarity-ada-001": "r50k_base",
+    "text-search-davinci-doc-001": "r50k_base",
+    "text-search-curie-doc-001": "r50k_base",
+    "text-search-babbage-doc-001": "r50k_base",
+    "text-search-ada-doc-001": "r50k_base",
+    "code-search-babbage-code-001": "r50k_base",
+    "code-search-ada-code-001": "r50k_base",
+    # open source
+    "gpt2": "gpt2",
+    "gpt-2": "gpt2",  # Maintains consistency with gpt-4
+}

flashrag/utils/pred_parse.py

@@ -0,0 +1,25 @@
+def selfask_pred_parse(pred):
+    """Parsing the prediction results of self-ask format."""
+    FINAL_ANSWER_PREFIX = "So the final answer is: "
+
+    lines = pred.split("\n")
+    answer = ""
+    for line in lines:
+        if FINAL_ANSWER_PREFIX in line:
+            answer = line.split(FINAL_ANSWER_PREFIX)[1].strip()
+            break
+
+    return answer
+
+
+def ircot_pred_parse(pred):
+    FINAL_ANSWER_PREFIX = "So the answer is:"
+    if FINAL_ANSWER_PREFIX in pred:
+        answer = pred.split(FINAL_ANSWER_PREFIX)[1].strip()
+    else:
+        answer = pred
+    return answer
+
+
+def basic_pred_parse(pred):
+    return pred.split("\n")[0].strip()

flashrag/utils/utils.py

@@ -0,0 +1,135 @@
+import os
+import importlib
+from transformers import AutoConfig
+from flashrag.dataset.dataset import Dataset
+
+
+def get_dataset(config):
+    """Load dataset from config."""
+
+    dataset_path = config["dataset_path"]
+    all_split = config["split"]
+
+    split_dict = {split: None for split in all_split}
+
+    for split in all_split:
+        split_path = os.path.join(dataset_path, f"{split}.jsonl")
+        if not os.path.exists(split_path):
+            print(f"{split} file not exists!")
+            continue
+        if split in ["test", "val", "dev"]:
+            split_dict[split] = Dataset(
+                config, split_path, sample_num=config["test_sample_num"], random_sample=config["random_sample"]
+            )
+        else:
+            split_dict[split] = Dataset(config, split_path)
+
+    return split_dict
+
+
+def get_generator(config, **params):
+    """Automatically select generator class based on config."""
+    if config["framework"] == "vllm":
+        return getattr(importlib.import_module("flashrag.generator"), "VLLMGenerator")(config, **params)
+    elif config["framework"] == "fschat":
+        return getattr(importlib.import_module("flashrag.generator"), "FastChatGenerator")(config, **params)
+    elif config["framework"] == "hf":
+        model_config = AutoConfig.from_pretrained(config["generator_model_path"])
+        arch = model_config.architectures[0]
+        if "t5" in arch.lower() or "bart" in arch.lower():
+            return getattr(importlib.import_module("flashrag.generator"), "EncoderDecoderGenerator")(config, **params)
+        else:
+            return getattr(importlib.import_module("flashrag.generator"), "HFCausalLMGenerator")(config, **params)
+    elif config["framework"] == "openai":
+        return getattr(importlib.import_module("flashrag.generator"), "OpenaiGenerator")(config, **params)
+    else:
+        raise NotImplementedError
+
+
+def get_retriever(config):
+    r"""Automatically select retriever class based on config's retrieval method
+
+    Args:
+        config (dict): configuration with 'retrieval_method' key
+
+    Returns:
+        Retriever: retriever instance
+    """
+    if config["retrieval_method"] == "bm25":
+        return getattr(importlib.import_module("flashrag.retriever"), "BM25Retriever")(config)
+    else:
+        return getattr(importlib.import_module("flashrag.retriever"), "DenseRetriever")(config)
+
+
+def get_reranker(config):
+    model_path = config["rerank_model_path"]
+    # get model config
+    model_config = AutoConfig.from_pretrained(model_path)
+    arch = model_config.architectures[0]
+    if "forsequenceclassification" in arch.lower():
+        return getattr(importlib.import_module("flashrag.retriever"), "CrossReranker")(config)
+    else:
+        return getattr(importlib.import_module("flashrag.retriever"), "BiReranker")(config)
+
+
+def get_judger(config):
+    judger_name = config["judger_name"]
+    if "skr" in judger_name.lower():
+        return getattr(importlib.import_module("flashrag.judger"), "SKRJudger")(config)
+    elif "adaptive" in judger_name.lower():
+        return getattr(importlib.import_module("flashrag.judger"), "AdaptiveJudger")(config)
+    else:
+        assert False, "No implementation!"
+
+
+def get_refiner(config, retriever=None, generator=None):
+    # 预定义默认路径字典
+    DEFAULT_PATH_DICT = {
+        "recomp_abstractive_nq": "fangyuan/nq_abstractive_compressor",
+        "recomp:abstractive_tqa": "fangyuan/tqa_abstractive_compressor",
+        "recomp:abstractive_hotpotqa": "fangyuan/hotpotqa_abstractive",
+    }
+    REFINER_MODULE = importlib.import_module("flashrag.refiner")
+
+    refiner_name = config["refiner_name"]
+    refiner_path = (
+        config["refiner_model_path"]
+        if config["refiner_model_path"] is not None
+        else DEFAULT_PATH_DICT.get(refiner_name, None)
+    )
+
+    try:
+        model_config = AutoConfig.from_pretrained(refiner_path)
+        arch = model_config.architectures[0].lower()
+    except:
+        model_config, arch = "", ""
+
+    if "recomp" in refiner_name or "bert" in arch:
+        if model_config.model_type == "t5":
+            refiner_class = "AbstractiveRecompRefiner"
+        else:
+            refiner_class = "ExtractiveRefiner"
+    elif "lingua" in refiner_name:
+        refiner_class = "LLMLinguaRefiner"
+    elif "selective-context" in refiner_name or "sc" in refiner_name:
+        refiner_class = "SelectiveContextRefiner"
+    elif "kg-trace" in refiner_name:
+        return getattr(REFINER_MODULE, "KGTraceRefiner")(config, retriever, generator)
+    else:
+        raise ValueError("No implementation!")
+
+    return getattr(REFINER_MODULE, refiner_class)(config)
+
+
+def hash_object(o) -> str:
+    """Returns a character hash code of arbitrary Python objects."""
+    import hashlib
+    import io
+    import dill
+    import base58
+
+    m = hashlib.blake2b()
+    with io.BytesIO() as buffer:
+        dill.dump(o, buffer)
+        m.update(buffer.getbuffer())
+        return base58.b58encode(m.digest()).decode()

flashrag_dev-0.1.1.dist-info/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Jiajie Jin
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.