refacto reranker

.env.example

@@ -40,8 +40,8 @@ EMBEDDER_API_KEY=EMPTY
 # RERANKER
 RERANKER_ENABLED=false
 RERANKER_MODEL=Alibaba-NLP/gte-multilingual-reranker-base # or jinaai/jina-reranker-v2-base-multilingual or jinaai/jina-colbert-v2 if you want
-RERANKER_BASE_URL=http://reranker:7997 # If using infinity
-RERANKER_MODEL_TYPE=crossencoder # colbert or infinity
+RERANKER_PORT=7997 # Port for the reranker service
+# RERANKER_BASE_URL=http://reranker:7997 # you can uncomment this if you want to use a specific base URL for the reranker service. In that, comment out the reranker docker service in docker-compose.yaml
 RERANKER_TOP_K=5 # Number of documents to return after reranking. upgrade to 8 for better results if your llm has a wider context window
 
 # Prompts
@@ -53,9 +53,9 @@ XDG_CACHE_HOME=/app/model_weights
 # If using MarkerLoader
 MARKER_MAX_TASKS_PER_CHILD=100
 MARKER_MAX_PROCESSES=10
-MARKER_MIN_PROCESSES=3
-MARKER_POOL_SIZE=3
-MARKER_NUM_GPUS=0.6
+MARKER_MIN_PROCESSES=1
+# MARKER_POOL_SIZE=1 # Value au increment if you have a cluster of machines
+MARKER_NUM_GPUS=0.01
 
 # Audio
 WHISPER_MODEL=base
@@ -80,14 +80,13 @@ RAY_RUNTIME_ENV_HOOK=ray._private.runtime_env.uv_runtime_env_hook.hook
 SAVE_UPLOADED_FILES=true # usefull for chainlit source viewing
 
 
-## Variable to add to activate indexer-ui
+## Variables to add to activate indexer-ui
 # INDEXERUI_COMPOSE_FILE=extern/indexer-ui/docker-compose.yaml  # Required path to the docker-compose file
 # INDEXERUI_PORT=8067                                             # Port to expose the Indexer UI (default is 3042)
 # INDEXERUI_URL='http://X.X.X.X:INDEXERUI_PORT'                   # Base URL of the Indexer UI (required to prevent CORS issues)
 # VITE_API_BASE_URL='http://X.X.X.X:APP_PORT'                     # Base URL of your FastAPI backend. Used by the frondend
 
 
-
 ## Specific to Ray cluster
 # SHARED_ENV=/ray_mount/.env
 # MODEL_WEIGHTS_VOLUME=/ray_mount/model_weights

.hydra_config/config.yaml

@@ -43,7 +43,6 @@ vectordb:
 reranker:
   enable: ${oc.decode:${oc.env:RERANKER_ENABLED, true}}
   model_name: ${oc.env:RERANKER_MODEL, Alibaba-NLP/gte-multilingual-reranker-base}
