Merge pull request #10 from AnswerDotAI/feat/flashrank_and_mixedbread

feat: support mixedbread API and flashrank

Author: bclavie

.gitignore

@@ -4,6 +4,8 @@ __pycache__/
 *.py[cod]
 *$py.class
 
+.flashrank_cache
+
 # C extensions
 *.so
 

pyproject.toml

@@ -52,10 +52,11 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
-all = ["transformers", "torch", "litellm", "requests", "sentencepiece", "protobuf"]
+all = ["transformers", "torch", "litellm", "requests", "sentencepiece", "protobuf", "flashrank"]
 transformers = ["transformers", "torch", "sentencepiece", "protobuf"]
 api = ["requests"]
 gpt = ["litellm"]
+flashrank = ["flashrank"]
 dev = ["ruff", "isort", "pytest", "ipyprogress", "ipython", "ranx", "ir_datasets", "srsly"]
 
 [project.urls]

rerankers/models/__init__.py

@@ -31,3 +31,10 @@
     AVAILABLE_RANKERS["ColBERTRanker"] = ColBERTRanker
 except ImportError:
     pass
+
+try:
+    from rerankers.models.flashrank_ranker import FlashRankRanker
+
+    AVAILABLE_RANKERS["FlashRankRanker"] = FlashRankRanker
+except ImportError:
+    pass

rerankers/models/api_rankers.py

@@ -11,7 +11,7 @@
     "cohere": "https://api.cohere.ai/v1/rerank",
     "jina": "https://api.jina.ai/v1/rerank",
     "voyage": "https://api.voyageai.com/v1/rerank",
-    "mixedbread": NotImplemented,
+    "mixedbread.ai": "https://api.mixedbread.ai/v1/reranking",
 }
 
 
@@ -29,24 +29,36 @@ def __init__(self, model: str, api_key: str, api_provider: str, verbose: int = 1
         }
         self.url = URLS[self.api_provider]
 
+    def _get_document_text(self, r: dict) -> str:
+        if self.api_provider == "voyage":
+            return r["document"]
+        elif self.api_provider == "mixedbread.ai":
+            return r["input"]
+        else:
+            return r["document"]["text"]
+
+    def _get_score(self, r: dict) -> float:
+        if self.api_provider == "mixedbread.ai":
+            return r["score"]
+        return r["relevance_score"]
+
     def _parse_response(
         self, response: dict, doc_ids: Union[List[str], List[int]]
     ) -> RankedResults:
         ranked_docs = []
-        results_key = "results" if self.api_provider != "voyage" else "data"
+        results_key = (
+            "results"
+            if self.api_provider not in ["voyage", "mixedbread.ai"]
+            else "data"
+        )
         print(response)
 
         for i, r in enumerate(response[results_key]):
-            document_text = (
-                r["document"]
-                if self.api_provider == "voyage"
-                else r["document"]["text"]
-            )
             ranked_docs.append(
                 Result(
                     doc_id=doc_ids[r["index"]],
-                    text=document_text,
-                    score=r["relevance_score"],
+                    text=self._get_document_text(r),
+                    score=self._get_score(r),
                     rank=i + 1,
                 )
             )
