updated latest dense retriever models including ST, huggingface and LLM2Vec

Author: thakur-nandan

README.md

@@ -83,7 +83,7 @@ Tested with python versions 3.9+
 
 - Preprocess your own IR dataset or use one of the already-preprocessed 17 benchmark datasets
 - Wide settings included, covers diverse benchmarks useful for both academia and industry
-- Includes well-known retrieval architectures (lexical, dense, sparse and reranking-based)
+- Evaluates well-known retrieval architectures (lexical, dense, sparse and reranking-based)
 - Add and evaluate your own model in a easy framework using different state-of-the-art evaluation metrics
 
 ## :beers: Quick Example
@@ -132,14 +132,15 @@ results = retriever.retrieve(corpus, queries)
 
 #### Evaluate your model with NDCG@k, MAP@K, Recall@K and Precision@K  where k = [1,3,5,10,100,1000]
 ndcg, _map, recall, precision = retriever.evaluate(qrels, results, retriever.k_values)
+mrr = retriever.evaluate_custom(qrels, results, retriever.k_values, metric="mrr")
 
 ### If you want to save your results and runfile (useful for reranking)
 results_dir = os.path.join(pathlib.Path(__file__).parent.absolute(), "results")
 os.makedirs(results_dir, exist_ok=True)
 
 #### Save the evaluation runfile & results
 util.save_runfile(os.path.join(results_dir, f"{dataset}.run.trec"), results)
-util.save_results(os.path.join(results_dir, f"{dataset}.json"), ndcg, _map, recall, precision)
+util.save_results(os.path.join(results_dir, f"{dataset}.json"), ndcg, _map, recall, precision, mrr)
 ```
 
 ## :beers: Available Datasets
@@ -227,13 +228,22 @@ If you find this repository helpful, feel free to cite our publication [BEIR: A
 
 If you use any baseline score from the BEIR leaderboard, feel free to cite our publication [Resources for Brewing BEIR: Reproducible Reference Models and an Official Leaderboard](https://arxiv.org/abs/2306.07471)
 ```