-  reranker_type: ${oc.decode:${oc.env:RERANKER_MODEL_TYPE, colbert}}
   top_k: ${oc.decode:${oc.env:RERANKER_TOP_K, 4}}
   base_url: ${oc.env:RERANKER_BASE_URL, http://reranker:${oc.env:RERANKER_PORT, 7997}}
 
@@ -96,6 +95,6 @@ loader:
   marker_num_gpus: ${oc.decode:${oc.env:MARKER_NUM_GPUS, 0.01}}
 
 ray:
-  num_gpus: ${oc.decode:${oc.env:RAY_NUM_GPUS, 0.1}}
+  num_gpus: ${oc.decode:${oc.env:RAY_NUM_GPUS, 0.01}}
   pool_size: ${oc.decode:${oc.env:RAY_POOL_SIZE, 1}}
   max_tasks_per_worker: ${oc.decode:${oc.env:RAY_MAX_TASKS_PER_WORKER, 5}}

pyproject.toml

@@ -9,7 +9,6 @@ dependencies = [
     "chainlit>=2.2.1",
     "docling>=2.24.0",
     "einops>=0.8.1",
-    "fastembed-gpu>=0.4.2",
     "hydra-core>=1.3.2",
     "langchain-community>=0.3.18",
     "langchain-core>=0.3.39",
@@ -22,7 +21,6 @@ dependencies = [
     "markitdown>=0.0.2",
     "pydub>=0.25.1",
     "pymupdf4llm>=0.0.17",
-    "ragatouille>=0.0.9",
     "spire-doc>=13.1.0",
     "markitdown>=0.0.2",
     "openai>=1.64.0",

ragondin/components/indexer/loaders/pdf_loaders/marker.py

@@ -19,7 +19,13 @@
 config = load_config()
 
 
-@ray.remote(num_gpus=config.loader.get("marker_num_gpus", 0))
+if torch.cuda.is_available():
+    MARKER_NUM_GPUS = config.loader.get("marker_num_gpus", 0.01)
+else:  # On CPU
+    MARKER_NUM_GPUS = 0
+
+
+@ray.remote(num_gpus=MARKER_NUM_GPUS)
 class MarkerWorker:
     def __init__(self):
         import os

ragondin/components/indexer/vectordb/vectordb.py

@@ -345,6 +345,7 @@ async def async_add_documents(self, chunks: list[Document]) -> None:
                 )
 
         await self.vector_store.aadd_documents(chunks)
+        # asyncio.create_task(self.vector_store.aadd_documents(chunks)) # for prods
 
         for partition, file_id in partition_file_list:
             self.partition_file_manager.add_file_to_partition(

ragondin/components/reranker.py

@@ -1,135 +1,54 @@
 import asyncio
-import copy
-import gc
-from enum import Enum
-
-import torch
 from infinity_client import Client
 from infinity_client.api.default import rerank
 from infinity_client.models import RerankInput, ReRankResult
 from langchain_core.documents.base import Document
-from sentence_transformers import CrossEncoder
-
-from .utils import SingletonMeta
-
 
-class RerankerType(Enum):
-    CROSSENCODER = "crossencoder"
-    COLBERT = "colbert"
-    INFINITY = "infinity"
 
-
-class Reranker(metaclass=SingletonMeta):
+class Reranker:
     def __init__(self, logger, config):
-        reranker_type = config.reranker["reranker_type"]
         self.model_name = config.reranker["model_name"]
-
+        self.client = Client(base_url=config.reranker["base_url"])
         self.logger = logger
         self.semaphore = asyncio.Semaphore(
             5
         )  # Only allow 5 reranking operation at a time
-
-        self.reranker_type = RerankerType(reranker_type)
-
-        match self.reranker_type:
-            case RerankerType.CROSSENCODER:
-                self.model = CrossEncoder(
-                    model_name=self.model_name,
-                    device="cuda" if torch.cuda.is_available() else "cpu",
-                    trust_remote_code=True,
-                )
-            case RerankerType.COLBERT:
-                # Initialize ColBERT model here
-                from ragatouille import RAGPretrainedModel
-
-                self.model = RAGPretrainedModel.from_pretrained(
-                    pretrained_model_name_or_path=self.model_name
-                )
-            case RerankerType.INFINITY:
-                self.client = Client(base_url=config.reranker["base_url"])
-            case _:
-                raise ValueError(
-                    "reranker_type must be either 'crossencoder', 'colbert' or 'infinity'."
-                )
-
-        self.logger.debug(
-            "Reranker initialized", type=self.reranker_type, model_name=self.model_name
-        )
+        self.logger.debug("Reranker initialized", model_name=self.model_name)
 
     async def rerank(
-        self, query: str, documents: list[Document], top_k: int = 6
-    ) -> list[Document]:
-        self.logger.debug(
-            "Reranking documents", documents_count=len(documents), top_k=top_k
-        )
-        top_k = min(top_k, len(documents))
-
-        async with self.semaphore:
-            match self.reranker_type:
-                case RerankerType.CROSSENCODER:
-                    reranked_docs = await asyncio.to_thread(
-                        lambda: self.___crossencoder_rerank(query, documents, top_k)
-                    )
-                    return reranked_docs
-
-                case RerankerType.COLBERT:
-                    reranked_docs = await asyncio.to_thread(
-                        lambda: self.___colbert_rerank(query, documents, top_k)
-                    )
-                    return reranked_docs
-
-                case RerankerType.INFINITY:
-                    reranked_docs = await asyncio.to_thread(
-                        lambda: self.___infinity_rerank(query, documents, top_k)
-                    )
-                    return reranked_docs
-
-    def ___crossencoder_rerank(
         self, query: str, documents: list[Document], top_k: int
     ) -> list[Document]:
-        with torch.no_grad():
-            docs_txt = [doc.page_content for doc in documents]
-            results = self.model.rank(query=query, documents=docs_txt, top_k=top_k)
-
-        gc.collect()
-        torch.cuda.empty_cache()
-        return [documents[r["corpus_id"]] for r in results]
-
-    def ___colbert_rerank(
-        self, query: str, documents: list[Document], top_k: int
-    ) -> list[Document]:
-        with torch.no_grad():
-            docs_txt = [doc.page_content for doc in documents]
-            results = self.model.rerank(
-                query=query, documents=docs_txt, k=top_k, bsize="auto"
+        async with self.semaphore:
+            self.logger.debug(
+                "Reranking documents", documents_count=len(documents), top_k=top_k
             )
-
-        gc.collect()
-        torch.cuda.empty_cache()
-        return [doc for doc in original_docs(results, documents)]
-
-    def ___infinity_rerank(
-        self, query: str, documents: list[Document], top_k: int
-    ) -> list[Document]:
-        rerank_input = RerankInput.from_dict(
-            {
-                "model": self.model_name,
-                "query": query,
-                "documents": [doc.page_content for doc in documents],
-                "top_n": top_k,
-                "return_documents": True,
-            }
-        )
-        rerank_result: ReRankResult = rerank.sync(client=self.client, body=rerank_input)
-        results = [{"content": item.document} for item in rerank_result.results]
-        return [doc for doc in original_docs(results, documents)]
-
-
-def original_docs(ranked_txt: list[str], docs: list[Document]):
-    docs = copy.deepcopy(docs)
-    for doc_txt in ranked_txt:
-        for doc in docs:
-            if doc_txt["content"] == doc.page_content:
-                yield doc
-                docs.remove(doc)
-                break
+            top_k = min(top_k, len(documents))
+            rerank_input = RerankInput.from_dict(
+                {
+                    "model": self.model_name,
+                    "query": query,
+                    "documents": [doc.page_content for doc in documents],
+                    "top_n": top_k,
+                    "return_documents": True,
+                    "raw_scores": True,
+                }
+            )
+            try:
+                rerank_result: ReRankResult = await rerank.asyncio(
+                    client=self.client, body=rerank_input
+                )
+                output = []
+                for rerank_res in rerank_result.results:
+                    doc = documents[rerank_res.index]
+                    doc.metadata["relevance_score"] = rerank_res.relevance_score
+                    output.append(doc)
+                return output
+
+            except Exception as e:
+                self.logger.error(
+                    "Reranking failed",
+                    error=str(e),
+                    model_name=self.model_name,
+                    documents_count=len(documents),
+                )
+                raise e