@@ -67,13 +79,22 @@ def rank(
         return RankedResults(results=results, query=query, has_scores=True)
 
     def _format_payload(self, query: str, docs: List[str]) -> str:
-        top_key = "top_n" if self.api_provider != "voyage" else "top_k"
+        top_key = (
+            "top_n" if self.api_provider not in ["voyage", "mixedbread.ai"] else "top_k"
+        )
+        documents_key = "documents" if self.api_provider != "mixedbread.ai" else "input"
+        return_documents_key = (
+            "return_documents"
+            if self.api_provider != "mixedbread.ai"
+            else "return_input"
+        )
+
         payload = {
             "model": self.model,
             "query": query,
-            "documents": docs,
+            documents_key: docs,
             top_key: len(docs),
-            "return_documents": True,
+            return_documents_key: True,
         }
         return json.dumps(payload)
 

rerankers/models/flashrank_ranker.py

@@ -0,0 +1,64 @@
+from rerankers.models.ranker import BaseRanker
+
+from flashrank import Ranker, RerankRequest
+
+
+from typing import Union, List, Optional, Tuple
+from rerankers.utils import (
+    vprint,
+    ensure_docids,
+    ensure_docs_list,
+)
+from rerankers.results import RankedResults, Result
+
+
+class FlashRankRanker(BaseRanker):
+    def __init__(
+        self,
+        model_name_or_path: str,
+        verbose: int = 1,
+        cache_dir: str = "./.flashrank_cache",
+    ):
+        self.verbose = verbose
+        vprint(
+            f"Loading model FlashRank model {model_name_or_path}...", verbose=verbose
+        )
+        self.model = Ranker(model_name=model_name_or_path, cache_dir=cache_dir)
+        self.ranking_type = "pointwise"
+
+    def tokenize(self, inputs: Union[str, List[str], List[Tuple[str, str]]]):
+        return self.tokenizer(
+            inputs, return_tensors="pt", padding=True, truncation=True
+        ).to(self.device)
+
+    def rank(
+        self,
+        query: str,
+        docs: List[str],
+        doc_ids: Optional[List[Union[str, int]]] = None,
+    ) -> RankedResults:
+        docs = ensure_docs_list(docs)
+        doc_ids = ensure_docids(doc_ids, len(docs))
+        passages = [{"id": doc_id, "text": doc} for doc_id, doc in zip(doc_ids, docs)]
+
+        rerank_request = RerankRequest(query=query, passages=passages)
+        flashrank_results = self.model.rerank(rerank_request)
+
+        ranked_results = [
+            Result(
+                doc_id=result["id"],
+                text=result["text"],
+                score=result["score"],
+                rank=idx + 1,
+            )
+            for idx, result in enumerate(flashrank_results)
+        ]
+        return RankedResults(results=ranked_results, query=query, has_scores=True)
+
+    def score(self, query: str, doc: str) -> float:
+        rerank_request = RerankRequest(
+            query=query, passages=[{"id": "temp_id", "text": doc}]
+        )
+        flashrank_result = self.model.rerank(rerank_request)
+        score = flashrank_result[0]["score"]
+        return score

rerankers/reranker.py

@@ -5,8 +5,9 @@
 
 DEFAULTS = {
     "jina": {"en": "jina-reranker-v1-base-en"},
-    "cohere": {"en": "rerank-english-v2.0", "other": "rerank-multilingual-v2.0"},
+    "cohere": {"en": "rerank-english-v3.0", "other": "rerank-multilingual-v3.0"},
     "voyage": {"en": "rerank-lite-1"},
+    "mixedbread.ai": {"en": "mixedbread-ai/mxbai-rerank-large-v1"},
     "cross-encoder": {
         "en": "mixedbread-ai/mxbai-rerank-base-v1",
         "fr": "antoinelouis/crossencoder-camembert-base-mmarcoFR",
@@ -26,6 +27,7 @@
         "ja": "bclavie/JaColBERTv2",
         "es": "AdrienB134/ColBERTv2.0-spanish-mmarcoES",
     },
+    "flashrank": {"en": "ms-marco-MiniLM-L-12-v2", "other": "ms-marco-MultiBERT-L-12"},
 }
 
 DEPS_MAPPING = {
@@ -35,9 +37,10 @@
     "RankGPTRanker": "gpt",
     "APIRanker": "api",
     "ColBERTRanker": "transformers",
+    "FlashRankRanker": "flashrank",
 }
 
-PROVIDERS = ["cohere", "jina", "voyage"]
+PROVIDERS = ["cohere", "jina", "voyage", "mixedbread.ai"]
 
 
 def _get_api_provider(model_name: str, model_type: Optional[str] = None) -> str:
@@ -68,6 +71,7 @@ def _get_model_type(model_name: str, explicit_model_type: Optional[str] = None)
             "t5": "T5Ranker",
             "colbert": "ColBERTRanker",
             "cross-encoder": "TransformerRanker",
+            "flashrank": "FlashRankRanker",
         }
         return model_mapping.get(explicit_model_type, explicit_model_type)
     else:
@@ -82,12 +86,18 @@ def _get_model_type(model_name: str, explicit_model_type: Optional[str] = None)
             "cohere": "APIRanker",
             "jina": "APIRanker",
             "voyage": "APIRanker",
+            "ms-marco-minilm-l-12-v2": "FlashRankRanker",
+            "ms-marco-multibert-l-12": "FlashRankRanker",
         }
         for key, value in model_mapping.items():
             if key in model_name:
                 return value
-        if any(
-            keyword in model_name for keyword in ["minilm", "bert", "cross-encoders/"]
+        if (
+            any(
+                keyword in model_name
+                for keyword in ["minilm", "bert", "cross-encoders/"]
+            )
+            and "/" in model_name
         ):
             return "TransformerRanker"
         print(