Cohere Rerank 3 vs BGE-Reranker-Large: Precision & Latency Benchmarks

import os
import logging
from typing import List, Dict, Any, Optional
import cohere
from sentence_transformers import CrossEncoder

logger = logging.getLogger("RAGPipeline")

class ResilientTwoStageReranker:
    def __init__(
        self,
        cohere_api_key: Optional[str] = None,
        local_model_name: str = "BAAI/bge-reranker-large",
        top_n: int = 5,
        timeout_seconds: float = 2.0
    ):
        self.top_n = top_n
        self.timeout_seconds = timeout_seconds
        
        # Initialize Cohere client if API key is present
        self.cohere_client = None
        key = cohere_api_key or os.getenv("COHERE_API_KEY")
        if key:
            self.cohere_client = cohere.ClientV2(api_key=key)
            logger.info("Initialized Cohere Rerank 3 client.")
            
        # Initialize local HuggingFace CrossEncoder as backup or primary
        logger.info(f"Loading local CrossEncoder fallback: {local_model_name}")
        self.local_model = CrossEncoder(local_model_name, max_length=512)

    def rerank(
        self,
        query: str,
        candidate_docs: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """
        Reranks Stage 1 candidates using Cohere Rerank 3 with automatic
        fallback to local BGE-Reranker-Large on API failures.
        """
        if not candidate_docs:
            return []

        doc_texts = [d["text"] for d in candidate_docs]

        # Primary Path: Cohere Rerank 3
        if self.cohere_client:
            try:
                response = self.cohere_client.rerank(
                    model="rerank-v3.5",
                    query=query,
                    documents=doc_texts,
                    top_n=self.top_n,
                    return_documents=False
                )
                
                reranked_results = []
                for hit in response.results:
                    original_doc = candidate_docs[hit.index].copy()
                    original_doc["rerank_score"] = float(hit.relevance_score)
                    original_doc["reranker_engine"] = "cohere-rerank-3.5"
                    reranked_results.append(original_doc)
                    
                return reranked_results

            except Exception as e:
                logger.warning(f"Cohere API call failed or timed out: {e}. Executing local BGE fallback.")

        # Fallback / Local Path: BGE-Reranker-Large
        query_doc_pairs = [[query, text] for text in doc_texts]
        scores = self.local_model.predict(query_doc_pairs)
        
        # Pair docs with predicted scores and sort descending
        scored_docs = []
        for idx, score in enumerate(scores):
            doc_copy = candidate_docs[idx].copy()
            doc_copy["rerank_score"] = float(score)
            doc_copy["reranker_engine"] = "bge-reranker-large"
            scored_docs.append(doc_copy)

        scored_docs.sort(key=lambda x: x["rerank_score"], reverse=True)
        return scored_docs[:self.top_n]

# Example Demonstration
if __name__ == "__main__":
    reranker = ResilientTwoStageReranker()
    test_query = "What are the tax implications of Section 127 FSIE in Malaysia?"
    test_candidates = [
        {"id": 1, "text": "Malaysia offers sunny tropical weather and beaches in Penang and Langkawi."},
        {"id": 2, "text": "Under Section 127 of the Income Tax Act 1967, qualifying foreign-sourced income is exempt from local taxation."},
        {"id": 3, "text": "The DE Rantau visa requires a minimum annual income of $24,000 for foreign tech freelancers."},
    ]
    results = reranker.rerank(test_query, test_candidates)
    for rank, doc in enumerate(results, 1):
        print(f"Rank {rank} [Score: {doc['rerank_score']:.4f} via {doc['reranker_engine']}]: {doc['text']}")