-@misc{kamalloo2023resources,
-      title={Resources for Brewing BEIR: Reproducible Reference Models and an Official Leaderboard},
-      author={Ehsan Kamalloo and Nandan Thakur and Carlos Lassance and Xueguang Ma and Jheng-Hong Yang and Jimmy Lin},
-      year={2023},
-      eprint={2306.07471},
-      archivePrefix={arXiv},
-      primaryClass={cs.IR}
+@inproceedings{kamalloo:2024,
+    author = {Kamalloo, Ehsan and Thakur, Nandan and Lassance, Carlos and Ma, Xueguang and Yang, Jheng-Hong and Lin, Jimmy},
+    title = {Resources for Brewing BEIR: Reproducible Reference Models and Statistical Analyses},
+    year = {2024},
+    isbn = {9798400704314},
+    publisher = {Association for Computing Machinery},
+    address = {New York, NY, USA},
+    url = {https://doi.org/10.1145/3626772.3657862},
+    doi = {10.1145/3626772.3657862},
+    abstract = {BEIR is a benchmark dataset originally designed for zero-shot evaluation of retrieval models across 18 different domain/task combinations. In recent years, we have witnessed the growing popularity of models based on representation learning, which naturally begs the question: How effective are these models when presented with queries and documents that differ from the training data? While BEIR was designed to answer this question, our work addresses two shortcomings that prevent the benchmark from achieving its full potential: First, the sophistication of modern neural methods and the complexity of current software infrastructure create barriers to entry for newcomers. To this end, we provide reproducible reference implementations that cover learned dense and sparse models. Second, comparisons on BEIR are performed by reducing scores from heterogeneous datasets into a single average that is difficult to interpret. To remedy this, we present meta-analyses focusing on effect sizes across datasets that are able to accurately quantify model differences. By addressing both shortcomings, our work facilitates future explorations in a range of interesting research questions.},
+    booktitle = {Proceedings of the 47th International ACM SIGIR Conference on Research and Development in Information Retrieval},
+    pages = {1431–1440},
+    numpages = {10},
+    keywords = {domain generalization, evaluation, reproducibility},
+    location = {Washington DC, USA},
+    series = {SIGIR '24}
 }
 ```
 

beir/retrieval/models/__init__.py

@@ -2,6 +2,8 @@
 
 from .bpr import BinarySentenceBERT
 from .huggingface import HuggingFace
+from .llm2vec import LLM2Vec
+from .nvembed import NVEmbed
 from .sentence_bert import SentenceBERT
 from .sparta import SPARTA
 from .splade import SPLADE
@@ -11,6 +13,8 @@
 __all__ = [
     "BinarySentenceBERT",
     "HuggingFace",
+    "LLM2Vec",
+    "NVEmbed",
     "SentenceBERT",
     "SPARTA",
     "SPLADE",

beir/retrieval/models/huggingface.py

@@ -21,12 +21,21 @@
 POOL_FUNC = {"cls": cls_pooling, "mean": mean_pooling, "eos": eos_pooling}
 
 
-def get_peft_model(peft_model_name: str) -> PeftModel:
+def get_peft_model(peft_model_name: str, **kwargs) -> tuple[PeftModel, str]:
     config = PeftConfig.from_pretrained(peft_model_name)
-    base_model = AutoModel.from_pretrained(config.base_model_name_or_path)
+    logger.info(f"Loading Auto Model from {config.base_model_name_or_path} for PEFT model")
+    base_model = AutoModel.from_pretrained(
+        config.base_model_name_or_path,
+        device_map="auto",
+        attn_implementation=kwargs.get("attn_implementation", "eager"),
+        torch_dtype=kwargs.get("torch_dtype", "auto"),
+        trust_remote_code=True,
+        cache_dir=kwargs.get("cache_dir", None),
+    )
+    logger.info(f"Loading PEFT model from {peft_model_name}")
     model = PeftModel.from_pretrained(base_model, peft_model_name)
     model = model.merge_and_unload()
-    return model
+    return model, config.base_model_name_or_path
 
 
 class HuggingFace:
@@ -43,18 +52,23 @@ def __init__(
         **kwargs,
     ):
         self.sep = sep
-        self.tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)
-        if self.tokenizer.pad_token_id is None:
-            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
-        self.tokenizer.padding_side = "right"
-
         if peft_model_path:
-            self.model = get_peft_model(peft_model_path)
+            self.model, base_model_path = get_peft_model(peft_model_path, **kwargs)
+            self.tokenizer = AutoTokenizer.from_pretrained(base_model_path, use_fast=True)
         else:
             self.model = AutoModel.from_pretrained(
-                model_path, device_map="auto", torch_dtype=kwargs.get("torch_dtype", "auto"), trust_remote_code=True
+                model_path,
+                device_map="auto",
+                torch_dtype=kwargs.get("torch_dtype", "auto"),
+                trust_remote_code=True,
+                attn_implementation=kwargs.get("attn_implementation", "default"),
+                cache_dir=kwargs.get("cache_dir", None),
             )
+            self.tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True)
         self.model.eval()
+        if self.tokenizer.pad_token_id is None:
+            self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+        self.tokenizer.padding_side = "right"
         self.max_length = max_length if max_length else self.tokenizer.model_max_length
         self.normalize = normalize  # Normalize the embeddings
         self.append_eos_token = append_eos_token  # Add eos token to the input

beir/retrieval/models/llm2vec.py

@@ -0,0 +1,114 @@
+from __future__ import annotations
+
+import importlib.util
+import logging
+
+if importlib.util.find_spec("llm2vec") is not None:
+    from llm2vec import LLM2Vec as LLM2VecOriginal
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from tqdm.autonotebook import trange
+
+from .util import extract_corpus_sentences
+
+logger = logging.getLogger(__name__)
+
+POOLING_MODES = {
+    "mean": "mean",
+    "weighted_mean": "weighted_mean",
+    "eos": "eos_token",
+    "bos_token": "bos_token",
+    "last_token": "last_token",
+}
+
+
+class LLM2Vec:
+    def __init__(
+        self,
+        model_path: str | tuple = None,
+        max_length: int = None,
+        sep: str = " ",
+        pooling: str = "mean",
+        normalize: bool = True,
+        prompts: dict[str, str] = None,
+        peft_model_path: str = None,
+        **kwargs,
+    ):
+        self.sep = sep
+        self.normalize = normalize
+        if pooling not in POOLING_MODES:
+            raise ValueError(f"Pooling mode {pooling} not supported. Choose from {list(POOLING_MODES.keys())}")
+
+        self.model = LLM2VecOriginal.from_pretrained(
+            base_model_name_or_path=model_path,
+            peft_model_name_or_path=peft_model_path,
+            pooling_mode=POOLING_MODES[pooling],
+            max_length=max_length,
+            **kwargs,
+        )
+
+        if prompts:
+            self.query_prefix = prompts.get("query", "")
+            self.doc_prefix = prompts.get("passage", "")
+
+    def _append_eos_token(self, texts, pad_to_multiple_of: int = 16):
+        """Tokenizes the input texts and pads the tokenized input to the max_length with the eos token"""
+        collated_texts = self.tokenizer(
+            texts,
+            padding=False,
+            truncation=True,
+            max_length=self.max_length - 1 if self.append_eos_token else self.max_length,
+            return_attention_mask=False,
+            return_token_type_ids=False,
+            add_special_tokens=True,
+        )
+        collated_texts["input_ids"] = [x + [self.tokenizer.eos_token_id] for x in collated_texts["input_ids"]]
+        collated_texts = self.tokenizer.pad(
+            collated_texts,
+            padding=True,
+            pad_to_multiple_of=pad_to_multiple_of,
+            return_attention_mask=True,
+            return_tensors="pt",
+        )
+        return collated_texts
+
+    def encode_queries(self, queries: list[str], batch_size: int = 16, **kwargs) -> list[Tensor] | np.ndarray | Tensor:
+        query_embeddings = []
+
+        with torch.no_grad():
+            for start_idx in trange(0, len(queries), batch_size):
+                sub_queries = [[self.query_prefix, query] for query in queries[start_idx : start_idx + batch_size]]
+                query_embeddings += self.model.encode(sub_queries, batch_size=batch_size, show_progress_bar=False)
+
+        query_embeddings = torch.stack(query_embeddings)
+
+        if self.normalize:
+            query_embeddings = F.normalize(query_embeddings, p=2, dim=1)
+
+        return query_embeddings
+
+    def encode_corpus(
+        self, corpus: list[dict[str, str]] | dict[str, list] | list[str], batch_size: int = 8, **kwargs
+    ) -> list[Tensor] | np.ndarray | Tensor:
+        corpus_embeddings = []
+        sentences = extract_corpus_sentences(corpus=corpus, sep=self.sep)
+
+        with torch.no_grad():
+            for start_idx in trange(0, len(sentences), batch_size):
+                if self.doc_prefix:
+                    sub_sentences = [
+                        [self.doc_prefix, sentence] for sentence in sentences[start_idx : start_idx + batch_size]
+                    ]
+                else:
+                    sub_sentences = sentences[start_idx : start_idx + batch_size]
+                corpus_embeddings += self.model.encode(sub_sentences, batch_size=batch_size, show_progress_bar=False)
+
+            corpus_embeddings = torch.stack(corpus_embeddings)
+
+            if self.normalize:
+                corpus_embeddings = F.normalize(corpus_embeddings, p=2, dim=1)
+
+            return corpus_embeddings

beir/retrieval/models/nvembed.py

@@ -0,0 +1,84 @@
+from __future__ import annotations
+
+import logging
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from tqdm.autonotebook import trange
+from transformers import AutoModel
+
+from .pooling import cls_pooling, eos_pooling, mean_pooling
+from .util import extract_corpus_sentences
+
+logger = logging.getLogger(__name__)
+
+POOL_FUNC = {"cls": cls_pooling, "mean": mean_pooling, "eos": eos_pooling}
+
+
+class NVEmbed:
+    def __init__(
+        self,
+        model_path: str | tuple = None,
+        max_length: int = None,
+        sep: str = " ",
+        pooling: str = "mean",
+        normalize: bool = False,
+        prompts: dict[str, str] = None,
+        **kwargs,
+    ):
+        self.sep = sep
+        self.model = AutoModel.from_pretrained(
+            model_path, device_map="auto", torch_dtype=kwargs.get("torch_dtype", "auto"), trust_remote_code=True
+        )
+        # self.model.eval()
+        self.max_length = max_length if max_length else self.tokenizer.model_max_length
+        self.normalize = normalize  # Normalize the embeddings
+
+        if pooling not in ["cls", "mean", "eos"]:
+            raise ValueError("Supported Pooling techniques should be either 'cls', 'mean' or 'eos'")
+        self.pooling_func = POOL_FUNC[pooling]
+
+        if prompts:
+            self.query_prefix = prompts.get("query", "")
+            self.doc_prefix = prompts.get("passage", "")
+
+    def encode_queries(self, queries: list[str], batch_size: int = 16, **kwargs) -> list[Tensor] | np.ndarray | Tensor:
+        query_embeddings = []
+
+        with torch.no_grad():
+            for start_idx in trange(0, len(queries), batch_size):
+                sub_queries = [self.query_prefix + query for query in queries[start_idx : start_idx + batch_size]]
+                query_embeddings += self.model.encode(
+                    sub_queries, instruction=self.query_prefix, max_length=self.max_length
+                )
+
+        query_embeddings = torch.stack(query_embeddings)
+
+        if self.normalize:
+            query_embeddings = F.normalize(query_embeddings, p=2, dim=1)
+
+        return query_embeddings
+
+    def encode_corpus(
+        self, corpus: list[dict[str, str]] | dict[str, list] | list[str], batch_size: int = 8, **kwargs
+    ) -> list[Tensor] | np.ndarray | Tensor:
+        corpus_embeddings = []
+        sentences = extract_corpus_sentences(corpus=corpus, sep=self.sep)
+
+        with torch.no_grad():
+            for start_idx in trange(0, len(sentences), batch_size):
+                sub_sentences = [
+                    self.doc_prefix + sentence for sentence in sentences[start_idx : start_idx + batch_size]
+                ]
+                corpus_embeddings += self.model.encode(
+                    sub_sentences, instruction=self.doc_prefix, max_length=self.max_length
+                )
+
+            corpus_embeddings = torch.stack(corpus_embeddings)
+
+            if self.normalize:
+                corpus_embeddings = F.normalize(corpus_embeddings, p=2, dim=1)
+
+            return corpus_embeddings

beir/retrieval/models/sentence_bert.py

@@ -17,28 +17,42 @@ class SentenceBERT:
     def __init__(
         self,
         model_path: str | tuple = None,
+        max_length: int = None,
         sep: str = " ",
         prompts: dict[str, str] = None,
+        prompt_names: dict[str, str] = None,
         **kwargs,
     ):
         self.sep = sep
+        self.max_length = max_length
 
         if isinstance(model_path, str):
-            self.q_model = SentenceTransformer(model_path, kwargs)
+            self.q_model = SentenceTransformer(model_path, **kwargs)
             self.doc_model = self.q_model
 
         elif isinstance(model_path, tuple):
-            self.q_model = SentenceTransformer(model_path[0], kwargs)
-            self.doc_model = SentenceTransformer(model_path[1], kwargs)
+            self.q_model = SentenceTransformer(model_path[0], **kwargs)
+            self.doc_model = SentenceTransformer(model_path[1], **kwargs)
 
-        self.query_prefix = ""
-        self.doc_prefix = ""
+        if self.max_length:
+            self.q_model.max_seq_length = self.max_length
+            self.doc_model.max_seq_length = self.max_length
+
+        self.query_prefix, self.query_prompt_name = None, None
+        self.doc_prefix, self.doc_prompt_name = None, None
 
         # Checks if prompts are not set in Sentence Transformers but required during inference
         if prompts and (len(self.q_model.prompts) or len(self.doc_model.prompts) == 0):
             self.query_prefix = prompts["query"]
             self.doc_prefix = prompts["passage"]
 
+        if prompt_names:
+            self.query_prompt_name = prompt_names["query"]
+            self.doc_prompt_name = prompt_names["passage"]
+
+        logger.info(f"Query prompt: {self.query_prefix}, Passage prompt: {self.doc_prefix}")
+        logger.info(f"Query prompt name: {self.query_prompt_name}, Passage prompt name: {self.doc_prompt_name}")
+
     def get_similarity(self):
         return self.q_model.similarity
 
@@ -74,7 +88,9 @@ def stop_multi_process_pool(self, pool: dict[str, object]):
 
     def encode_queries(self, queries: list[str], batch_size: int = 16, **kwargs) -> list[Tensor] | np.ndarray | Tensor:
         return self.q_model.encode(
-            [self.query_prefix + query for query in queries],
+            queries,
+            prompt=self.query_prefix,
+            prompt_name=self.query_prompt_name,
             batch_size=batch_size,
             **kwargs,
         )
@@ -87,7 +103,9 @@ def encode_corpus(
     ) -> list[Tensor] | np.ndarray | Tensor:
         sentences = extract_corpus_sentences(corpus=corpus, sep=self.sep)
         return self.doc_model.encode(
-            [self.doc_prefix + sentence for sentence in sentences],
+            sentences,
+            prompt=self.doc_prefix,
+            prompt_name=self.doc_prompt_name,
             batch_size=batch_size,
             **kwargs,
         )