reranker integration (optional)

app/main.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from .retriever import retrieve_context, get_vectorstore
+from .retriever import get_context, get_vectorstore
 
 # Initialize vector store at startup
 print("Initializing vector store connection...")
@@ -40,7 +40,8 @@ def retrieve_mcp(
     year = [y.strip() for y in year_filter.split(",") if y.strip()] if year_filter else None
     
     # Call retriever function and return raw results
-    results = retrieve_context(
+    results = get_context(
+        vectorstore=vectorstore,
         query=query,
         reports=reports,
         sources=sources,
@@ -64,7 +65,8 @@ def retrieve_ui(query, reports_filter="", sources_filter="", subtype_filter="",
     year = [y.strip() for y in year_filter.split(",") if y.strip()] if year_filter else None
     
     # Call retriever function
-    results = retrieve_context(
+    results = get_context(
+        vectorstore=vectorstore,
         query=query,
         reports=reports,
         sources=sources,

app/retriever.py

@@ -1,6 +1,8 @@
 from typing import List, Dict, Any, Optional
 from qdrant_client.http import models as rest
 from langchain.schema import Document
+from langchain_community.cross_encoders import HuggingFaceCrossEncoder
+from langchain.retrievers.document_compressors import CrossEncoderReranker
 from .utils import getconfig
 from .vectorstore_interface import create_vectorstore, VectorStoreInterface
 import logging
@@ -12,18 +14,39 @@ config = getconfig("params.cfg")
 RETRIEVER_TOP_K = int(config.get("retriever", "TOP_K"))
 SCORE_THRESHOLD = float(config.get("retriever", "SCORE_THRESHOLD"))
 
-# Initialize vector store connection at module import time
-logging.info("Initializing vector store connection...")
-vectorstore = create_vectorstore(config)
-logging.info("Vector store connection initialized successfully")
+# Reranker settings from config
+RERANKER_ENABLED = config.getboolean("reranker", "ENABLED", fallback=False)
+RERANKER_MODEL = config.get("reranker", "MODEL_NAME", fallback="cross-encoder/ms-marco-MiniLM-L-6-v2")
+RERANKER_TOP_K = int(config.get("reranker", "TOP_K", fallback=5))
+RERANKER_TOP_K_SCALE_FACTOR = int(config.get("reranker", "TOP_K_SCALE_FACTOR", fallback=2))
+
+# # Initialize vector store connection at module import time
+# logging.info("Initializing vector store connection...")
+# vectorstore = create_vectorstore(config)
+# logging.info("Vector store connection initialized successfully")
+
+# Initialize reranker if enabled
+reranker = None
+if RERANKER_ENABLED:
+    try:
+        logging.info(f"Initializing reranker with model: {RERANKER_MODEL}")
+        model = HuggingFaceCrossEncoder(model_name=RERANKER_MODEL)
+        reranker = CrossEncoderReranker(model=model, top_n=RERANKER_TOP_K)
+        logging.info("Reranker initialized successfully")
+    except Exception as e:
+        logging.error(f"Failed to initialize reranker: {str(e)}")
+        reranker = None
 
 def get_vectorstore() -> VectorStoreInterface:
     """
-    Return the pre-initialized vector store connection.
+    Create and return a vector store connection.
     
     Returns:
         VectorStoreInterface instance
     """
+    logging.info("Initializing vector store connection...")
+    vectorstore = create_vectorstore(config)
+    logging.info("Vector store connection initialized successfully")
     return vectorstore
 
 def create_filter(
@@ -89,48 +112,111 @@ def create_filter(
         return rest.Filter(must=conditions)
     return None
 
-def retrieve_context(
+def rerank_documents(query: str, documents: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """
+    Rerank documents using cross-encoder (specify in params.cfg)
+    
+    Args:
+        query: The search query
+        documents: List of documents to rerank 
+    
+    Returns:
+        Reranked list of documents in original format
+    """
+    if not reranker or not documents:
+        return documents
+    
+    try:
+        logging.info(f"Starting reranking of {len(documents)} documents")
+        
+        # Convert to LangChain Document format using correct keys (need to review this later for portability)
+        langchain_docs = []
+        for doc in documents:
+            # Use correct keys from the data storage test module
+            content = doc.get('answer', '')
+            metadata = doc.get('answer_metadata', {})
+            
+            if not content:
+                logging.warning(f"Document missing content: {doc}")
+                continue
+                
+            langchain_doc = Document(
+                page_content=content,
+                metadata=metadata
+            )
+            langchain_docs.append(langchain_doc)
+        
+        if not langchain_docs:
+            logging.warning("No valid documents found for reranking")
+            return documents
+        
+        # Rerank documents
+        logging.info(f"Reranking {len(langchain_docs)} documents")
+        reranked_docs = reranker.compress_documents(langchain_docs, query)
+        
+        # Convert back to original format
+        result = []
+        for doc in reranked_docs:
+            result.append({
+                'answer': doc.page_content,
+                'answer_metadata': doc.metadata,
+            })
+        
+        logging.info(f"Successfully reranked {len(documents)} documents to top {len(result)}")
+        return result
+        
+    except Exception as e:
+        logging.error(f"Error during reranking: {str(e)}")
+        # Return original documents if reranking fails
+        return documents
+
+def get_context(
+    vectorstore: VectorStoreInterface,
     query: str,
     reports: List[str] = None,
     sources: str = None,
     subtype: str = None,
-    year: List[str] = None,
-    top_k: int = None
+    year: List[str] = None
 ) -> List[Dict[str, Any]]:
     """
-    Retrieve semantically similar documents from the vector database.
+    Retrieve semantically similar documents from the vector database with optional reranking.
     
     Args:
+        vectorstore: The vector store interface to search
         query: The search query
-        vectorstore: Pre-initialized vector store instance
         reports: List of specific report filenames to search within
         sources: Source type to filter by
         subtype: Document subtype to filter by
         year: List of years to filter by
-        top_k: Number of results to return (defaults to config value)
     
     Returns:
-        List of dictionaries with 'page_content' and 'metadata' keys
+        List of dictionaries with 'answer', 'answer_metadata', and 'score' keys
     """
     try:
-        # Use the passed vector store instead of calling get_vectorstore()
-        k = top_k or RETRIEVER_TOP_K
+        # Use a higher k for initial retrieval if reranking is enabled (more candidates docs)
+        top_k = RETRIEVER_TOP_K
+        if RERANKER_ENABLED and reranker:
+            top_k = top_k * RERANKER_TOP_K_SCALE_FACTOR
+            logging.info(f"Reranking enabled, retrieving {top_k} candidates")
         
-        # For Hugging Face Spaces, we pass the model name from config
         search_kwargs = {
             "model_name": config.get("embeddings", "MODEL_NAME")
         }
         
-        # Note: Filtering is currently limited for Hugging Face Spaces
-        # as the API doesn't expose filtering capabilities
-        if any([reports, sources, subtype, year]):
-            logging.warning("Filtering not supported for Hugging Face Spaces API")
-        
-        # Perform retrieval
-        retrieved_docs = vectorstore.search(query, k, **search_kwargs)
+        # Perform initial retrieval
+        retrieved_docs = vectorstore.search(query, top_k, **search_kwargs)
         
         logging.info(f"Retrieved {len(retrieved_docs)} documents for query: {query[:50]}...")
         
+        # Apply reranking if enabled
+        if RERANKER_ENABLED and reranker and retrieved_docs:
+            logging.info("Applying reranking...")
+            retrieved_docs = rerank_documents(query, retrieved_docs)
+            
+            # Trim to final desired k
+            retrieved_docs = retrieved_docs[:RERANKER_TOP_K]
+        
+        logging.info(f"Returning {len(retrieved_docs)} final documents")
         return retrieved_docs
         
     except Exception as e:

params.cfg

@@ -1,7 +1,3 @@
-[retriever]
-TOP_K = 10
-SCORE_THRESHOLD = 0.6
-
 [vectorstore]
 TYPE = huggingface_spaces
 SPACE_URL = GIZ/audit_data
@@ -15,4 +11,16 @@ COLLECTION_NAME = docling
 
 [embeddings]
 MODEL_NAME = BAAI/bge-m3
-# DEVICE = cpu
+# DEVICE = cpu
+
+[retriever]
+TOP_K = 10
+SCORE_THRESHOLD = 0.6
+
+[reranker]
+MODEL_NAME = cross-encoder/ms-marco-MiniLM-L-6-v2
+TOP_K = 5
+ENABLED = true
+# use this to scale out the total docs retrieved prior to reranking (i.e. retriever top_k * TOP_K_SCALE_FACTOR)
+TOP_K_SCALE_FACTOR = 2
+

requirements.txt

@@ -4,4 +4,5 @@ langchain-community
 qdrant-client
 sentence-transformers
 gradio_client>=0.10.0
-huggingface_hub>=0.20.0
+huggingface_hub>=0.20.0
+torch