Retriever

rankify.retrievers.retriever

METHOD_MAP = {'bm25': BM25Retriever, 'dpr-multi': DenseRetriever, 'dpr-single': DenseRetriever, 'ance-multi': ANCERetriever, 'bpr-single': DenseRetriever, 'bge': BGERetriever, 'colbert': ColBERTRetriever, 'contriever': ContrieverRetriever, 'online': OnlineRetriever, 'hyde': HydeRetriever, 'diver-dense': DiverDenseRetriever, 'diver-bm25': DiverBM25Retriever, 'reasonir': ReasonIRRetriever, 'reason-embed': ReasonEmbedRetriever, 'bge-reasoner-embed': BgeReasonerRetriever, 'unicoil': UniCOILRetriever, 'unicoil-noexp': UniCOILRetriever, 'splade-v2': SpladeV2Retriever, 'openai-embedding': APIEmbeddingRetriever, 'cohere-embedding': APIEmbeddingRetriever, 'voyage-embedding': APIEmbeddingRetriever} module-attribute

Document

Represents a document consisting of a question, answers, and contexts.

Attributes:

Name	Type	Description
question	Question	The question associated with the document.
answers	Answer	The answers to the question.
contexts	list[Context]	A list of related contexts.
reorder_contexts	list[Context] or None	A reordered list of contexts based on relevance.

Source code in rankify/dataset/dataset.py

class Document:
    """
    Represents a document consisting of a question, answers, and contexts.

    Attributes:
        question (Question): The question associated with the document.
        answers (Answer): The answers to the question.
        contexts (list[Context]): A list of related contexts.
        reorder_contexts (list[Context] or None): A reordered list of contexts based on relevance.
    """
    def __init__(self, question: Question, answers: Answer, contexts: list = None , id: int = None) -> None:
        """
        Initializes a Document instance.

        Args:
            question (Question): The question associated with the document.
            answers (Answer): The answers to the question.
            contexts (list[Context], optional): A list of contexts related to the question.

        Example:
            ```python
            q = Question("What is the capital of France?")
            a = Answer(["Paris"])
            c1 = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
            c2 = Context(score=0.5, has_answer=False, id=2, title="Berlin", text="Berlin is the capital of Germany.")
            d = Document(question=q, answers=a, contexts=[c1, c2])
            print(d)
            ```
        """
        self.question: Question = question
        self.answers: Answer = answers
        self.contexts: List[Context] = contexts
        self.reorder_contexts: List[Context] = None
        self.id = str(id) 

    @classmethod
    def from_dict(cls, data: dict,n_docs:int=100) -> 'Document':
        """
        Creates a Document instance from a dictionary.

        Args:
            data (dict): A dictionary containing the question, answers, and contexts.
            n_docs (int, optional): The number of contexts to include. Defaults to 100.

        Returns:
            Document: A new Document instance.

        Example:
            ```python
            data = {
                "question": "What is the capital of France?",
                "answers": ["Paris"],
                "ctxs": [
                    {"score": 0.9, "has_answer": True, "id": 1, "title": "Paris", "text": "The capital of France is Paris."},
                    {"score": 0.5, "has_answer": False, "id": 2, "title": "Berlin", "text": "Berlin is the capital of Germany."}
                ]
            }
            d = Document.from_dict(data)
            print(d.question)
            ```
        """
        question = Question(data["question"])
        if "answers" in data:
            answers = Answer(data["answers"])
        else:
            answers =Answer('')

        if "query_id" in data:
            id = data["query_id"]
        else:
            id = None
        contexts = [Context(**ctx) for ctx in data["ctxs"][:n_docs]]
        return cls(question, answers, contexts, id=id)

    def to_dict(self) -> Dict[str, Optional[object]]:
        """
        Converts the document into a dictionary representation.

        Returns:
            dict: A dictionary containing the question, answers, and contexts.
        """
        return {
            "question": self.question.question,
            "answers": self.answers.answers,
            "contexts": [ctx.to_dict() for ctx in self.contexts]
        }
    def to_dict_reoreder(self) -> Dict[str,Optional[object]]:
        return {
            "question" : self.question.question,
            "answers" : self.answers.answers,
            "contexts" : [ctx.to_dict() for ctx in self.reorder_contexts]
        }
    def __str__(self) -> str:
        """
        Returns a string representation of the Document instance.

        Returns:
            str: The formatted document information.

        Example:
            ```python
            d = Document(Question("What is the capital of France?"), Answer(["Paris"]))
            print(d)
            ```
        """
        contexts_str = "\n\n".join([str(ctx) for ctx in self.contexts])
        reorder_contexts_str= ''
        if self.reorder_contexts is not None:
            reorder_contexts_str = "\n\n".join([str(ctx) for ctx in self.reorder_contexts])
        return f"{self.question}\n\n{self.answers}\n\nContext: \n\n{contexts_str}\nReorder contexts: \n\n{reorder_contexts_str}"

__init__(question, answers, contexts=None, id=None)

Initializes a Document instance.

Parameters:

Name	Type	Description	Default
question	Question	The question associated with the document.	required
answers	Answer	The answers to the question.	required
contexts	list[Context]	A list of contexts related to the question.	None

Example

q = Question("What is the capital of France?")
a = Answer(["Paris"])
c1 = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
c2 = Context(score=0.5, has_answer=False, id=2, title="Berlin", text="Berlin is the capital of Germany.")
d = Document(question=q, answers=a, contexts=[c1, c2])
print(d)

Source code in rankify/dataset/dataset.py

def __init__(self, question: Question, answers: Answer, contexts: list = None , id: int = None) -> None:
    """
    Initializes a Document instance.

    Args:
        question (Question): The question associated with the document.
        answers (Answer): The answers to the question.
        contexts (list[Context], optional): A list of contexts related to the question.

    Example:
        ```python
        q = Question("What is the capital of France?")
        a = Answer(["Paris"])
        c1 = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
        c2 = Context(score=0.5, has_answer=False, id=2, title="Berlin", text="Berlin is the capital of Germany.")
        d = Document(question=q, answers=a, contexts=[c1, c2])
        print(d)
        ```
    """
    self.question: Question = question
    self.answers: Answer = answers
    self.contexts: List[Context] = contexts
    self.reorder_contexts: List[Context] = None
    self.id = str(id) 

from_dict(data, n_docs=100) classmethod

Creates a Document instance from a dictionary.

Parameters:

Name	Type	Description	Default
data	dict	A dictionary containing the question, answers, and contexts.	required
n_docs	int	The number of contexts to include. Defaults to 100.	100

Returns:

Name	Type	Description
Document	Document	A new Document instance.

Example

data = {
    "question": "What is the capital of France?",
    "answers": ["Paris"],
    "ctxs": [
        {"score": 0.9, "has_answer": True, "id": 1, "title": "Paris", "text": "The capital of France is Paris."},
        {"score": 0.5, "has_answer": False, "id": 2, "title": "Berlin", "text": "Berlin is the capital of Germany."}
    ]
}
d = Document.from_dict(data)
print(d.question)

Source code in rankify/dataset/dataset.py

@classmethod
def from_dict(cls, data: dict,n_docs:int=100) -> 'Document':
    """
    Creates a Document instance from a dictionary.

    Args:
        data (dict): A dictionary containing the question, answers, and contexts.
        n_docs (int, optional): The number of contexts to include. Defaults to 100.

    Returns:
        Document: A new Document instance.

    Example:
        ```python
        data = {
            "question": "What is the capital of France?",
            "answers": ["Paris"],
            "ctxs": [
                {"score": 0.9, "has_answer": True, "id": 1, "title": "Paris", "text": "The capital of France is Paris."},
                {"score": 0.5, "has_answer": False, "id": 2, "title": "Berlin", "text": "Berlin is the capital of Germany."}
            ]
        }
        d = Document.from_dict(data)
        print(d.question)
        ```
    """
    question = Question(data["question"])
    if "answers" in data:
        answers = Answer(data["answers"])
    else:
        answers =Answer('')

    if "query_id" in data:
        id = data["query_id"]
    else:
        id = None
    contexts = [Context(**ctx) for ctx in data["ctxs"][:n_docs]]
    return cls(question, answers, contexts, id=id)

to_dict()

Converts the document into a dictionary representation.

Returns:

Name	Type	Description
dict	Dict[str, Optional[object]]	A dictionary containing the question, answers, and contexts.

Source code in rankify/dataset/dataset.py

def to_dict(self) -> Dict[str, Optional[object]]:
    """
    Converts the document into a dictionary representation.

    Returns:
        dict: A dictionary containing the question, answers, and contexts.
    """
    return {
        "question": self.question.question,
        "answers": self.answers.answers,
        "contexts": [ctx.to_dict() for ctx in self.contexts]
    }

__str__()

Returns a string representation of the Document instance.

Returns:

Name	Type	Description
str	str	The formatted document information.

Example

d = Document(Question("What is the capital of France?"), Answer(["Paris"]))
print(d)

Source code in rankify/dataset/dataset.py

def __str__(self) -> str:
    """
    Returns a string representation of the Document instance.

    Returns:
        str: The formatted document information.

    Example:
        ```python
        d = Document(Question("What is the capital of France?"), Answer(["Paris"]))
        print(d)
        ```
    """
    contexts_str = "\n\n".join([str(ctx) for ctx in self.contexts])
    reorder_contexts_str= ''
    if self.reorder_contexts is not None:
        reorder_contexts_str = "\n\n".join([str(ctx) for ctx in self.reorder_contexts])
    return f"{self.question}\n\n{self.answers}\n\nContext: \n\n{contexts_str}\nReorder contexts: \n\n{reorder_contexts_str}"

BaseRetriever

Bases: ABC

Abstract base class for all retrieval methods in the rankify framework.

This class defines the common interface that all retrievers must implement, ensuring consistency across different retrieval methods (BM25, DPR, etc.).

Source code in rankify/retrievers/base_retriever.py

class BaseRetriever(ABC):
    """
    Abstract base class for all retrieval methods in the rankify framework.

    This class defines the common interface that all retrievers must implement,
    ensuring consistency across different retrieval methods (BM25, DPR, etc.).
    """

    def __init__(self, n_docs: int = 10, batch_size: int = 36, threads: int = 30, **kwargs):
        """
        Initialize the base retriever.

        Args:
            n_docs (int): Number of documents to retrieve per query
            batch_size (int): Number of queries to process in a batch
            threads (int): Number of parallel threads for retrieval
            **kwargs: Additional parameters specific to each retriever
        """
        self.n_docs = n_docs
        self.batch_size = batch_size
        self.threads = threads
        self.config = kwargs

    @abstractmethod
    def _initialize_searcher(self) -> Any:
        """Initialize the specific searcher for this retriever type."""
        pass

    @abstractmethod
    def retrieve(self, documents: List[Document]) -> List[Document]:
        """
        Retrieve relevant contexts for the given documents.

        Args:
            documents (List[Document]): List of documents containing queries

        Returns:
            List[Document]: Documents updated with retrieved contexts
        """
        pass

    def _preprocess_query(self, query: str) -> str:
        """
        Preprocess query text. Can be overridden by specific retrievers.

        Args:
            query (str): Raw query text

        Returns:
            str: Preprocessed query text
        """
        return query.strip()

__init__(n_docs=10, batch_size=36, threads=30, **kwargs)

Initialize the base retriever.

Parameters:

Name	Type	Description	Default
n_docs	int	Number of documents to retrieve per query	10
batch_size	int	Number of queries to process in a batch	36
threads	int	Number of parallel threads for retrieval	30
**kwargs		Additional parameters specific to each retriever	{}

Source code in rankify/retrievers/base_retriever.py

def __init__(self, n_docs: int = 10, batch_size: int = 36, threads: int = 30, **kwargs):
    """
    Initialize the base retriever.

    Args:
        n_docs (int): Number of documents to retrieve per query
        batch_size (int): Number of queries to process in a batch
        threads (int): Number of parallel threads for retrieval
        **kwargs: Additional parameters specific to each retriever
    """
    self.n_docs = n_docs
    self.batch_size = batch_size
    self.threads = threads
    self.config = kwargs

retrieve(documents) abstractmethod

Retrieve relevant contexts for the given documents.

Parameters:

Name	Type	Description	Default
documents	List[Document]	List of documents containing queries	required

Returns:

Type	Description
List[Document]	List[Document]: Documents updated with retrieved contexts

Source code in rankify/retrievers/base_retriever.py

@abstractmethod
def retrieve(self, documents: List[Document]) -> List[Document]:
    """
    Retrieve relevant contexts for the given documents.

    Args:
        documents (List[Document]): List of documents containing queries

    Returns:
        List[Document]: Documents updated with retrieved contexts
    """
    pass

BM25Retriever

Bases: BaseRetriever

BM25 retriever implementation using Pyserini's LuceneSearcher.

Implements probabilistic ranking model BM25 for document retrieval.

Source code in rankify/retrievers/bm25_retriever.py

class BM25Retriever(BaseRetriever):
    """
    BM25 retriever implementation using Pyserini's LuceneSearcher.

    Implements probabilistic ranking model BM25 for document retrieval.
    """

    def __init__(self, index_type: str = "wiki", index_folder: str = None, **kwargs):
        super().__init__(**kwargs)
        self.index_type = index_type
        self.index_folder = index_folder
        self.tokenizer = SimpleTokenizer()

        # Initialize index manager
        self.index_manager = IndexManager()

        # Get index path
        if index_folder:
            self.index_path = index_folder
        else:
            self.index_path = self.index_manager.get_index_path("bm25", index_type)

        # Initialize searcher
        self.searcher = self._initialize_searcher()
        self.rev_id_mapping = self._load_reverse_mapping()
    def _load_reverse_mapping(self):
        # prefer explicit reverse file; fall back to inverting id_mapping.json
        base = self.index_path
        rev_path = os.path.join(base, "id_mapping_rev.json")
        fwd_path = os.path.join(base, "id_mapping.json")
        m = {}
        if os.path.exists(rev_path):
            with open(rev_path, "r", encoding="utf-8") as f:
                m = json.load(f)                     # { "123": "orig_id" }
        elif os.path.exists(fwd_path):
            with open(fwd_path, "r", encoding="utf-8") as f:
                fwd = json.load(f)                   # { "orig_id": 123 }
            m = {str(v): k for k, v in fwd.items()}  # ensure string keys
        return m
    def _initialize_searcher(self) -> LuceneSearcher:
        """Initialize Lucene searcher."""
        if self.index_path.startswith("wikipedia-") or "prebuilt" in self.index_path:
            return LuceneSearcher.from_prebuilt_index(self.index_path)
        else:
            # Adjust path for local indexes
            if self.index_type == "wiki":
                actual_path = os.path.join(self.index_path, "bm25_index")
            else:
                actual_path = os.path.join(self.index_path, f"bm25_index_{self.index_type}")
            return LuceneSearcher(actual_path)

    def retrieve(self, documents: List[Document]) -> List[Document]:
        """Retrieve relevant contexts using BM25."""
        queries = [doc.question.question for doc in documents]
        qids = [str(i) for i in range(len(queries))]

        print(f"Retrieving {len(documents)} documents with BM25...")

        # Perform batch search
        batch_results = self._batch_search(queries, qids)

        # Process results
        for i, document in enumerate(tqdm(documents, desc="Processing documents")):
            contexts = []
            hits = batch_results.get(str(i), [])

            for hit in hits:
                try:
                    context = self._create_context_from_hit(hit, document)
                    if context:
                        contexts.append(context)
                except Exception as e:
                    print(f"Error processing document ID {hit.docid}: {e}")

            document.contexts = contexts

        return documents

    def _create_context_from_hit(self, hit, document) -> Context:
        """Create Context object from search hit."""
        lucene_doc = self.searcher.doc(hit.docid)
        raw_content = json.loads(lucene_doc.raw())

        contents = raw_content.get("contents", "")
        if "\n" in contents:
            title, text = contents.split("\n", 1)
        else:
            title, text = "No Title", contents

        # Map the (string) docid back to original ID
        doc_id = self.rev_id_mapping.get(str(hit.docid), str(hit.docid))

        return Context(
            id=doc_id,
            title=title,
            text=text,
            score=hit.score,
            has_answer=has_answers(text, document.answers.answers, self.tokenizer),
        )

    def _batch_search(self, queries: List[str], qids: List[str]) -> dict:
        """Perform batch search using Lucene searcher."""
        batch_results = {}

        for start in tqdm(range(0, len(queries), self.batch_size), desc="Batch search"):
            end = min(start + self.batch_size, len(queries))
            batch_queries = queries[start:end]
            batch_qids = qids[start:end]

            try:
                batch_hits = self.searcher.batch_search(
                    batch_queries, batch_qids, k=self.n_docs, threads=self.threads
                )
                batch_results.update(batch_hits)
            except Exception as e:
                print(f"Batch search failed for queries {batch_queries}: {e}")

        return batch_results

retrieve(documents)

Retrieve relevant contexts using BM25.

Source code in rankify/retrievers/bm25_retriever.py

def retrieve(self, documents: List[Document]) -> List[Document]:
    """Retrieve relevant contexts using BM25."""
    queries = [doc.question.question for doc in documents]
    qids = [str(i) for i in range(len(queries))]

    print(f"Retrieving {len(documents)} documents with BM25...")

    # Perform batch search
    batch_results = self._batch_search(queries, qids)

    # Process results
    for i, document in enumerate(tqdm(documents, desc="Processing documents")):
        contexts = []
        hits = batch_results.get(str(i), [])

        for hit in hits:
            try:
                context = self._create_context_from_hit(hit, document)
                if context:
                    contexts.append(context)
            except Exception as e:
                print(f"Error processing document ID {hit.docid}: {e}")

        document.contexts = contexts

    return documents

DenseRetriever

Bases: BaseRetriever

Dense retriever implementation supporting DPR, ANCE, and BPR models.

Uses FAISS indexes for efficient dense vector retrieval.

Source code in rankify/retrievers/dense_retriever.py

class DenseRetriever(BaseRetriever):
    """
    Dense retriever implementation supporting DPR, ANCE, and BPR models.

    Uses FAISS indexes for efficient dense vector retrieval.
    """

    MSMARCO_CORPUS_URL = "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/msmarco-passage-corpus/msmarco-passage-corpus.tsv?download=true"

    def __init__(self, method: str = "dpr-multi", index_type: str = "wiki", 
                 index_folder: str = None, **kwargs):
        super().__init__(**kwargs)
        self.method = method
        self.index_type = index_type
        self.index_folder = index_folder
        self.tokenizer = SimpleTokenizer()

        # Initialize index manager
        self.index_manager = IndexManager()

        # Handle custom index folder case
        if index_folder:
            # For custom indexes, get encoder from method and default index_type
            if method in self.index_manager.index_configs:
                # Use wiki as default if the provided index_type is a path
                default_index_type = "wiki" if (index_type.startswith("./") or 
                                               index_type.startswith("/") or 
                                               os.path.exists(index_type)) else index_type
                if default_index_type in self.index_manager.index_configs[method]:
                    self.query_encoder = self.index_manager.index_configs[method][default_index_type]["encoder"]
                else:
                    # Fallback to first available encoder for this method
                    available_types = list(self.index_manager.index_configs[method].keys())
                    self.query_encoder = self.index_manager.index_configs[method][available_types[0]]["encoder"]
                    print(f"Warning: Using default encoder {self.query_encoder} for custom index")
            else:
                raise ValueError(f"Unsupported method: {method}")
        else:
            # Get query encoder for standard indexes
            if method not in self.index_manager.index_configs:
                raise ValueError(f"Unsupported method: {method}")
            if index_type not in self.index_manager.index_configs[method]:
                raise ValueError(f"Unsupported index_type '{index_type}' for method '{method}'")
            self.query_encoder = self.index_manager.index_configs[method][index_type]["encoder"]

        self.question_tokenizer = AutoTokenizer.from_pretrained(self.query_encoder)

        # Initialize variables
        self.id_mapping = None
        self.corpus = {}

        # Handle custom vs prebuilt indexes
        if index_folder:
            self.index_path = index_folder
            self.id_mapping = self.index_manager.load_id_mapping(index_folder)
            self.corpus = self.index_manager.load_corpus(index_folder)
            print(f"Loaded corpus with {len(self.corpus)} documents")
            if self.id_mapping:
                print(f"Loaded ID mapping with {len(self.id_mapping)} entries")
        else:
            self.index_path = self.index_manager.get_index_path(method, index_type)

        # Initialize searcher
        self.searcher = self._initialize_searcher()
        if self.searcher is None:
            raise RuntimeError("Failed to initialize searcher")

        # Load corpus for MSMARCO if needed
        if index_type == "msmarco" and not index_folder:
            self._load_msmarco_corpus()

    def _initialize_searcher(self):
        """
        Initialize FAISS searcher for all cases.
        This is PURE DPR - no LuceneSearcher anywhere!
        """
        try:
            # Case 1: Custom index folder provided
            if self.index_folder:
                print(f"Initializing FaissSearcher with custom index folder: {self.index_folder}")
                if not os.path.exists(self.index_folder):
                    raise ValueError(f"Custom index folder '{self.index_folder}' does not exist.")
                return FaissSearcher(self.index_folder, self.query_encoder)

            # Case 2: Standard indexes (prebuilt or downloaded)
            print(f"Index path: {self.index_path}")

            # Case 2a: URL - need to download first
            if self.index_path.startswith("http"):
                print(f"Downloading index from URL: {self.index_path}")
                local_dir = self.index_manager.download_and_extract_index(self.index_path)
                print(f"Downloaded index to: {local_dir}")
                return FaissSearcher(local_dir, self.query_encoder)

            # Case 2b: Local file path (already downloaded)
            elif (os.path.exists(self.index_path) or 
                  self.index_path.startswith('/') or 
                  self.index_path.startswith('./')):
                print(f"Using local index path: {self.index_path}")
                if not os.path.exists(self.index_path):
                    raise ValueError(f"Local index path '{self.index_path}' does not exist.")
                return FaissSearcher(self.index_path, self.query_encoder)

            # Case 2c: Prebuilt index name (like "wikipedia-dpr-100w.dpr-multi")
            else:
                print(f"Using prebuilt index: {self.index_path}")
                return FaissSearcher.from_prebuilt_index(self.index_path, self.query_encoder)

        except Exception as e:
            print(f"ERROR initializing searcher:")
            print(f"  - Method: {self.method}")
            print(f"  - Index type: {self.index_type}")
            print(f"  - Index folder: {self.index_folder}")
            print(f"  - Index path: {self.index_path}")
            print(f"  - Query encoder: {self.query_encoder}")
            print(f"  - Error: {e}")
            import traceback
            traceback.print_exc()
            return None

    def _load_msmarco_corpus(self):
        """Load MSMARCO corpus into memory."""
        corpus_file = os.path.join(self.index_manager.cache_dir, "msmarco-passage-corpus.tsv")

        if not os.path.exists(corpus_file):
            print("Downloading MSMARCO corpus...")
            self._download_file(self.MSMARCO_CORPUS_URL, corpus_file)

        print("Loading MSMARCO corpus...")
        self.corpus = {}
        with open(corpus_file, "r", encoding="utf-8") as f:
            next(f)  # Skip header
            for line in f:
                parts = line.strip().split("\t")
                if len(parts) >= 3:
                    doc_id, text, title = parts[0], parts[1], parts[2]
                    self.corpus[doc_id] = {"text": text, "title": title}

    def _download_file(self, url: str, save_path: str):
        """Download a file from URL."""
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        print(f"Downloading {os.path.basename(save_path)}...")
        response = requests.get(url, stream=True)
        response.raise_for_status()

        with open(save_path, "wb") as f:
            for chunk in tqdm(response.iter_content(chunk_size=1024), 
                            desc=f"Downloading {os.path.basename(save_path)}"):
                f.write(chunk)

    def retrieve(self, documents: List[Document]) -> List[Document]:
        """Retrieve relevant contexts using dense retrieval."""
        queries = [doc.question.question for doc in documents]
        qids = [str(i) for i in range(len(queries))]

        print(f"Retrieving {len(documents)} documents with {self.method}...")

        # Perform batch search
        batch_results = self._batch_search(queries, qids)

        # Process results
        for i, document in enumerate(tqdm(documents, desc="Processing documents")):
            contexts = []
            hits = batch_results.get(str(i), [])

            for hit in hits:
                try:
                    context = self._create_context_from_hit(hit, document)
                    if context:
                        contexts.append(context)
                except Exception as e:
                    print(f"Error processing document ID {hit.docid}: {e}")

            document.contexts = contexts
            #break

        return documents

    def _create_context_from_hit(self, hit, document) -> Context:
        """Create Context object from search hit."""
        doc_id = hit.docid
        text = ""
        title = ""

        try:
            # Strategy 1: Use loaded corpus (MSMARCO or custom)
            if self.corpus:
                if self.index_type == "msmarco":
                    # MSMARCO corpus format
                    doc_data = self.corpus.get(str(hit.docid), {})
                    if doc_data:
                        text = doc_data.get("text", "")
                        title = doc_data.get("title", "")
                    else:
                        text = f"Document {hit.docid} not found in MSMARCO corpus"
                        title = f"Document {hit.docid}"
                else:
                    # Custom corpus format
                    doc_data = self.corpus.get(str(hit.docid), {})
                    if doc_data:
                        text = doc_data.get("contents") or doc_data.get("text", "") or ""
                        title = doc_data.get("title", "")
                        # Use ID mapping if available
                        if self.id_mapping:
                            doc_id = self.id_mapping.get(int(hit.docid), hit.docid)
                    else:
                        text = f"Document {hit.docid} not found in custom corpus"
                        title = f"Document {hit.docid}"

            # Strategy 2: Try to get document from searcher
            elif hasattr(self.searcher, 'doc'):
                try:
                    doc = self.searcher.doc(hit.docid)
                    raw_content = json.loads(doc.raw())
                    content = raw_content.get("contents", "")

                    if '\n' in content:
                        lines = content.split('\n', 1)
                        title = lines[0].strip()
                        text = lines[1].strip() if len(lines) > 1 else ""
                    else:
                        title = "No Title"
                        text = content

                except Exception as e:
                    print(f"Error retrieving document {hit.docid} from searcher: {e}")
                    text = f"Error retrieving document {hit.docid}"
                    title = f"Document {hit.docid}"

            # Strategy 3: Fallback
            else:
                print(f"Warning: No way to retrieve content for document {hit.docid}")
                text = f"Document {hit.docid}"
                title = f"Document {hit.docid}"

        except Exception as e:
            print(f"Error processing hit {hit.docid}: {e}")
            text = f"Error: {str(e)}"
            title = f"Document {hit.docid}"

        return Context(
            id=doc_id,
            title=title,
            text=text,
            score=hit.score,
            has_answer=has_answers(text, document.answers.answers, self.tokenizer)
        )

    def _batch_search(self, queries: List[str], qids: List[str]) -> dict:
        """Perform batch search using FAISS searcher."""
        if self.searcher is None:
            raise RuntimeError("Searcher not initialized properly")

        if not hasattr(self.searcher, 'batch_search'):
            raise AttributeError(f"Searcher {type(self.searcher)} does not have batch_search method")

        batch_results = {}
        batch_qids, batch_queries = [], []

        for idx, (qid, query) in enumerate(tqdm(zip(qids, queries), desc="Batch search", total=len(qids))):
            query = self._preprocess_query(query)
            batch_qids.append(qid)
            batch_queries.append(query)

            if (idx + 1) % self.batch_size == 0 or idx == len(qids) - 1:
                try:
                    results = self.searcher.batch_search(batch_queries, batch_qids, 
                                                       self.n_docs, self.threads)
                    batch_results.update(results)
                except Exception as e:
                    print(f"Batch search failed:")
                    print(f"  - Searcher type: {type(self.searcher)}")
                    print(f"  - Queries: {batch_queries}")
                    print(f"  - Error: {e}")
                    raise e

                batch_qids.clear()
                batch_queries.clear()

        return batch_results

    def _preprocess_query(self, query: str) -> str:
        """Preprocess query to handle token length limits."""
        try:
            tokenized = self.question_tokenizer(query, add_special_tokens=True)

            if len(tokenized["input_ids"]) > 511:
                truncated = self.question_tokenizer.decode(
                    tokenized["input_ids"][:511], skip_special_tokens=True
                )
                return truncated
            return query
        except Exception as e:
            print(f"Error preprocessing query '{query}': {e}")
            return query

retrieve(documents)

Retrieve relevant contexts using dense retrieval.

Source code in rankify/retrievers/dense_retriever.py

def retrieve(self, documents: List[Document]) -> List[Document]:
    """Retrieve relevant contexts using dense retrieval."""
    queries = [doc.question.question for doc in documents]
    qids = [str(i) for i in range(len(queries))]

    print(f"Retrieving {len(documents)} documents with {self.method}...")

    # Perform batch search
    batch_results = self._batch_search(queries, qids)

    # Process results
    for i, document in enumerate(tqdm(documents, desc="Processing documents")):
        contexts = []
        hits = batch_results.get(str(i), [])

        for hit in hits:
            try:
                context = self._create_context_from_hit(hit, document)
                if context:
                    contexts.append(context)
            except Exception as e:
                print(f"Error processing document ID {hit.docid}: {e}")

        document.contexts = contexts
        #break

    return documents

ANCERetriever

Bases: BaseRetriever

ANCE retriever implementation following the exact same pattern as DenseRetriever.

Supports both prebuilt and custom indices using FaissSearcher like DPR.

Parameters:

Name	Type	Description	Default
index_type	str	Type of prebuilt index ("wiki", "msmarco").	'wiki'
index_folder	str	Path to custom ANCE index folder.	None
encoder_name	str	ANCE model name (auto-detected if not provided).	None
method	str	Method variant ("ance", "ance-multi", "ance-msmarco").	'ance-multi'
n_docs	int	Number of documents to retrieve per query.	required
batch_size	int	Number of queries to process in a batch.	required
threads	int	Number of parallel threads for processing.	required
device	str	Device to use for encoding ("cpu" or "cuda").	'cuda'

Source code in rankify/retrievers/ance_retriever.py

class ANCERetriever(BaseRetriever):
    """
    ANCE retriever implementation following the exact same pattern as DenseRetriever.

    Supports both prebuilt and custom indices using FaissSearcher like DPR.

    Args:
        index_type (str): Type of prebuilt index ("wiki", "msmarco"). 
        index_folder (str, optional): Path to custom ANCE index folder.
        encoder_name (str, optional): ANCE model name (auto-detected if not provided).
        method (str): Method variant ("ance", "ance-multi", "ance-msmarco").
        n_docs (int): Number of documents to retrieve per query.
        batch_size (int): Number of queries to process in a batch.
        threads (int): Number of parallel threads for processing.
        device (str): Device to use for encoding ("cpu" or "cuda").
    """

    def __init__(self, index_type: str = "wiki", index_folder: str = None,
                encoder_name: str = None, device: str = "cuda", method: str = "ance-multi", **kwargs):
        super().__init__(**kwargs)

        self.method = method
        self.index_type = index_type
        self.index_folder = index_folder
        self.device = device if torch.cuda.is_available() else "cpu"
        self.tokenizer_simple = SimpleTokenizer()
        self.batch_size = 128

        # Initialize index manager (following DenseRetriever pattern exactly)
        self.index_manager = IndexManager()

        # Initialize query_encoder to None first
        self.query_encoder = None

        # Get query encoder (following DenseRetriever logic exactly)
        try:
            if index_folder:
                # For custom indexes, get encoder from method and default index_type
                if method in self.index_manager.index_configs:
                    # Use wiki as default if the provided index_type is a path
                    default_index_type = "wiki" if (index_type.startswith("./") or 
                                                index_type.startswith("/") or 
                                                os.path.exists(index_type)) else index_type

                    if default_index_type in self.index_manager.index_configs[method]:
                        self.query_encoder = self.index_manager.index_configs[method][default_index_type]["encoder"]
                    else:
                        # Fallback to first available encoder for this method
                        available_types = list(self.index_manager.index_configs[method].keys())
                        if available_types:
                            self.query_encoder = self.index_manager.index_configs[method][available_types[0]]["encoder"]
                            #print(f"Warning: Using default encoder {self.query_encoder} for custom index")
                        else:
                            raise ValueError(f"No encoders available for method: {method}")
                else:
                    # Fallback: use encoder_name if provided, otherwise use default ANCE encoder
                    if encoder_name:
                        self.query_encoder = encoder_name
                    else:
                        self.query_encoder = "castorini/ance-dpr-question-multi"  # Default ANCE encoder
                    #print(f"Warning: Method '{method}' not found in index configs, using encoder: {self.query_encoder}")
            else:
                # Get query encoder for standard indexes
                if method not in self.index_manager.index_configs:
                    # Fallback: use encoder_name if provided, otherwise use default ANCE encoder
                    if encoder_name:
                        self.query_encoder = encoder_name
                    else:
                        self.query_encoder = "castorini/ance-dpr-context-multi"  # Default ANCE encoder
                    #print(f"Warning: Method '{method}' not found in index configs, using encoder: {self.query_encoder}")
                elif index_type not in self.index_manager.index_configs[method]:
                    # Try to get any available encoder for this method
                    available_types = list(self.index_manager.index_configs[method].keys())
                    if available_types:
                        self.query_encoder = self.index_manager.index_configs[method][available_types[0]]["encoder"]
                        #print(f"Warning: Index type '{index_type}' not found for method '{method}', using encoder: {self.query_encoder}")
                    else:
                        # Final fallback
                        if encoder_name:
                            self.query_encoder = encoder_name
                        else:
                            self.query_encoder = "castorini/ance-dpr-question-multi"
                        #print(f"Warning: No configs found, using encoder: {self.query_encoder}")
                else:
                    self.query_encoder = self.index_manager.index_configs[method][index_type]["encoder"]

        except Exception as e:
            #print(f"Error getting query encoder: {e}")
            # Final fallback
            if encoder_name:
                self.query_encoder = encoder_name
            else:
                self.query_encoder = "castorini/ance-dpr-question-multi"  # Default ANCE encoder
            #print(f"Using fallback encoder: {self.query_encoder}")

        # Ensure query_encoder is set
        if self.query_encoder is None:
            if encoder_name:
                self.query_encoder = encoder_name
            else:
                self.query_encoder = "castorini/ance-dpr-question-multi"  # Default ANCE encoder
            #print(f"Final fallback: Using encoder {self.query_encoder}")

        # Initialize variables (following DenseRetriever pattern exactly)
        self.id_mapping = None
        self.docid_to_faiss = None  # Add reverse mapping
        self.corpus = {}

        # Handle custom vs prebuilt indexes (following DenseRetriever pattern exactly)
        if index_folder:
            self.index_path = index_folder

            # FIXED: Load both ID mappings
            faiss_to_docid_file = os.path.join(index_folder, "faiss_to_docid.json")
            if os.path.exists(faiss_to_docid_file):
                with open(faiss_to_docid_file, 'r', encoding='utf-8') as f:
                    self.id_mapping = json.load(f)
                #print(f"Loaded faiss_to_docid mapping with {len(self.id_mapping)} entries")
            else:
                #print(f"Warning: faiss_to_docid mapping file not found: {faiss_to_docid_file}")
                self.id_mapping = None

            # Load reverse mapping (docid_to_faiss)
            docid_to_faiss_file = os.path.join(index_folder, "docid_to_faiss.json")
            if os.path.exists(docid_to_faiss_file):
                with open(docid_to_faiss_file, 'r', encoding='utf-8') as f:
                    self.docid_to_faiss = json.load(f)
                #print(f"Loaded docid_to_faiss mapping with {len(self.docid_to_faiss)} entries")
            else:
                #print(f"Warning: docid_to_faiss mapping file not found: {docid_to_faiss_file}")
                self.docid_to_faiss = None

            # Load corpus metadata
            corpus_file = os.path.join(index_folder, "corpus_metadata.json")
            if os.path.exists(corpus_file):
                with open(corpus_file, 'r', encoding='utf-8') as f:
                    self.corpus = json.load(f)
                print(f"Loaded corpus with {len(self.corpus)} documents")
            else:
                print(f"Warning: Corpus metadata file not found: {corpus_file}")
                self.corpus = {}
        else:
            # FIXED: Use the actual method name (ance-multi) to get the correct prebuilt index
            try:
                self.index_path = self.index_manager.get_index_path(self.method, index_type)
            except Exception as e:
                #print(f"Error getting index path: {e}")
                # Fallback to a reasonable default
                self.index_path = f"{method}-{index_type}"

            #print(f"🔍 Method '{self.method}' + Index type '{index_type}' -> Index path: '{self.index_path}'")

        # Initialize searcher (following DenseRetriever pattern exactly)
        try:
            self.searcher = self._initialize_searcher()
            if self.searcher is None:
                raise RuntimeError("Failed to initialize searcher")
        except Exception as e:
            #print(f"Error initializing searcher: {e}")
            raise RuntimeError(f"Failed to initialize searcher: {e}")

        print(f"✅ ANCE Retriever initialized")
        print(f"  📚 Index: {self.index_path}")
        print(f"  🏷️ Type: {index_type}")
        print(f"  🔧 Method: {self.method}")
        print(f"  🤖 Model: {self.query_encoder}")
        print(f"  📊 Corpus: {len(self.corpus)} documents")
        if self.id_mapping:
            print(f"  🗂️ Mappings: {len(self.id_mapping)} entries")
        print(f"  🖥️ Device: {self.device}")

    def _initialize_searcher(self):
        """
        Initialize FAISS searcher following DenseRetriever pattern exactly.
        This uses FaissSearcher for both prebuilt and custom indices.
        """
        try:
            # Case 1: Custom index folder provided
            if self.index_folder:
                print(f"Initializing FaissSearcher with custom index folder: {self.index_folder}")
                if not os.path.exists(self.index_folder):
                    raise ValueError(f"Custom index folder '{self.index_folder}' does not exist.")
                return FaissSearcher(self.index_folder, self.query_encoder)

            # Case 2: Standard indexes (prebuilt or downloaded)
            #print(f"Index path: {self.index_path}")

            # Case 2a: URL - need to download first
            if self.index_path.startswith("http"):
                #print(f"Downloading index from URL: {self.index_path}")
                local_dir = self.index_manager.download_and_extract_index(self.index_path)
                #print(f"Downloaded index to: {local_dir}")
                return FaissSearcher(local_dir, self.query_encoder)

            # Case 2b: Local file path (already downloaded)
            elif (os.path.exists(self.index_path) or 
                  self.index_path.startswith('/') or 
                  self.index_path.startswith('./')):
                #print(f"Using local index path: {self.index_path}")
                if not os.path.exists(self.index_path):
                    raise ValueError(f"Local index path '{self.index_path}' does not exist.")
                return FaissSearcher(self.index_path, self.query_encoder)

            # Case 2c: Prebuilt index name (like "wikipedia-dpr-100w.ance-multi")
            else:

                print(f"Using prebuilt index: {self.index_path}")
                # IMPORTANT: For ance-multi, this should be a prebuilt identifier like "wikipedia-dpr-100w.ance-multi"
                if self.method == "ance-multi" and self.index_path.startswith("/"):
                    raise ValueError(f"ERROR: Method 'ance-multi' should use prebuilt index, but got local path: {self.index_path}. "
                                   f"Check IndexManager configuration for method='{self.method}', index_type='{self.index_type}'")
                return FaissSearcher.from_prebuilt_index(self.index_path, self.query_encoder)

        except Exception as e:
            print(f"ERROR initializing searcher:")
            print(f"  - Method: {self.method}")
            print(f"  - Index type: {self.index_type}")
            print(f"  - Index folder: {self.index_folder}")
            print(f"  - Index path: {self.index_path}")
            print(f"  - Query encoder: {self.query_encoder}")
            print(f"  - Error: {e}")
            import traceback
            traceback.print_exc()
            return None

    def retrieve(self, documents: List[Document]) -> List[Document]:
        """Retrieve relevant contexts using ANCE (following DenseRetriever pattern exactly)."""
        queries = [doc.question.question for doc in documents]
        qids = [str(i) for i in range(len(queries))]

        print(f"🔍 Retrieving {len(documents)} documents with {self.method}...")

        # Perform batch search (following DenseRetriever pattern exactly)
        batch_results = self._batch_search(queries, qids)

        # Process results (following DenseRetriever pattern exactly)
        for i, document in enumerate(tqdm(documents, desc="Processing documents")):
            contexts = []
            hits = batch_results.get(str(i), [])

            for hit in hits:
                try:
                    context = self._create_context_from_hit(hit, document)
                    if context:
                        contexts.append(context)
                except Exception as e:
                    print(f"Error processing document ID {hit.docid}: {e}")

            document.contexts = contexts
            #break

        print(f"✅ Retrieved contexts for {len(documents)} documents")
        return documents

    def _batch_search(self, queries: List[str], qids: List[str]) -> dict:
        """Perform batch search using FAISS searcher (following DenseRetriever pattern exactly)."""
        if self.searcher is None:
            raise RuntimeError("Searcher not initialized properly")

        if not hasattr(self.searcher, 'batch_search'):
            raise AttributeError(f"Searcher {type(self.searcher)} does not have batch_search method")

        batch_results = {}
        batch_qids, batch_queries = [], []

        for idx, (qid, query) in enumerate(tqdm(zip(qids, queries), desc="Batch search", total=len(qids))):
            query = self._preprocess_query(query)
            batch_qids.append(qid)
            batch_queries.append(query)

            if (idx + 1) % self.batch_size == 0 or idx == len(qids) - 1:
                try:
                    results = self.searcher.batch_search(batch_queries, batch_qids, 
                                                       self.n_docs, self.threads)
                    batch_results.update(results)
                except Exception as e:
                    print(f"Batch search failed:")
                    print(f"  - Searcher type: {type(self.searcher)}")
                    print(f"  - Queries: {batch_queries}")
                    print(f"  - Error: {e}")
                    raise e

                batch_qids.clear()
                batch_queries.clear()
            #break
        return batch_results

    def _create_context_from_hit(self, hit, document: Document) -> Context:
        """Create Context object from search hit with correct ID handling."""

        # FIXED: Determine if hit.docid is a FAISS ID or original document ID
        hit_docid = str(hit.docid)

        # Strategy: Try to determine the type of ID we received
        is_faiss_id = False
        original_doc_id = hit_docid
        faiss_id = None

        # Check if hit_docid looks like a FAISS sequential ID (numeric)
        try:
            # If it's numeric, it might be a FAISS ID
            potential_faiss_id = int(hit_docid)

            # If we have an ID mapping and this looks like a FAISS ID
            if self.id_mapping and hit_docid in self.id_mapping:
                # This is definitely a FAISS ID, convert to original
                original_doc_id = self.id_mapping[hit_docid]
                faiss_id = potential_faiss_id
                is_faiss_id = True
                #print(f"DEBUG: FAISS ID {hit_docid} -> Original ID {original_doc_id}")
            else:
                # This might be a numeric original document ID
                original_doc_id = hit_docid
                # Try to find the corresponding FAISS ID
                if self.docid_to_faiss and hit_docid in self.docid_to_faiss:
                    faiss_id = int(self.docid_to_faiss[hit_docid])
                #print(f"DEBUG: Numeric original ID {hit_docid}")

        except ValueError:
            # hit_docid is not numeric, so it's definitely an original document ID
            original_doc_id = hit_docid
            # Try to find the corresponding FAISS ID
            if self.docid_to_faiss and hit_docid in self.docid_to_faiss:
                faiss_id = int(self.docid_to_faiss[hit_docid])
            #print(f"DEBUG: Non-numeric original ID {hit_docid}")

        text = ""
        title = ""

        try:
            # Strategy 1: Use loaded corpus (custom)
            if self.corpus:
                # Look up document content using the original document ID
                doc_data = self.corpus.get(original_doc_id, {})
                if doc_data:
                    text = doc_data.get("contents", "")
                    title = doc_data.get("title", "")
                else:
                    text = f"Document {original_doc_id} not found in custom corpus"
                    title = f"Document {original_doc_id}"

            # Strategy 2: Try to get document from searcher (prebuilt)
            elif hasattr(self.searcher, 'doc'):
                try:
                    # We need a FAISS ID to get document from searcher
                    doc_id_for_searcher = faiss_id if faiss_id is not None else hit_docid

                    # Use the appropriate ID to get document from searcher
                    doc = self.searcher.doc(doc_id_for_searcher)
                    raw_content = json.loads(doc.raw())
                    content = raw_content.get("contents", "")

                    if '\n' in content:
                        lines = content.split('\n', 1)
                        title = lines[0].strip()
                        text = lines[1].strip() if len(lines) > 1 else ""
                    else:
                        title = "No Title"
                        text = content

                except Exception as e:
                    print(f"Error retrieving document {doc_id_for_searcher} from searcher: {e}")
                    text = f"Error retrieving document {original_doc_id}"
                    title = f"Document {original_doc_id}"

            # Strategy 3: Fallback
            else:
                print(f"Warning: No way to retrieve content for document {original_doc_id}")
                text = f"Document {original_doc_id}"
                title = f"Document {original_doc_id}"

        except Exception as e:
            print(f"Error processing hit {hit_docid}: {e}")
            text = f"Error: {str(e)}"
            title = f"Document {original_doc_id}"

        # Check if document contains answer (following DenseRetriever pattern exactly)
        has_answer = has_answers(text, document.answers.answers, self.tokenizer_simple)

        return Context(
            id=original_doc_id,  # Return the original document ID
            title=title,
            text=text,
            score=hit.score,
            has_answer=has_answer
        )

    def _preprocess_query(self, query: str) -> str:
        """Preprocess query to handle token length limits (following DenseRetriever pattern)."""
        # For ANCE, we don't need to load the tokenizer here since FaissSearcher handles encoding
        # Just return the query as-is, following the same pattern as DenseRetriever
        return query.strip()

retrieve(documents)

Retrieve relevant contexts using ANCE (following DenseRetriever pattern exactly).

Source code in rankify/retrievers/ance_retriever.py

def retrieve(self, documents: List[Document]) -> List[Document]:
    """Retrieve relevant contexts using ANCE (following DenseRetriever pattern exactly)."""
    queries = [doc.question.question for doc in documents]
    qids = [str(i) for i in range(len(queries))]

    print(f"🔍 Retrieving {len(documents)} documents with {self.method}...")

    # Perform batch search (following DenseRetriever pattern exactly)
    batch_results = self._batch_search(queries, qids)

    # Process results (following DenseRetriever pattern exactly)
    for i, document in enumerate(tqdm(documents, desc="Processing documents")):
        contexts = []
        hits = batch_results.get(str(i), [])

        for hit in hits:
            try:
                context = self._create_context_from_hit(hit, document)
                if context:
                    contexts.append(context)
            except Exception as e:
                print(f"Error processing document ID {hit.docid}: {e}")

        document.contexts = contexts
        #break

    print(f"✅ Retrieved contexts for {len(documents)} documents")
    return documents

BGERetriever

Bases: BaseRetriever

BGE retriever implementation using precomputed embeddings and FAISS indexing.

Implements BGE (Beijing Academy of Artificial Intelligence General Embedding) model for dense passage retrieval with efficient FAISS-based search.

Source code in rankify/retrievers/bge_retriever.py

class BGERetriever(BaseRetriever):
    """
    BGE retriever implementation using precomputed embeddings and FAISS indexing.

    Implements BGE (Beijing Academy of Artificial Intelligence General Embedding) model
    for dense passage retrieval with efficient FAISS-based search.
    """

    def __init__(self, model: str = "BAAI/bge-large-en-v1.5", index_type: str = "wiki", 
                 index_folder: str = None, device: str = "cuda", **kwargs):
        super().__init__(**kwargs)
        self.model_name = model
        self.index_type = index_type
        self.index_folder = index_folder
        self.device = device
        self.tokenizer_simple = SimpleTokenizer()

        # Initialize index manager
        self.index_manager = IndexManager()

        # Setup paths and download if needed
        if index_folder:
            self.index_path = index_folder
            self.passage_path = os.path.join(self.index_folder, 'passages.tsv')
        else:
            self.index_path = os.path.join(self.index_manager.cache_dir, "index", f"bge_index_{index_type}")
            self._ensure_index_and_passages_downloaded()
            self.passage_path = self._get_passage_path()

        # Load components
        self.doc_ids = self._load_document_ids()
        self.doc_texts = self._load_tsv()
        self.index = self._initialize_searcher()

        # Load model and tokenizer
        self.model_hf = AutoModel.from_pretrained(self.model_name).to(self.device).eval()
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)

    def _initialize_searcher(self):
        """Initialize FAISS index for BGE retrieval."""
        return self._build_faiss_index()

    def _ensure_index_and_passages_downloaded(self):
        """Download and extract BGE index and passages if needed."""
        os.makedirs(self.index_path, exist_ok=True)

        if self.index_type not in self.index_manager.index_configs.get("bge", {}):
            raise ValueError(f"Unsupported BGE index type: {self.index_type}")

        config = self.index_manager.index_configs["bge"][self.index_type]

        # Check if required files exist
        required_files = [
            os.path.join(self.index_path, "bge_doc_ids.pkl"),
            os.path.join(self.index_path, "bge_embeddings.h5")
        ]

        if not all(os.path.exists(f) for f in required_files):
            urls = config.get("urls")
            if isinstance(urls, list):
                # Multi-part download
                for url in urls:
                    filename = self._extract_filename_from_url(url)
                    local_path = os.path.join(self.index_path, filename)
                    if not os.path.exists(local_path):
                        self._download_file(url, local_path)
                self._extract_multi_part_archive()
            else:
                # Single file download
                zip_path = os.path.join(self.index_path, "index.zip")
                if not os.path.exists(zip_path):
                    self._download_file(urls, zip_path)
                self._extract_zip_files()

        # Download passages if needed
        passages_url = config.get("passages_url")
        if passages_url:
            passage_filename = self._extract_filename_from_url(passages_url)
            passage_path = os.path.join(self.index_manager.cache_dir, passage_filename)
            if not os.path.exists(passage_path):
                self._download_file(passages_url, passage_path)

    def _get_passage_path(self):
        """Get path to passages file."""
        if self.index_type not in self.index_manager.index_configs.get("bge", {}):
            raise ValueError(f"Unsupported BGE index type: {self.index_type}")

        config = self.index_manager.index_configs["bge"][self.index_type]
        passages_url = config.get("passages_url")
        if passages_url:
            passage_filename = self._extract_filename_from_url(passages_url)
            return os.path.join(self.index_manager.cache_dir, passage_filename)
        return None

    def _extract_filename_from_url(self, url):
        """Extract filename from URL."""
        parsed_url = urlparse(url)
        return os.path.basename(parsed_url.path).split('?')[0]

    def _download_file(self, url: str, save_path: str):
        """Download file from URL with progress bar."""
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        response = requests.get(url, stream=True)
        response.raise_for_status()

        with open(save_path, "wb") as f:
            for chunk in tqdm(response.iter_content(chunk_size=1024), 
                            desc=f"Downloading {os.path.basename(save_path)}"):
                f.write(chunk)

    def _extract_zip_files(self):
        """Extract ZIP files in the index folder."""
        zip_files = [f for f in os.listdir(self.index_path) if f.endswith(".zip")]

        for zip_file in zip_files:
            zip_path = os.path.join(self.index_path, zip_file)

            with zipfile.ZipFile(zip_path, "r") as zip_ref:
                for member in zip_ref.namelist():
                    filename = os.path.basename(member)
                    if not filename:
                        continue

                    extracted_path = os.path.join(self.index_path, filename)
                    with zip_ref.open(member) as source, open(extracted_path, "wb") as target:
                        shutil.copyfileobj(source, target)

            print(f"Extracted {zip_file}")
            os.remove(zip_path)

    def _extract_multi_part_archive(self):
        """Extract multi-part tar.gz archives."""
        parts = sorted([f for f in os.listdir(self.index_path) if f.startswith("bgb_index.tar.")],
                      key=lambda x: x.split('.')[-1])

        if not parts:
            return

        combined_path = os.path.join(self.index_path, "combined.tar.gz")

        # Combine parts
        with open(combined_path, "wb") as combined:
            for part in parts:
                with open(os.path.join(self.index_path, part), "rb") as part_file:
                    shutil.copyfileobj(part_file, combined)

        try:
            # Decompress and extract
            tar_file = combined_path.replace(".gz", "")
            with gzip.open(combined_path, "rb") as f_in, open(tar_file, "wb") as f_out:
                shutil.copyfileobj(f_in, f_out)

            shutil.unpack_archive(tar_file, self.index_path)
            os.remove(tar_file)

        except Exception as e:
            raise RuntimeError(f"Error extracting multi-part archive: {e}")
        finally:
            # Cleanup
            if os.path.exists(combined_path):
                os.remove(combined_path)
            for part in parts:
                part_path = os.path.join(self.index_path, part)
                if os.path.exists(part_path):
                    os.remove(part_path)

    def _build_faiss_index(self):
        """Build or load FAISS index."""
        print("Handling FAISS index...")

        index_path = os.path.join(self.index_path, "faiss_index.bin")
        embeddings_path = os.path.join(self.index_path, "bge_embeddings.h5")

        if os.path.exists(index_path):
            print(f"Loading existing FAISS index from {index_path}...")
            index = faiss.read_index(index_path)
        else:
            print(f"Building FAISS index from embeddings at {embeddings_path}...")
            with h5py.File(embeddings_path, "r") as f:
                embeddings = f["embeddings"][:].astype(np.float32)

            embedding_dim = embeddings.shape[1]
            index = faiss.IndexFlatIP(embedding_dim)

            # Add embeddings in chunks
            chunk_size = 50000
            for start in tqdm(range(0, embeddings.shape[0], chunk_size), 
                            desc="Adding embeddings to FAISS index"):
                end = min(start + chunk_size, embeddings.shape[0])
                chunk = embeddings[start:end]
                index.add(chunk)

            print(f"Saving FAISS index to {index_path}...")
            faiss.write_index(index, index_path)

        print(f"FAISS index loaded with {index.ntotal} embeddings.")
        return index

    def _load_document_ids(self):
        """Load document IDs from pickle file."""
        doc_ids_path = os.path.join(self.index_path, "bge_doc_ids.pkl")
        print(f"Loading document IDs from {doc_ids_path}...")

        doc_ids = []
        with open(doc_ids_path, "rb") as f:
            while True:
                try:
                    doc_ids.extend(pickle.load(f))
                except EOFError:
                    break

        print(f"Loaded {len(doc_ids)} document IDs.")
        return doc_ids

    def _load_tsv(self):
        """Load document texts from TSV file."""
        if not self.passage_path or not os.path.exists(self.passage_path):
            print("Warning: Passage file not found, using empty corpus")
            return {}

        doc_texts = {}
        with open(self.passage_path, "r", encoding="utf-8") as f:
            next(f)  # Skip header
            for line in f:
                try:
                    doc_id, passage, title = line.strip().split("\t")
                    doc_texts[doc_id] = {"text": passage, "title": title}
                except ValueError:
                    continue  # Skip malformed lines

        print(f"Loaded {len(doc_texts)} passages.")
        return doc_texts

    def _encode_queries(self, queries: List[str]):
        """Encode queries into dense embeddings."""
        all_embeddings = []
        with torch.no_grad():
            for i in tqdm(range(0, len(queries), self.batch_size), desc="Encoding queries"):
                batch = queries[i:i + self.batch_size]
                tokenized = self.tokenizer(batch, padding=True, truncation=True, 
                                         return_tensors="pt").to(self.device)
                model_output = self.model_hf(**tokenized)
                embeddings = model_output.last_hidden_state[:, 0, :]  # CLS token
                embeddings = torch.nn.functional.normalize(embeddings, p=2, dim=1)
                all_embeddings.append(embeddings.cpu().numpy())

        return np.vstack(all_embeddings)

    def retrieve(self, documents: List[Document]) -> List[Document]:
        """Retrieve relevant contexts using BGE."""
        queries = [doc.question.question for doc in documents]
        print(f"Retrieving {len(documents)} documents with BGE...")

        # Encode queries
        query_embeddings = self._encode_queries(queries)

        # Check dimension compatibility
        if query_embeddings.shape[1] != self.index.d:
            raise ValueError(f"Dimension mismatch: queries={query_embeddings.shape[1]}, "
                           f"index={self.index.d}")

        # Batch FAISS search
        all_distances = []
        all_indices = []

        for start_idx in tqdm(range(0, len(query_embeddings), self.batch_size), 
                            desc="FAISS search"):
            end_idx = min(start_idx + self.batch_size, len(query_embeddings))
            batch_embeddings = query_embeddings[start_idx:end_idx]

            distances, indices = self.index.search(batch_embeddings, self.n_docs)
            all_distances.append(distances)
            all_indices.append(indices)

        # Combine results
        all_distances = np.vstack(all_distances)
        all_indices = np.vstack(all_indices)

        # Process results
        for i, document in enumerate(tqdm(documents, desc="Processing documents")):
            contexts = []
            for dist, idx in zip(all_distances[i], all_indices[i]):
                try:
                    context = self._create_context_from_result(idx, dist, document)
                    if context:
                        contexts.append(context)
                except Exception as e:
                    print(f"Error processing result {idx}: {e}")

            document.contexts = contexts

        return documents

    def _create_context_from_result(self, idx: int, score: float, document: Document) -> Context:
        """Create Context object from search result."""
        doc_id = self.doc_ids[idx]
        doc_data = self.doc_texts.get(str(doc_id), {
            "text": "Text not found", 
            "title": "No Title"
        })

        return Context(
            id=doc_id,
            title=doc_data["title"],
            text=doc_data["text"],
            score=float(score),
            has_answer=has_answers(doc_data["text"], document.answers.answers, self.tokenizer_simple)
        )

retrieve(documents)

Retrieve relevant contexts using BGE.

Source code in rankify/retrievers/bge_retriever.py

def retrieve(self, documents: List[Document]) -> List[Document]:
    """Retrieve relevant contexts using BGE."""
    queries = [doc.question.question for doc in documents]
    print(f"Retrieving {len(documents)} documents with BGE...")

    # Encode queries
    query_embeddings = self._encode_queries(queries)

    # Check dimension compatibility
    if query_embeddings.shape[1] != self.index.d:
        raise ValueError(f"Dimension mismatch: queries={query_embeddings.shape[1]}, "
                       f"index={self.index.d}")

    # Batch FAISS search
    all_distances = []
    all_indices = []

    for start_idx in tqdm(range(0, len(query_embeddings), self.batch_size), 
                        desc="FAISS search"):
        end_idx = min(start_idx + self.batch_size, len(query_embeddings))
        batch_embeddings = query_embeddings[start_idx:end_idx]

        distances, indices = self.index.search(batch_embeddings, self.n_docs)
        all_distances.append(distances)
        all_indices.append(indices)

    # Combine results
    all_distances = np.vstack(all_distances)
    all_indices = np.vstack(all_indices)

    # Process results
    for i, document in enumerate(tqdm(documents, desc="Processing documents")):
        contexts = []
        for dist, idx in zip(all_distances[i], all_indices[i]):
            try:
                context = self._create_context_from_result(idx, dist, document)
                if context:
                    contexts.append(context)
            except Exception as e:
                print(f"Error processing result {idx}: {e}")

        document.contexts = contexts

    return documents

ColBERTRetriever

Bases: BaseRetriever

ColBERT retriever with backward compatibility for prebuilt indices.

Supports two modes: 1. Prebuilt indices (wiki, msmarco) - Original format with passages.tsv 2. Custom indices - New format with collection.tsv + ID mappings

Source code in rankify/retrievers/colbert_retriever.py

class ColBERTRetriever(BaseRetriever):
    """
    ColBERT retriever with backward compatibility for prebuilt indices.

    Supports two modes:
    1. Prebuilt indices (wiki, msmarco) - Original format with passages.tsv
    2. Custom indices - New format with collection.tsv + ID mappings
    """

    def __init__(self, model: str = "colbert-ir/colbertv2.0", index_type: str = "wiki", 
                 index_folder: str = None, **kwargs):
        super().__init__(**kwargs)
        self.model = model
        self.index_type = index_type
        self.index_folder = index_folder
        self.tokenizer = SimpleTokenizer()

        # Initialize index manager
        self.index_manager = IndexManager()

        # Initialize flag (will be properly set in _setup_paths)
        self.is_custom_index = False

        # Setup paths and determine index type
        self._setup_paths()

        # Load ID mappings (only for custom indices)
        self.id_mappings = self._load_id_mappings()

        # Initialize searcher
        self.searcher = self._initialize_searcher()

        # Load passages
        self.passages = self._load_passages()

    def _setup_paths(self):
        """Setup file paths based on whether it's a custom or prebuilt index."""
        if self.index_folder:
            # Check if it's actually a custom index by looking for custom files
            collection_file = os.path.join(self.index_folder, "collection.tsv")
            id_mapping_file = os.path.join(self.index_folder, "id_mapping.json")

            if os.path.exists(collection_file) and os.path.exists(id_mapping_file):
                # True custom index
                self.is_custom_index = True
                self.index_path = self.index_folder
                self.passages_file = os.path.join(self.index_folder, "passages.tsv")
                self.collection_file = collection_file
                self.id_mapping_file = id_mapping_file
                print(f"Using custom index: {self.index_folder}")
            else:
                # Prebuilt index with index_folder specified
                self.is_custom_index = False
                self.index_path = self.index_folder
                passages_file = os.path.join(self.index_folder, "passages.tsv")
                if os.path.exists(passages_file):
                    self.passages_file = passages_file
                else:
                    # Try to find passages file in parent directories
                    self._ensure_index_and_passages_downloaded()
                    self.passages_file = self._get_passages_file_path()
                self.collection_file = self.passages_file  # Same file for prebuilt indices
                self.id_mapping_file = None
                print(f"Using prebuilt index at: {self.index_folder}")
        else:
            # Standard prebuilt index
            self.is_custom_index = False
            self.index_path = os.path.join(self.index_manager.cache_dir, "index", "colbert")
            self._ensure_index_and_passages_downloaded()
            self.passages_file = self._get_passages_file_path()
            self.collection_file = self.passages_file  # Same file for prebuilt indices
            self.id_mapping_file = None
            print(f"Using prebuilt index: {self.index_type}")

    def _load_id_mappings(self):
        """Load ID mappings for custom indices only."""
        if not self.is_custom_index or not self.id_mapping_file:
            return None

        if not os.path.exists(self.id_mapping_file):
            print(f"Warning: ID mapping file not found: {self.id_mapping_file}")
            return None

        try:
            with open(self.id_mapping_file, 'r', encoding='utf-8') as f:
                mappings = json.load(f)

            # Convert string keys back to integers for sequential_to_original
            sequential_to_original = {int(k): v for k, v in mappings.get("sequential_to_original", {}).items()}
            original_to_sequential = mappings.get("original_to_sequential", {})

            print(f"✅ Loaded ID mappings: {len(original_to_sequential)} entries")
            return {
                "sequential_to_original": sequential_to_original,
                "original_to_sequential": original_to_sequential
            }

        except Exception as e:
            print(f"Warning: Error loading ID mappings: {e}")
            return None

    def _initialize_searcher(self):
        """Initialize ColBERT searcher with backward compatibility."""
        with Run().context(RunConfig(nranks=1, experiment="colbert")):

            if self.is_custom_index:
                # Custom index - use our format
                index_root = self.index_folder
                collection_path = self.collection_file

                # Verify collection file exists
                if not os.path.exists(collection_path):
                    raise FileNotFoundError(f"ColBERT collection file not found: {collection_path}")

                print(f"Initializing custom ColBERT searcher...")
                print(f"  Index root: {index_root}")
                print(f"  Collection: {collection_path}")

            else:
                # Prebuilt index - use original format
                if self.index_folder:
                    index_root = self.index_folder
                else:
                    index_root = os.path.join(self.index_path, self.index_type)
                collection_path = self.passages_file

                # Verify passages file exists (this is what prebuilt indices use)
                if not os.path.exists(collection_path):
                    raise FileNotFoundError(f"Passages file not found: {collection_path}")

                print(f"Initializing prebuilt ColBERT searcher...")
                print(f"  Index root: {index_root}")
                print(f"  Collection: {collection_path}")

            config = ColBERTConfig(
                root=index_root,
                collection=collection_path
            )

            return Searcher(index=index_root, config=config)

    def _ensure_index_and_passages_downloaded(self):
        """Download and extract ColBERT index and passages if needed (prebuilt indices only)."""
        if self.is_custom_index or self.index_folder:
            return  # Skip for custom indices or when index_folder is provided

        os.makedirs(self.index_path, exist_ok=True)

        if self.index_type not in self.index_manager.index_configs.get("colbert", {}):
            raise ValueError(f"Unsupported ColBERT index type: {self.index_type}")

        config = self.index_manager.index_configs["colbert"][self.index_type]

        # Check if index exists
        index_dir = os.path.join(self.index_path, self.index_type)
        if not os.path.exists(index_dir):
            urls = config.get("urls")
            if isinstance(urls, list):
                # Multi-part download
                for url in urls:
                    filename = self._extract_filename_from_url(url)
                    local_path = os.path.join(self.index_path, filename)
                    if not os.path.exists(local_path):
                        self._download_file(url, local_path)
                self._extract_multi_part_zip()
            else:
                # Single file download
                zip_path = os.path.join(self.index_path, "index.zip")
                if not os.path.exists(zip_path):
                    self._download_file(urls, zip_path)
                self._extract_zip_files()

        # Download passages if needed
        passages_url = config.get("passages_url")
        if passages_url:
            passages_filename = self._extract_filename_from_url(passages_url)
            passages_path = os.path.join(self.index_manager.cache_dir, passages_filename)
            if not os.path.exists(passages_path):
                self._download_file(passages_url, passages_path)

    def _get_passages_file_path(self):
        """Get path to passages file for prebuilt indices."""
        if self.index_folder and not self.is_custom_index:
            # Direct path provided for prebuilt index
            passages_path = os.path.join(self.index_folder, "passages.tsv")
            if os.path.exists(passages_path):
                return passages_path

        if self.index_type not in self.index_manager.index_configs.get("colbert", {}):
            raise ValueError(f"Unsupported ColBERT index type: {self.index_type}")

        config = self.index_manager.index_configs["colbert"][self.index_type]
        passages_url = config.get("passages_url")
        if passages_url:
            passages_filename = self._extract_filename_from_url(passages_url)
            return os.path.join(self.index_manager.cache_dir, passages_filename)
        return None

    def _extract_filename_from_url(self, url):
        """Extract filename from URL, ignoring query parameters."""
        parsed_url = urlparse(url)
        return os.path.basename(parsed_url.path).split('?')[0]

    def _download_file(self, url: str, save_path: str):
        """Download file from URL with progress bar."""
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        response = requests.get(url, stream=True)
        response.raise_for_status()

        with open(save_path, "wb") as f:
            for chunk in tqdm(response.iter_content(chunk_size=1024), 
                            desc=f"Downloading {os.path.basename(save_path)}"):
                f.write(chunk)

    def _extract_zip_files(self):
        """Extract ZIP files in the index folder."""
        zip_files = [f for f in os.listdir(self.index_path) if f.endswith(".zip")]

        for zip_file in zip_files:
            zip_path = os.path.join(self.index_path, zip_file)

            print(f"Extracting {zip_file}...")
            with zipfile.ZipFile(zip_path, "r") as zip_ref:
                zip_ref.extractall(self.index_path)

            os.remove(zip_path)

    def _extract_multi_part_zip(self):
        """Extract multi-part ZIP files by combining and decompressing them."""
        # Find all parts (handle both .001, .002 format and other naming patterns)
        all_files = os.listdir(self.index_path)
        zip_parts = []

        # Look for wiki.zip.001, wiki.zip.002, etc.
        for f in all_files:
            if f.startswith("wiki.zip.") and f.split('.')[-1].isdigit():
                zip_parts.append(f)

        if not zip_parts:
            print("No multi-part ZIP files found")
            return

        # Sort by part number
        zip_parts.sort(key=lambda x: int(x.split('.')[-1]))

        combined_zip_path = os.path.join(self.index_path, "combined.zip")

        print(f"Combining {len(zip_parts)} parts into {combined_zip_path}...")

        # Combine all parts
        with open(combined_zip_path, "wb") as combined:
            for part in zip_parts:
                part_path = os.path.join(self.index_path, part)
                print(f"Adding {part} to combined archive...")
                with open(part_path, "rb") as part_file:
                    # Use shutil.copyfileobj for better memory efficiency
                    import shutil
                    shutil.copyfileobj(part_file, combined)

        # Extract the combined ZIP
        print(f"Extracting combined ZIP file: {combined_zip_path}...")
        try:
            with zipfile.ZipFile(combined_zip_path, "r") as zip_ref:
                zip_ref.extractall(self.index_path)
            print("Extraction completed successfully")
        except zipfile.BadZipFile as e:
            print(f"Error: Combined file is not a valid ZIP: {e}")
            raise
        finally:
            # Clean up: remove combined file and parts
            if os.path.exists(combined_zip_path):
                os.remove(combined_zip_path)
            for part in zip_parts:
                part_path = os.path.join(self.index_path, part)
                if os.path.exists(part_path):
                    os.remove(part_path)

    def _load_passages(self):
        """Load passages with format detection."""
        if not self.passages_file or not os.path.exists(self.passages_file):
            print("Warning: Passages file not found, using empty passages dict")
            return {}

        passages = {}
        print(f"Loading passages from {self.passages_file}...")

        # Detect if the first line has integer or string IDs
        with open(self.passages_file, "r", encoding="utf-8") as f:
            first_line = f.readline().strip()
            f.seek(0)  # Reset to beginning

            # Check if header exists
            has_header = first_line.startswith('id\t') or first_line.lower().startswith('id\t')

            if has_header:
                next(f)  # Skip header

            # Process lines
            for line_num, line in enumerate(f, 1):
                try:
                    line = line.strip()
                    if not line:
                        continue

                    parts = line.split("\t")
                    if len(parts) >= 2:
                        pid = parts[0]  # Keep as string (works for both int and string IDs)
                        text = parts[1] if len(parts) > 1 else ""
                        title = parts[2] if len(parts) > 2 else ""

                        passages[pid] = {"text": text, "title": title}
                    else:
                        print(f"Warning: Malformed line {line_num} in passages file")

                except Exception as e:
                    print(f"Error processing line {line_num}: {e}")
                    continue

        print(f"✅ Loaded {len(passages)} passages")
        return passages

    def _convert_result_id(self, result_id: int) -> str:
        """Convert ColBERT result ID to the appropriate format."""
        if self.is_custom_index and self.id_mappings:
            # Custom index: convert sequential ID back to original string ID
            return self.id_mappings["sequential_to_original"].get(result_id, str(result_id))
        else:
            # Prebuilt index: IDs are already in correct format
            return str(result_id)

    def retrieve(self, documents: List[Document]) -> List[Document]:
        """Retrieve relevant contexts using ColBERT."""
        print(f"Retrieving {len(documents)} documents with ColBERT...")
        print(f"Mode: {'Custom index' if self.is_custom_index else f'Prebuilt index ({self.index_type})'}")

        for i, document in enumerate(tqdm(documents, desc="Processing documents")):
            query = document.question.question

            try:
                # ColBERT search returns (document_ids, ranks, scores)
                results = self.searcher.search(query, k=self.n_docs)
                document_ids, ranks, scores = results

                contexts = []
                for result_id, rank, score in zip(document_ids, ranks, scores):
                    try:
                        context = self._create_context_from_result(result_id, score, document)
                        if context:
                            contexts.append(context)
                    except Exception as e:
                        print(f"Error processing result {result_id}: {e}")

                document.contexts = contexts

            except Exception as e:
                print(f"Error searching for document {i}: {e}")
                document.contexts = []

        return documents

    def _create_context_from_result(self, result_id: int, score: float, document: Document) -> Context:
        """Create Context object from ColBERT search result."""
        # Convert result ID to appropriate format
        passage_id = self._convert_result_id(result_id)

        # Find passage by ID
        if passage_id in self.passages:
            passage = self.passages[passage_id]
            return Context(
                id=passage_id,
                title=passage["title"],
                text=passage["text"],
                score=float(score),
                has_answer=has_answers(passage["text"], document.answers.answers, self.tokenizer)
            )
        else:
            print(f"Warning: Passage {passage_id} (result_id: {result_id}) not found in passages dict")
            return Context(
                id=passage_id,
                title="Passage not found",
                text=f"Passage {passage_id} not found in corpus",
                score=float(score),
                has_answer=False
            )

retrieve(documents)

Retrieve relevant contexts using ColBERT.

Source code in rankify/retrievers/colbert_retriever.py

def retrieve(self, documents: List[Document]) -> List[Document]:
    """Retrieve relevant contexts using ColBERT."""
    print(f"Retrieving {len(documents)} documents with ColBERT...")
    print(f"Mode: {'Custom index' if self.is_custom_index else f'Prebuilt index ({self.index_type})'}")

    for i, document in enumerate(tqdm(documents, desc="Processing documents")):
        query = document.question.question

        try:
            # ColBERT search returns (document_ids, ranks, scores)
            results = self.searcher.search(query, k=self.n_docs)
            document_ids, ranks, scores = results

            contexts = []
            for result_id, rank, score in zip(document_ids, ranks, scores):
                try:
                    context = self._create_context_from_result(result_id, score, document)
                    if context:
                        contexts.append(context)
                except Exception as e:
                    print(f"Error processing result {result_id}: {e}")

            document.contexts = contexts

        except Exception as e:
            print(f"Error searching for document {i}: {e}")
            document.contexts = []

    return documents

ContrieverRetriever

Bases: BaseRetriever

FIXED Contriever retriever implementation with proper string ID handling.

Key fixes: - Handles both string and integer document IDs - Robust ID mapping and lookup - Better error handling for missing documents

Source code in rankify/retrievers/contriever_retriever.py

class ContrieverRetriever(BaseRetriever):
    """
    FIXED Contriever retriever implementation with proper string ID handling.

    Key fixes:
    - Handles both string and integer document IDs
    - Robust ID mapping and lookup
    - Better error handling for missing documents
    """

    def __init__(self, model: str = "facebook/contriever-msmarco", index_type: str = "wiki", 
                 index_folder: str = None, device: str = "cuda", **kwargs):
        super().__init__(**kwargs)
        self.model_path = model
        self.index_type = index_type
        self.index_folder = index_folder
        self.device = device
        self.tokenizer_simple = SimpleTokenizer()

        # Initialize index manager
        self.index_manager = IndexManager()

        # Setup paths
        if index_folder:
            self.index_path = index_folder
            self.passage_path = os.path.join(self.index_folder, "passages.tsv")
        else:
            self.embeddings_dir = os.path.join(self.index_manager.cache_dir, "index", "contriever_embeddings")
            self.index_path = os.path.join(self.embeddings_dir, self.index_type)
            self._ensure_index_and_passages_downloaded()
            self.passage_path = self._get_passage_path()

        # Load components
        self.index = self._load_index()
        print(f"Loading passages from: {self.passage_path}")
        self.passages = load_passages(self.passage_path)

        # FIXED: Handle both string and integer IDs properly
        self.passage_id_map = self._create_passage_id_map(self.passages)

        # Debug: Print some sample IDs to understand the format
        sample_ids = list(self.passage_id_map.keys())[:5]
        print(f"Sample passage IDs: {sample_ids}")
        print(f"Total passages loaded: {len(self.passage_id_map)}")

        # Load model
        self.model, self.tokenizer, _ = load_retriever(self.model_path)
        self.model = self.model.to(self.device).eval()

    def _create_passage_id_map(self, passages: List[Dict]) -> Dict[Union[str, int], Dict]:
        """
        FIXED: Create passage ID mapping that handles both string and integer IDs.

        Args:
            passages: List of passage dictionaries

        Returns:
            Dictionary mapping IDs (strings or integers) to passage data
        """
        passage_map = {}

        for passage in passages:
            original_id = passage["id"]

            # Try to convert to int, but keep as string if it fails
            try:
                # First try direct integer conversion
                int_id = int(original_id)
                passage_map[int_id] = passage
            except (ValueError, TypeError):
                # If that fails, keep as string
                passage_map[str(original_id)] = passage

                # Also try extracting numbers from string IDs for compatibility
                try:
                    # Extract numeric part if it's a string like "doc123" or similar
                    numeric_part = ''.join(filter(str.isdigit, str(original_id)))
                    if numeric_part:
                        int_id = int(numeric_part)
                        # Store both string and integer versions for flexibility
                        if int_id not in passage_map:
                            passage_map[int_id] = passage
                except:
                    pass  # If extraction fails, just use string

        return passage_map

    def _initialize_searcher(self):
        """Initialize searcher - not needed for Contriever as it uses direct FAISS."""
        return None

    def _ensure_index_and_passages_downloaded(self):
        """Download and extract Contriever index and passages if needed."""
        if self.index_type not in self.index_manager.index_configs.get("contriever", {}):
            raise ValueError(f"Unsupported Contriever index type: {self.index_type}")

        config = self.index_manager.index_configs["contriever"][self.index_type]

        # Handle embeddings download and extraction
        embeddings_url = config.get("url")
        if embeddings_url and not os.path.exists(self.index_path):
            os.makedirs(self.index_path, exist_ok=True)

            archive_name = self._clean_filename(embeddings_url)
            archive_path = os.path.join(self.index_path, archive_name)

            print(f"Downloading embeddings for '{self.index_type}'...")
            self._download_file(embeddings_url, archive_path)

            # Extract based on file type
            if archive_name.endswith(".tar"):
                print(f"Extracting TAR archive for '{self.index_type}'...")
                with tarfile.open(archive_path, "r") as tar:
                    tar.extractall(path=self.index_path)
            elif archive_name.endswith(".zip"):
                print(f"Extracting ZIP archive for '{self.index_type}'...")
                with zipfile.ZipFile(archive_path, "r") as zip_ref:
                    zip_ref.extractall(self.index_path)
            else:
                raise ValueError(f"Unsupported archive format: {archive_name}")

            os.remove(archive_path)  # Clean up

        # Handle passages download
        passages_url = config.get("passages_url")
        if passages_url:
            passage_filename = self._clean_filename(passages_url)
            passage_path = os.path.join(self.index_manager.cache_dir, passage_filename)

            if not os.path.exists(passage_path):
                print(f"Downloading passages '{passage_filename}'...")
                self._download_file(passages_url, passage_path)

    def _get_passage_path(self):
        """Get path to passages file."""
        if self.index_type not in self.index_manager.index_configs.get("contriever", {}):
            raise ValueError(f"Unsupported Contriever index type: {self.index_type}")

        config = self.index_manager.index_configs["contriever"][self.index_type]
        passages_url = config.get("passages_url")
        if passages_url:
            passage_filename = self._clean_filename(passages_url)
            return os.path.join(self.index_manager.cache_dir, passage_filename)
        return None

    def _clean_filename(self, url):
        """Extract filename from URL, removing query parameters."""
        parsed_url = urlparse(url)
        return os.path.basename(parsed_url.path).split('?')[0]

    def _download_file(self, url: str, save_path: str):
        """Download file from URL with progress bar."""
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        response = requests.get(url, stream=True)
        response.raise_for_status()

        with open(save_path, "wb") as f:
            for chunk in tqdm(response.iter_content(chunk_size=1024), 
                            desc=f"Downloading {os.path.basename(save_path)}"):
                f.write(chunk)

    def _load_index(self) -> Indexer:
        """Load or build the FAISS index using Contriever utilities."""
        # Get vector size from config
        if self.index_type in self.index_manager.index_configs.get("contriever", {}):
            config = self.index_manager.index_configs["contriever"][self.index_type]
            vector_size = config.get("vector_size", 768)
        else:
            vector_size = 768

        index = Indexer(vector_sz=vector_size, n_subquantizers=0, n_bits=8)

        # Determine index folder
        if self.index_folder:
            index_folder = self.index_folder
        else:
            index_folder = self.index_path
            if self.index_type == "wiki":
                index_folder = os.path.join(index_folder, "wikipedia_embeddings")

        index_path = os.path.join(index_folder, "index.faiss")

        if os.path.exists(index_path):
            print(f"Loading existing FAISS index from {index_folder}...")
            index.deserialize_from(index_folder)
        else:
            print(f"Building FAISS index from embeddings in {index_folder}...")
            # Look for both .pkl files and passages_XX files (Contriever format)
            embeddings_files = glob.glob(os.path.join(index_folder, "*.pkl"))
            if not embeddings_files:
                # Try Contriever's naming convention: passages_XX
                embeddings_files = glob.glob(os.path.join(index_folder, "passages_*"))

            embeddings_files = sorted(embeddings_files)

            if not embeddings_files:
                raise FileNotFoundError(f"No embedding files found in {index_folder}. "
                                      f"Looking for .pkl files or passages_* files.")

            print(f"Found {len(embeddings_files)} embedding files: {[os.path.basename(f) for f in embeddings_files[:5]]}...")
            self._index_encoded_data(index, embeddings_files)
            index.serialize(index_folder)

        return index

    def _debug_file_format(self, file_path):
        """Debug helper to understand file format."""
        try:
            with open(file_path, "rb") as f:
                first_bytes = f.read(100)
                print(f"First 100 bytes of {file_path}: {first_bytes}")

            # Check file size
            file_size = os.path.getsize(file_path)
            print(f"File size: {file_size} bytes")

        except Exception as e:
            print(f"Error reading file {file_path}: {e}")

    def _index_encoded_data(self, index, embedding_files, indexing_batch_size=1000000):
        """Load and index encoded passage embeddings into FAISS."""
        allids = []
        allembeddings = np.array([])

        # Debug first file to understand format
        if embedding_files:
            print(f"Debugging first file format...")
            self._debug_file_format(embedding_files[0])

        for i, file_path in enumerate(embedding_files):
            print(f"Loading embeddings from {file_path}")
            try:
                with open(file_path, "rb") as fin:
                    ids, embeddings = pickle.load(fin)
            except (pickle.UnpicklingError, EOFError) as e:
                print(f"Error loading pickle file {file_path}: {e}")
                print(f"Trying alternative loading method...")
                try:
                    # Try numpy format
                    data = np.load(file_path, allow_pickle=True)
                    if isinstance(data, tuple) and len(data) == 2:
                        ids, embeddings = data
                    else:
                        print(f"Unexpected data format in {file_path}")
                        continue
                except Exception as e2:
                    print(f"Failed to load {file_path} with alternative method: {e2}")
                    continue
            except Exception as e:
                print(f"Unexpected error loading {file_path}: {e}")
                continue

            allembeddings = np.vstack((allembeddings, embeddings)) if allembeddings.size else embeddings
            allids.extend(ids)

            # Process in batches to manage memory
            while allembeddings.shape[0] > indexing_batch_size:
                allembeddings, allids = self._add_embeddings(index, allembeddings, allids, indexing_batch_size)

        # Process remaining embeddings
        while allembeddings.shape[0] > 0:
            allembeddings, allids = self._add_embeddings(index, allembeddings, allids, indexing_batch_size)

        print("Data indexing completed.")

    def _add_embeddings(self, index, embeddings, ids, indexing_batch_size):
        """Add a batch of embeddings to the index."""
        end_idx = min(indexing_batch_size, embeddings.shape[0])
        ids_toadd = ids[:end_idx]
        embeddings_toadd = embeddings[:end_idx]
        ids = ids[end_idx:]
        embeddings = embeddings[end_idx:]
        index.index_data(ids_toadd, embeddings_toadd)
        return embeddings, ids

    def _embed_queries(self, queries: List[str]) -> np.ndarray:
        """Embed queries using the Contriever model with text normalization."""
        self.model.eval()
        embeddings = []
        batch_queries = []

        with torch.no_grad():
            for i, query in enumerate(queries):
                # Apply Contriever-specific preprocessing
                query = query.lower()
                query = normalize(query)
                batch_queries.append(query)

                if len(batch_queries) == self.batch_size or i == len(queries) - 1:
                    encoded_batch = self.tokenizer.batch_encode_plus(
                        batch_queries, 
                        return_tensors="pt", 
                        max_length=512, 
                        padding=True, 
                        truncation=True
                    )

                    # Move to device
                    for k, v in encoded_batch.items():
                        encoded_batch[k] = v.to(self.device)

                    output = self.model(**encoded_batch)
                    embeddings.append(output.cpu())
                    batch_queries = []

        embeddings = torch.cat(embeddings, dim=0)
        return embeddings.numpy()

    def retrieve(self, documents: List[Document]) -> List[Document]:
        """Retrieve relevant contexts using Contriever."""
        print(f"Retrieving {len(documents)} documents with Contriever...")

        # Prepare queries (remove question marks as in original)
        queries = [doc.question.question.replace("?", "") for doc in documents]

        # Embed queries
        query_embeddings = self._embed_queries(queries)

        # Search using FAISS index
        top_ids_and_scores = self.index.search_knn(
            query_embeddings, 
            self.n_docs, 
            index_batch_size=self.batch_size
        )

        # Process results
        for i, document in enumerate(tqdm(documents, desc="Processing documents")):
            top_ids, scores = top_ids_and_scores[i]
            contexts = []

            for doc_id, score in zip(top_ids, scores):
                try:
                    context = self._create_context_from_result(doc_id, score, document)
                    if context:
                        contexts.append(context)
                except (IndexError, KeyError, ValueError) as e:
                    print(f"Error processing result {doc_id}: {e}")
                    continue

            document.contexts = contexts

        return documents

    def _create_context_from_result(self, doc_id, score: float, document: Document) -> Context:
        """
        FIXED: Create Context object from Contriever search result with robust ID handling.
        """
        # Try multiple strategies to find the passage
        passage = None
        original_doc_id = doc_id

        # Strategy 1: Direct lookup (string or int)
        if doc_id in self.passage_id_map:
            passage = self.passage_id_map[doc_id]

        # Strategy 2: Try as string if it's currently an int
        elif str(doc_id) in self.passage_id_map:
            passage = self.passage_id_map[str(doc_id)]

        # Strategy 3: Try as int if it's currently a string  
        elif isinstance(doc_id, str):
            try:
                int_doc_id = int(doc_id)
                if int_doc_id in self.passage_id_map:
                    passage = self.passage_id_map[int_doc_id]
            except ValueError:
                pass

        # Strategy 4: Handle "doc" prefix removal
        if passage is None:
            try:
                clean_id = str(doc_id).replace("doc", "")
                if clean_id in self.passage_id_map:
                    passage = self.passage_id_map[clean_id]
                elif clean_id.isdigit() and int(clean_id) in self.passage_id_map:
                    passage = self.passage_id_map[int(clean_id)]
            except (ValueError, AttributeError):
                pass

        # Strategy 5: Extract numeric part from string IDs
        if passage is None and isinstance(doc_id, str):
            try:
                numeric_part = ''.join(filter(str.isdigit, str(doc_id)))
                if numeric_part:
                    numeric_id = int(numeric_part)
                    if numeric_id in self.passage_id_map:
                        passage = self.passage_id_map[numeric_id]
            except ValueError:
                pass

        if passage is None:
            print(f"Warning: Document {original_doc_id} not found in passage map")
            print(f"Available ID types: {type(list(self.passage_id_map.keys())[0]) if self.passage_id_map else 'None'}")
            return None

        # Get text content (handle different field names)
        text_content = passage.get("contents", passage.get("text", ""))

        return Context(
            id=original_doc_id,  # Keep original ID format
            title=passage.get("title", ""),
            text=text_content,
            score=float(score),
            has_answer=has_answers(
                text_content, 
                document.answers.answers, 
                self.tokenizer_simple, 
                regex=False
            )
        )

retrieve(documents)

Retrieve relevant contexts using Contriever.

Source code in rankify/retrievers/contriever_retriever.py

def retrieve(self, documents: List[Document]) -> List[Document]:
    """Retrieve relevant contexts using Contriever."""
    print(f"Retrieving {len(documents)} documents with Contriever...")

    # Prepare queries (remove question marks as in original)
    queries = [doc.question.question.replace("?", "") for doc in documents]

    # Embed queries
    query_embeddings = self._embed_queries(queries)

    # Search using FAISS index
    top_ids_and_scores = self.index.search_knn(
        query_embeddings, 
        self.n_docs, 
        index_batch_size=self.batch_size
    )

    # Process results
    for i, document in enumerate(tqdm(documents, desc="Processing documents")):
        top_ids, scores = top_ids_and_scores[i]
        contexts = []

        for doc_id, score in zip(top_ids, scores):
            try:
                context = self._create_context_from_result(doc_id, score, document)
                if context:
                    contexts.append(context)
            except (IndexError, KeyError, ValueError) as e:
                print(f"Error processing result {doc_id}: {e}")
                continue

        document.contexts = contexts

    return documents

OnlineRetriever

Source code in rankify/retrievers/online_retriever.py

class OnlineRetriever:
    def __init__(
        self,
        n_docs: int = 5,
        api_key: str = None,
        chunk_size: int = 500,
        chunk_overlap: int = 50,
        **kwags
    ) -> None:
        self.n_docs = n_docs
        self.searcher = WebSearchTool(search_provider_api_key=api_key)
        self.tokenizer = SimpleTokenizer()
        self.chunker = Chunker(
            chunk_size=chunk_size,
            chunk_overlap=chunk_overlap,
        )

    def _search_web(self, query: str) -> List[Any]:
        if not self.searcher.is_initialized:
            self.searcher.setup()
        return self.searcher.forward(query, num_result=self.n_docs)


    def retrieve(self, documents: List[Document]) -> List[Document]:
        for idx, doc in enumerate(tqdm(documents, desc="Fetching docs...")):
            question = doc.question.question
            logger.info(f"Fetching contexts for Q{idx}: {question}")

            sources = self._search_web(question)
            processed = []
            for s_idx, source in enumerate(sources):
                text = source.get('fit_markdown', '')
                # with open("log.txt", "a",  encoding="utf-8") as f:
                #     f.writelines(html)
                if text:
                    passages = self.chunker.split_text(text)
                    #print(len(passages), "passages found")
                    for p_idx, p in enumerate(passages):
                        #print(p)

                        #print(f"Processing snippet {s_idx} passage {p_idx}: {p['contents'][:50]}...")
                        processed.append({'id': f'{s_idx}_{p_idx}', 'contents': p})
                else:
                    snippet = source.get('snippet','')
                    if len(snippet.split()) > 20:
                        processed.append({'id': f's{ s_idx }_snip', 'contents': snippet})

            # Indexing
            with tempfile.TemporaryDirectory() as tmpdir:
                json.dump(processed, open(os.path.join(tmpdir,'docs.json'),'w'))
                subprocess.run([
                    'python','-m','pyserini.index.lucene',
                    '-collection','JsonCollection','-generator','DefaultLuceneDocumentGenerator',
                    '-input',tmpdir,'-index',tmpdir,
                    '-storePositions','-storeDocvectors','-storeRaw'
                ], check=True)

                searcher = LuceneSearcher(tmpdir)
                hits = searcher.search(question, k=self.n_docs)
                contexts: List[Context] = []
                for h_idx, hit in enumerate(hits):
                    raw = searcher.doc(hit.docid).raw()
                    data = json.loads(raw)
                    text = data.get('contents','')
                    title = "" #text.split("\n")[0]
                    try:
                        cid = int(hit.docid)
                    except ValueError:
                        cid = h_idx
                    contexts.append(Context(
                        id=cid,
                        title=title,
                        text=text,
                        score=hit.score,
                        has_answer=has_answers(text, doc.answers.answers, self.tokenizer)
                    ))
                doc.contexts = contexts
        return documents

HydeRetriever

Bases: BaseRetriever

Hypothetical Document Embedding (HyDE) Retriever implementation.

HyDE enhances document retrieval by generating hypothetical documents using an LLM and averaging their embeddings with the query embedding for improved search performance.

The retrieval process: 1. Generate hypothetical documents using an LLM based on the query 2. Encode both the original query and generated documents 3. Average the embeddings to obtain a refined query representation 4. Retrieve top documents using the averaged embedding

References

• Luyu Gao, Xueguang Ma, Jimmy Lin, and Jamie Callan. (2022): Precise Zero-Shot Dense Retrieval without Relevance Labels. https://arxiv.org/abs/2212.10496

Source code in rankify/retrievers/hyde_retriever.py

class HydeRetriever(BaseRetriever):
    """
    Hypothetical Document Embedding (HyDE) Retriever implementation.

    HyDE enhances document retrieval by generating hypothetical documents using an LLM
    and averaging their embeddings with the query embedding for improved search performance.

    The retrieval process:
    1. Generate hypothetical documents using an LLM based on the query
    2. Encode both the original query and generated documents
    3. Average the embeddings to obtain a refined query representation
    4. Retrieve top documents using the averaged embedding

    References:
        - Luyu Gao, Xueguang Ma, Jimmy Lin, and Jamie Callan. (2022): 
          Precise Zero-Shot Dense Retrieval without Relevance Labels.
          https://arxiv.org/abs/2212.10496
    """

    def __init__(self, model: str = "facebook/contriever-msmarco", index_type: str = "wiki", 
                 device: str = "cuda", api_key: str = None, **kwargs):
        super().__init__(**kwargs)
        self.model = model
        self.index_type = index_type
        self.device = device
        self.api_key = api_key

        # Initialize index manager
        self.index_manager = IndexManager()

        # Get configuration
        config = self._get_config()

        # HyDE-specific parameters
        self.task = kwargs.get("task", config.get("task", "web search"))
        self.llm_model = kwargs.get("llm_model", config.get("llm_model", "gpt-3.5-turbo-0125"))
        self.base_url = kwargs.get("base_url", config.get("base_url", None))
        self.num_generated_docs = kwargs.get("num_generated_docs", config.get("num_generated_docs", 1))
        self.max_token_generated_docs = kwargs.get("max_token_generated_docs", config.get("max_token_generated_docs", 512))
        self.temperature = kwargs.get("temperature", config.get("temperature", 0.7))

        # Validate API key if required
        if config.get("requires_api_key", True) and not self.api_key:
            raise ValueError("API key is required for HyDE retriever with LLM generation")

        # Initialize components
        self.tokenizer = SimpleTokenizer()
        self.promptor = Promptor(self.task)
        self.generator = OpenAIGenerator(
            model_name=self.llm_model,
            api_key=self.api_key,
            base_url=self.base_url,
            n=self.num_generated_docs,
            max_tokens=self.max_token_generated_docs,
            temperature=self.temperature,
            top_p=1,
            frequency_penalty=0.0,
            presence_penalty=0.0,
            stop=['\n\n\n'],
            wait_till_success=False
        )

        # Initialize the base retriever (Contriever by default)
        base_model = kwargs.get("base_model", config.get("base_model", model))
        base_index_type = kwargs.get("base_index_type", config.get("base_index_type", index_type))

        self.contriever = ContrieverRetriever(
            model=base_model,
            n_docs=self.n_docs,
            batch_size=self.batch_size,
            device=device,
            index_type=base_index_type
        )

        # Set searcher reference
        self.searcher = self.contriever.index

    def _get_config(self):
        """Get configuration for HyDE retriever."""
        hyde_configs = self.index_manager.index_configs.get("hyde", {})
        return hyde_configs.get(self.index_type, {
            "task": "web search",
            "llm_model": "gpt-3.5-turbo-0125",
            "num_generated_docs": 1,
            "max_token_generated_docs": 512,
            "temperature": 0.7,
            "requires_api_key": True
        })

    def _initialize_searcher(self):
        """Initialize searcher - handled through ContrieverRetriever."""
        return None

    def retrieve(self, documents: List[Document]) -> List[Document]:
        """
        Retrieve contexts using HyDE-based query expansion.

        The process involves:
        1. Generating hypothetical documents using an LLM
        2. Encoding both the original query and generated documents
        3. Averaging the embeddings to obtain a refined query representation
        4. Retrieving top documents using FAISS-based search

        Args:
            documents: List of Document objects containing queries

        Returns:
            List of Document objects with retrieved contexts
        """
        for i, document in enumerate(tqdm(documents, desc="Processing documents", unit="docs")):
            contexts = []

            # Build prompt and generate hypothetical documents
            prompt = self.promptor.build_prompt(document.question.question.replace("?", ""))
            hypothesis_documents = self.generator.generate(prompt)

            # Embed query and hypothetical documents
            all_emb_c = []
            for c in [prompt] + hypothesis_documents:
                c_emb = self.contriever._embed_queries([c])
                all_emb_c.append(c_emb.squeeze(0))

            # Average embeddings
            all_emb_c = np.array(all_emb_c)
            avg_emb_c = np.mean(all_emb_c, axis=0)
            hype_vector = avg_emb_c.reshape((1, len(avg_emb_c)))

            # Search using averaged embedding
            top_ids_and_scores = self.contriever.index.search_knn(
                hype_vector, 
                self.n_docs, 
                index_batch_size=self.batch_size
            )
            doc_ids, scores = top_ids_and_scores[0]

            # Create contexts from results
            for doc_id, score in zip(doc_ids, scores):
                try:
                    passage = self.contriever.passage_id_map[int(doc_id)]
                    context = Context(
                        id=int(doc_id),
                        title=passage["title"],
                        text=passage.get("text", passage.get("contents", "")),
                        score=score,
                        has_answer=has_answers(
                            passage.get("text", passage.get("contents", "")), 
                            document.answers.answers, 
                            self.tokenizer, 
                            regex=False
                        )
                    )
                    contexts.append(context)
                except (IndexError, KeyError):
                    # Log or handle the error, and continue with the next passage
                    continue

            document.contexts = contexts

        return documents

retrieve(documents)

Retrieve contexts using HyDE-based query expansion.

The process involves: 1. Generating hypothetical documents using an LLM 2. Encoding both the original query and generated documents 3. Averaging the embeddings to obtain a refined query representation 4. Retrieving top documents using FAISS-based search

Parameters:

Name	Type	Description	Default
documents	List[Document]	List of Document objects containing queries	required

Returns:

Type	Description
List[Document]	List of Document objects with retrieved contexts

Source code in rankify/retrievers/hyde_retriever.py

def retrieve(self, documents: List[Document]) -> List[Document]:
    """
    Retrieve contexts using HyDE-based query expansion.

    The process involves:
    1. Generating hypothetical documents using an LLM
    2. Encoding both the original query and generated documents
    3. Averaging the embeddings to obtain a refined query representation
    4. Retrieving top documents using FAISS-based search

    Args:
        documents: List of Document objects containing queries

    Returns:
        List of Document objects with retrieved contexts
    """
    for i, document in enumerate(tqdm(documents, desc="Processing documents", unit="docs")):
        contexts = []

        # Build prompt and generate hypothetical documents
        prompt = self.promptor.build_prompt(document.question.question.replace("?", ""))
        hypothesis_documents = self.generator.generate(prompt)

        # Embed query and hypothetical documents
        all_emb_c = []
        for c in [prompt] + hypothesis_documents:
            c_emb = self.contriever._embed_queries([c])
            all_emb_c.append(c_emb.squeeze(0))

        # Average embeddings
        all_emb_c = np.array(all_emb_c)
        avg_emb_c = np.mean(all_emb_c, axis=0)
        hype_vector = avg_emb_c.reshape((1, len(avg_emb_c)))

        # Search using averaged embedding
        top_ids_and_scores = self.contriever.index.search_knn(
            hype_vector, 
            self.n_docs, 
            index_batch_size=self.batch_size
        )
        doc_ids, scores = top_ids_and_scores[0]

        # Create contexts from results
        for doc_id, score in zip(doc_ids, scores):
            try:
                passage = self.contriever.passage_id_map[int(doc_id)]
                context = Context(
                    id=int(doc_id),
                    title=passage["title"],
                    text=passage.get("text", passage.get("contents", "")),
                    score=score,
                    has_answer=has_answers(
                        passage.get("text", passage.get("contents", "")), 
                        document.answers.answers, 
                        self.tokenizer, 
                        regex=False
                    )
                )
                contexts.append(context)
            except (IndexError, KeyError):
                # Log or handle the error, and continue with the next passage
                continue

        document.contexts = contexts

    return documents

DiverDenseRetriever

Bases: BaseRetriever

Comprehensive Diver-style retriever supporting: - SentenceTransformers (bge, sbert, nomic, instructor) - HF AutoModels (sf, qwen, e5, rader, contriever, m2) - GritLM (grit)

Source code in rankify/retrievers/diver_dense_retriever.py

class DiverDenseRetriever(BaseRetriever):
    """
    Comprehensive Diver-style retriever supporting:
    - SentenceTransformers (bge, sbert, nomic, instructor)
    - HF AutoModels (sf, qwen, e5, rader, contriever, m2)
    - GritLM (grit)
    """

    def __init__(
        self,
        corpus_path: str,
        model_id: str,      
        corpus_format: str = "jsonl",          
        text_field: str = "text",             
        title_field: str = "title",           
        id_field: str = "id",                

        task: str = "retrieval",
        instruction_query: str = "Represent this question for retrieving relevant passages: ",
        instruction_document: str = "Represent this passage for retrieval: ",
        doc_max_length: int = 2048, 
        query_max_length: int = 2048,

        cache_dir: str = "./cache",
        long_context: bool = False,
        encode_batch_size: int = 8,

        checkpoint: Optional[str] = None,
        device: Optional[str] = None,
        **kwargs,
    ):
        super().__init__(**kwargs)

        if corpus_path is None:
            raise ValueError("DiverDenseRetriever requires `corpus_path`")

        if model_id is None:
            raise ValueError("DiverDenseRetriever requires `model_id`")

        self.corpus_path = corpus_path
        self.corpus_format = corpus_format.lower()
        self.text_field = text_field
        self.title_field = title_field
        self.id_field = id_field

        self.model_id = model_id
        self.checkpoint = checkpoint
        self.task = task
        self.instruction_query = instruction_query
        self.instruction_document = instruction_document

        self.cache_dir = cache_dir
        self.long_context = long_context
        self.doc_max_length = doc_max_length
        self.query_max_length = query_max_length
        self.encode_batch_size = encode_batch_size
        if self.model_id == "grit":
            self.encode_batch_size = 1  # GritLM is sensitive to batching, force to 1

        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")

        self.tokenizer_simple = SimpleTokenizer()

        # Load corpus
        self.doc_ids, self.doc_texts, self.doc_titles = self._load_corpus()
        self._docid_to_idx = {str(did): i for i, did in enumerate(self.doc_ids)}

        # Load model
        self.model = None
        self.tokenizer = None
        self._load_model()

        self.doc_emb = self._load_or_build_doc_embeddings()

    def _load_model(self):
        # SentenceTransformer models
        if self.model_id == "bge":
            self.model = SentenceTransformer('BAAI/bge-large-en-v1.5', device=self.device)
        elif self.model_id == "sbert":
            self.model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2', device=self.device)
        elif self.model_id == "contriever_st":
            self.model = SentenceTransformer('nishimoto/contriever-sentencetransformer', device=self.device)
        elif self.model_id == "nomic":
            self.model = SentenceTransformer("nomic-ai/nomic-embed-text-v1", trust_remote_code=True, device=self.device)
        elif self.model_id == "diver":
            self.model = SentenceTransformer("AQ-MedAI/Diver-Retriever-4B", trust_remote_code=True, device=self.device)
        elif self.model_id == "inst-l":
            self.model = SentenceTransformer("hkunlp/instructor-large", device=self.device)
            self.model.max_seq_length = self.doc_max_length
        elif self.model_id == "inst-xl":
            self.model = SentenceTransformer("hkunlp/instructor-xl", device=self.device)
            self.model.max_seq_length = self.doc_max_length

        # HF AutoModel models (sf, e5, rader)
        elif self.model_id in ["sf", "e5", "rader"]:
            checkpoint = self.checkpoint
            if not checkpoint:
                if self.model_id == 'sf': checkpoint = 'Salesforce/SFR-Embedding-Mistral'
                elif self.model_id == 'e5': checkpoint = 'intfloat/e5-mistral-7b-instruct'
                elif self.model_id == 'rader': checkpoint = 'Raderspace/RaDeR_Qwen_25_7B_instruct_MATH_LLMq_CoT_lexical'

            self.tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
            self.model = AutoModel.from_pretrained(checkpoint, device_map="auto", trust_remote_code=True).eval()

        # M2
        elif self.model_id == "m2":
            checkpoint = self.checkpoint or "togethercomputer/m2-bert-80M-32k-retrieval"
            self.tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased", model_max_length=32768)
            self.model = AutoModelForSequenceClassification.from_pretrained(checkpoint, trust_remote_code=True).eval()
            # M2 usually runs on CPU/GPU via HF, ensure device
            self.model.to(self.device)

        # Contriever (Original HF)
        elif self.model_id == "contriever":
            checkpoint = self.checkpoint or 'facebook/contriever-msmarco'
            self.tokenizer = AutoTokenizer.from_pretrained(checkpoint)
            self.model = AutoModel.from_pretrained(checkpoint).to(self.device).eval()

        # TAS-B: DistilBERT dot-product model trained with Topic-Aware Sampling
        elif self.model_id == "tas-b":
            checkpoint = self.checkpoint or 'sebastian-hofstaetter/distilbert-dot-tas_b-b256-msmarco'
            self.tokenizer = AutoTokenizer.from_pretrained(checkpoint)
            self.model = AutoModel.from_pretrained(checkpoint).to(self.device).eval()

        # GritLM
        elif self.model_id == "grit":
            if not GRITLM_AVAILABLE:
                raise ImportError("gritlm is required for 'grit' model. Please install it.")
            checkpoint = self.checkpoint or 'GritLM/GritLM-7B'
            self.model = GritLM(checkpoint, torch_dtype="auto", mode="embedding")

        else:
            raise ValueError(f"The model {self.model_id} is not supported")

        # Handle device placement for models that don't auto-map
        if hasattr(self.model, "to") and self.model_id not in ["grit", "sf", "e5", "rader"]:
            self.model = self.model.to(self.device)

    def _load_corpus(self):
        if not os.path.exists(self.corpus_path):
            raise FileNotFoundError(f"Corpus not found: {self.corpus_path}")

        doc_ids: List[str] = []
        doc_texts: List[str] = []
        doc_titles: List[str] = []

        if self.corpus_format == "jsonl":
            with open(self.corpus_path, "r", encoding="utf-8") as f:
                for line in f:
                    if not line.strip():
                        continue
                    obj = json.loads(line)
                    doc_id = str(obj.get(self.id_field, len(doc_ids)))
                    text = obj.get(self.text_field, "") or ""
                    title = obj.get(self.title_field, "") or ""
                    doc_ids.append(doc_id)
                    doc_texts.append(text)
                    doc_titles.append(title)
        else:
            raise ValueError("Corpus format must be 'jsonl'")

        return doc_ids, doc_texts, doc_titles

    def _corpus_id(self):
        return os.path.basename(os.path.dirname(self.corpus_path))

    def _cache_file(self):
        # match Diver folder naming so you can reuse caches
        batch_size = self.encode_batch_size
        corpus_id = self._corpus_id()
        folder = os.path.join(self.cache_dir, "doc_emb", corpus_id, self.model_id, self.task, f"long_{self.long_context}_{batch_size}")
        os.makedirs(folder, exist_ok=True)
        return os.path.join(folder, "0.npy")

    def _encode_hf_auto_model(self, texts, max_length, instruction=None):
        # For sf, qwen, e5, rader
        if instruction:
            texts = add_instruct_concatenate(texts, self.task, instruction)

        all_embeddings = []
        batch_size = self.encode_batch_size

        for i in tqdm(range(0, len(texts), batch_size), desc="Encoding"):
            batch_texts = texts[i:i+batch_size]
            batch_dict = self.tokenizer(batch_texts, max_length=max_length, padding=True, truncation=True, return_tensors='pt')
            batch_dict = {k: v.to(self.device) for k, v in batch_dict.items()}

            with torch.no_grad():
                outputs = self.model(**batch_dict)
                if self.model_id == 'rader':
                     embeddings = outputs.last_hidden_state[:, -1, :] # Last token
                else:
                    embeddings = last_token_pool(outputs.last_hidden_state, batch_dict['attention_mask'])

                embeddings = F.normalize(embeddings, p=2, dim=1)
                all_embeddings.append(embeddings.cpu().numpy())

        return np.vstack(all_embeddings)

    def _encode_m2(self, texts, max_length, instruction=None):
        if instruction:
            texts = add_instruct_concatenate(texts, self.task, instruction)

        all_embeddings = []
        batch_size = self.encode_batch_size

        for i in tqdm(range(0, len(texts), batch_size), desc="Encoding M2"):
            batch_texts = texts[i:i+batch_size]
            batch_dict = self.tokenizer(batch_texts, max_length=max_length, padding=True, truncation=True, return_tensors='pt')
            # Move to device if model is on device
            if next(self.model.parameters()).is_cuda:
                 batch_dict = {k: v.to(next(self.model.parameters()).device) for k, v in batch_dict.items()}

            with torch.no_grad():
                outputs = self.model(**batch_dict)
                embeddings = outputs['sentence_embedding']
                embeddings = F.normalize(embeddings, p=2, dim=1)
                all_embeddings.append(embeddings.cpu().numpy())

        return np.vstack(all_embeddings)

    def _encode_contriever(self, texts, max_length):
        # Original HF Contriever
        all_embeddings = []
        batch_size = self.encode_batch_size

        for i in tqdm(range(0, len(texts), batch_size), desc="Encoding Contriever"):
            batch_texts = texts[i:i+batch_size]
            inputs = self.tokenizer(batch_texts, padding=True, truncation=True, return_tensors='pt')
            inputs = {k: v.to(self.device) for k, v in inputs.items()}

            with torch.no_grad():
                outputs = self.model(**inputs)
                embeddings = mean_pooling(outputs[0], inputs['attention_mask'])
                embeddings = F.normalize(embeddings, p=2, dim=1)
                all_embeddings.append(embeddings.cpu().numpy())

        return np.vstack(all_embeddings)

    def _encode_tasb(self, texts, max_length):
        """TAS-B: CLS-token encoding with L2 normalisation and dot-product scoring."""
        all_embeddings = []
        batch_size = self.encode_batch_size

        for i in tqdm(range(0, len(texts), batch_size), desc="Encoding TAS-B"):
            batch_texts = texts[i : i + batch_size]
            inputs = self.tokenizer(
                batch_texts, padding=True, truncation=True,
                max_length=max_length, return_tensors="pt"
            )
            inputs = {k: v.to(self.device) for k, v in inputs.items()}
            with torch.no_grad():
                outputs = self.model(**inputs)
                embeddings = outputs.last_hidden_state[:, 0, :]  # CLS token
                embeddings = F.normalize(embeddings, p=2, dim=1)
                all_embeddings.append(embeddings.cpu().numpy())

        return np.vstack(all_embeddings)

    def _load_or_build_doc_embeddings(self):
        cache_file = self._cache_file()
        if os.path.isfile(cache_file):
            doc_emb = np.load(cache_file, allow_pickle=True)
            return doc_emb

        docs = self.doc_texts

        # Encoding logic based on model_id
        if self.model_id in ["bge", "sbert", "contriever_st", "nomic", "diver"]:
            doc_emb = self.model.encode(docs, show_progress_bar=True, batch_size=self.encode_batch_size, normalize_embeddings=True)

        elif self.model_id in ["inst-l", "inst-xl"]:
            docs = add_instruct_list(docs, task=self.task, instruction=self.instruction_document)
            doc_emb = self.model.encode(docs, show_progress_bar=True, batch_size=self.encode_batch_size, normalize_embeddings=True)

        elif self.model_id == "e5":  # e5 paper uses strict "passage: {document}" for encoding
            docs = [f"passage: {d}" for d in docs]
            doc_emb = self._encode_hf_auto_model(docs, self.doc_max_length, instruction=None)

        elif self.model_id in ["sf", "rader"]:
            if self.model_id == 'rader':
                 docs = [f"document: {t[:8192]}" for t in docs]
            doc_emb = self._encode_hf_auto_model(docs, self.doc_max_length, instruction=None)

        elif self.model_id == "m2":
            doc_emb = self._encode_m2(docs, self.doc_max_length, instruction=None)

        elif self.model_id == "contriever":
            doc_emb = self._encode_contriever(docs, self.doc_max_length)

        elif self.model_id == "tas-b":
            doc_emb = self._encode_tasb(docs, self.doc_max_length)

        elif self.model_id == "grit":
            doc_emb = self.model.encode(docs, instruction=self.instruction_document, batch_size=self.encode_batch_size, max_length=self.doc_max_length)

        else:
            raise ValueError(f"Encoding not implemented for {self.model_id}")

        np.save(cache_file, doc_emb)

        assert len(self.doc_ids) == doc_emb.shape[0], (
            f"Embedding mismatch: {len(self.doc_ids)} docs vs {doc_emb.shape[0]} embeddings"
        )

        return doc_emb

    def _initialize_searcher(self):
        return None

    def _encode_queries(self, queries: List[str]):
        if self.model_id in ["bge", "sbert", "contriever_st", "nomic", "diver"]:
            if self.model_id == "bge":
                queries = add_instruct_concatenate(texts=queries, task=self.task, instruction=self.instruction_query)
            return self.model.encode(queries, show_progress_bar=True, batch_size=self.encode_batch_size, normalize_embeddings=True)

        elif self.model_id == "e5":
            queries = [f"query: {q}" for q in queries]    # e5 paper requires strict "query: {query}" when encoding
            return self._encode_hf_auto_model(queries, self.query_max_length, instruction=None)

        elif self.model_id in ["inst-l", "inst-xl"]:
            queries = add_instruct_list(texts=queries, task=self.task, instruction=self.instruction_query)
            return self.model.encode(queries, show_progress_bar=True, batch_size=self.encode_batch_size, normalize_embeddings=True)

        elif self.model_id in ["sf", "rader"]:
            if self.model_id == 'rader':
                queries = [f"query: {q}" for q in queries]
            return self._encode_hf_auto_model(queries, self.query_max_length, instruction=self.instruction_query if self.model_id != 'rader' else None)

        elif self.model_id == "m2":
            return self._encode_m2(queries, self.query_max_length, instruction=self.instruction_query)

        elif self.model_id == "contriever":
            return self._encode_contriever(queries, self.query_max_length)

        elif self.model_id == "tas-b":
            return self._encode_tasb(queries, self.query_max_length)

        elif self.model_id == "grit":
            q_instr = self.instruction_query.format(task=self.task) if "{task}" in self.instruction_query else self.instruction_query
            return self.model.encode(queries, instruction=q_instr, batch_size=self.encode_batch_size, max_length=self.query_max_length)

        raise ValueError(f"Encoding not implemented for {self.model_id}")

    def retrieve(self, documents: List[Document]):
        queries = [d.question.question for d in documents]
        query_ids = [str(d.id) if d.id is not None else str(i) for i, d in enumerate(documents)]

        query_emb = self._encode_queries(queries)

        # Ensure doc_emb is loaded
        if self.doc_emb is None:
            self.doc_emb = self._load_or_build_doc_embeddings()

        # Use torchmetrics for cosine similarity
        scores = pairwise_cosine_similarity(torch.from_numpy(query_emb), torch.from_numpy(self.doc_emb)).tolist()

        results = get_scores(
            query_ids=query_ids,
            doc_ids=self.doc_ids,
            scores=scores,
            num_hits=self.n_docs,
        )

        # Fill contexts
        for d, query_id in zip(documents, query_ids):
            ctxs: List[Context] = []
            for doc_id, score in results[query_id].items():
                idx = self._docid_to_idx.get(str(doc_id))
                if idx is None:
                    continue

                text = self.doc_texts[idx]
                title = self.doc_titles[idx] if idx < len(self.doc_titles) else ""

                ctxs.append(
                    Context(
                        id=str(doc_id),
                        title=title,
                        text=text,
                        score=float(score),
                        has_answer=has_answers(text, d.answers.answers, self.tokenizer_simple),
                    )
                )
            d.contexts = ctxs

        return documents

DiverBM25Retriever

Bases: BaseRetriever

Diver-style BM25 retriever using Gensim's LuceneBM25Model.

Source code in rankify/retrievers/diver_bm25_retriever.py

class DiverBM25Retriever(BaseRetriever):
    """
    Diver-style BM25 retriever using Gensim's LuceneBM25Model.
    """

    def __init__(
        self,
        corpus_path: str,
        corpus_format: str = "jsonl",
        text_field: str = "text",
        title_field: str = "title",
        id_field: str = "id",
        normalize_scores: bool = False,
        k1: float = 0.9,
        b: float = 0.4,
        **kwargs,
    ):
        super().__init__(**kwargs)

        if corpus_path is None:
            raise ValueError("DiverBM25Retriever requires `corpus_path`")

        self.corpus_path = corpus_path
        self.corpus_format = corpus_format.lower()
        self.text_field = text_field
        self.title_field = title_field
        self.id_field = id_field
        self.normalize_scores = normalize_scores
        self.k1 = k1
        self.b = b

        self.tokenizer_simple = SimpleTokenizer()
        self.analyzer = analysis.Analyzer(analysis.get_lucene_analyzer())

        self.dictionary = None
        self.doc_ids, self.doc_texts, self.doc_titles = self._load_corpus()
        self._docid_to_idx = {str(did): i for i, did in enumerate(self.doc_ids)}

        self.model, self.bm25_index = self._initialize_searcher()

    def _load_corpus(self):
        if not os.path.exists(self.corpus_path):
            raise FileNotFoundError(f"Corpus not found: {self.corpus_path}")

        doc_ids: List[str] = []
        doc_texts: List[str] = []
        doc_titles: List[str] = []

        if self.corpus_format == "jsonl":
            with open(self.corpus_path, "r", encoding="utf-8") as f:
                for line in f:
                    if not line.strip():
                        continue
                    obj = json.loads(line)
                    doc_id = str(obj.get(self.id_field, len(doc_ids)))
                    text = obj.get(self.text_field, "") or ""
                    title = obj.get(self.title_field, "") or ""
                    doc_ids.append(doc_id)
                    doc_texts.append(text)
                    doc_titles.append(title)
        else:
            raise ValueError("Corpus format must be 'jsonl'")

        return doc_ids, doc_texts, doc_titles

    def _initialize_searcher(self):
        corpus_analyzed = [self.analyzer.analyze(x) for x in self.doc_texts]
        self.dictionary = Dictionary(corpus_analyzed)
        model = LuceneBM25Model(dictionary=self.dictionary, k1=self.k1, b=self.b)
        bm25_corpus = model[list(map(self.dictionary.doc2bow, corpus_analyzed))]
        bm25_index = SparseMatrixSimilarity(bm25_corpus, num_docs=len(corpus_analyzed), num_terms=len(self.dictionary),
                                            normalize_queries=False, normalize_documents=False)
        return model, bm25_index

    def retrieve(self, documents: List[Document]) -> List[Document]:
        queries = [d.question.question for d in documents]
        query_ids = [str(d.id) if d.id is not None else str(i) for i, d in enumerate(documents)]

        all_scores_list = []
        for query in tqdm(queries, desc="BM25 retrieval"):
            query_analyzed = self.analyzer.analyze(query)
            bm25_query = self.model[self.dictionary.doc2bow(query_analyzed)]
            scores = self.bm25_index[bm25_query].tolist()
            if self.normalize_scores:
                scores = minmax_normalize(scores)
            all_scores_list.append(scores)

        results = get_scores(
            query_ids=query_ids,
            doc_ids=self.doc_ids,
            scores=all_scores_list,
            num_hits=self.n_docs,
        )

        for d, query_id in zip(documents, query_ids):
            ctxs: List[Context] = []
            for doc_id, score in results[query_id].items():
                idx = self._docid_to_idx.get(str(doc_id))
                if idx is None:
                    continue

                text = self.doc_texts[idx]
                title = self.doc_titles[idx] if idx < len(self.doc_titles) else ""

                ctxs.append(
                    Context(
                        id=str(doc_id),
                        title=title,
                        text=text,
                        score=float(score),
                        has_answer=has_answers(text, d.answers.answers, self.tokenizer_simple),
                    )
                )
            d.contexts = ctxs

        return documents

ReasonIRRetriever

Bases: BaseRetriever

Rankify retriever wrapper for ReasonIR (uses model.encode with instructions) - caches doc embeddings in cache_dir/doc_emb/corpus_id/reasonir/0.npy - encodes queries - cosine similarity + top-k

Source code in rankify/retrievers/reasonir_retriever.py

class ReasonIRRetriever(BaseRetriever):
    """
    Rankify retriever wrapper for ReasonIR (uses model.encode with instructions)
    - caches doc embeddings in cache_dir/doc_emb/corpus_id/reasonir/0.npy
    - encodes queries
    - cosine similarity + top-k
    """

    def __init__(
        self,
        corpus_path: str,
        checkpoint: str = "reasonir/ReasonIR-8B",
        task: str = "retrieval",
        query_instruction: str = "Represent this question for retrieving relevant passages: ",
        doc_instruction: str = "Represent this passage for retrieval: ",
        cache_dir: str = "./cache",
        batch_size: int = 4,
        query_max_length: int = 32768,
        doc_max_length: int = 32768,
        id_field: str = "id",
        title_field: str = "title",
        text_field: str = "text",
        device: Optional[str] = None,
        normalize: bool = True,
        **kwargs,
    ):
        super().__init__(**kwargs)

        if corpus_path is None:
            raise ValueError("ReasonirRetriever requires `corpus_path`")

        self.corpus_path = corpus_path
        self.corpus_format = "jsonl"
        self.checkpoint = checkpoint
        self.task = task
        self.query_instruction = query_instruction
        self.doc_instruction = doc_instruction

        self.cache_dir = cache_dir
        self.batch_size = batch_size
        self.query_max_length = query_max_length
        self.doc_max_length = doc_max_length
        self.normalize = normalize

        self.id_field = id_field
        self.title_field = title_field
        self.text_field = text_field

        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.simple_tokenizer = SimpleTokenizer()

        # Load corpus
        self.doc_ids, self.doc_texts, self.doc_titles = self._load_corpus()

        self._docid_to_idx = {str(did): i for i, did in enumerate(self.doc_ids)}

        self.model = SentenceTransformer(self.checkpoint, trust_remote_code=True, device=self.device)
        self.tokenizer = None

        # Cache / load doc embeddings
        self.doc_emb = self._load_or_build_doc_embeddings()

    def _load_corpus(self):
        if not os.path.exists(self.corpus_path):
            raise FileNotFoundError(f"Corpus not found: {self.corpus_path}")

        doc_ids: List[str] = []
        doc_texts: List[str] = []
        doc_titles: List[str] = []

        if self.corpus_format == "jsonl":
            with open(self.corpus_path, "r", encoding="utf-8") as f:
                for line in f:
                    if not line.strip():
                        continue
                    obj = json.loads(line)
                    doc_id = str(obj.get(self.id_field, len(doc_ids)))
                    text = obj.get(self.text_field, "") or ""
                    title = obj.get(self.title_field, "") or ""
                    doc_ids.append(doc_id)
                    doc_texts.append(text)
                    doc_titles.append(title)
        else:
            raise ValueError("Corpus format must be 'jsonl'")

        return doc_ids, doc_texts, doc_titles


    def _corpus_id(self):
        return os.path.basename(os.path.dirname(self.corpus_path))

    def _cache_file(self):
        corpus_id = self._corpus_id()

        folder = os.path.join(
            self.cache_dir,
            "doc_emb",
            corpus_id,
            "reasonir",
        )
        os.makedirs(folder, exist_ok=True)
        return os.path.join(folder, "0.npy")

    @torch.no_grad()
    def _encode(self, texts: List[str], instruction: str, max_length: int):
        emb = self.model.encode(
            texts,
            instruction=instruction,
            batch_size=self.batch_size,
            max_length=max_length,
            convert_to_numpy=True,
            normalize_embeddings=self.normalize,
        )

        return emb

    def _load_or_build_doc_embeddings(self):
        cache_file = self._cache_file()
        if os.path.isfile(cache_file):
            return np.load(cache_file, allow_pickle=True)

        doc_instr = self.doc_instruction.format(task=self.task) if "{task}" in self.doc_instruction else self.doc_instruction
        doc_emb = self._encode(self.doc_texts, instruction=doc_instr, max_length=self.doc_max_length)
        np.save(cache_file, doc_emb)
        return doc_emb

    def _initialize_searcher(self):
    # Retrievers don't use an external searcher
        return None

    def retrieve(self, documents: List[Document]):
        # Queries + ids
        queries = [d.question.question for d in documents]
        query_ids = [str(d.id) if d.id is not None else str(i) for i, d in enumerate(documents)]

        q_instr = self.query_instruction.format(task=self.task) if "{task}" in self.query_instruction else self.query_instruction
        query_emb = self._encode(queries, instruction=q_instr, max_length=self.query_max_length)

        # Similarity
        scores = pairwise_cosine_similarity(torch.from_numpy(query_emb), torch.from_numpy(self.doc_emb)).tolist()

        results = get_scores(
            query_ids=query_ids,
            doc_ids=self.doc_ids,
            scores=scores,
            num_hits=self.n_docs,
        )

        # Fill contexts
        for d, query_id in zip(documents, query_ids):
            ctxs: List[Context] = []
            for doc_id, score in results[query_id].items():
                idx = self._docid_to_idx.get(str(doc_id))
                if idx is None:
                    continue

                text = self.doc_texts[idx]
                title = self.doc_titles[idx] if idx < len(self.doc_titles) else ""

                ctxs.append(
                    Context(
                        id=str(doc_id),
                        title=title,
                        text=text,
                        score=float(score),
                        has_answer=has_answers(text, d.answers.answers, self.simple_tokenizer),
                    )
                )

            d.contexts = ctxs

        return documents

ReasonEmbedRetriever

Bases: BaseRetriever

A dense retriever that leverages ReasonEmbed models for reasoning-heavy retrieval tasks.

Attributes:

Name	Type	Description
model_id	str	Identifier for the embedding backbone (e.g. 'qwen3-8b', 'llama-8b').
encode_batch_size	int	Batch size for embedding generation.
device	str	Device to run the model on ('cuda' or 'cpu').

Source code in rankify/retrievers/reasonembed_retriever.py

class ReasonEmbedRetriever(BaseRetriever):
    """
    A dense retriever that leverages ReasonEmbed models for reasoning-heavy retrieval tasks.

    Attributes:
        model_id (str): Identifier for the embedding backbone (e.g. 'qwen3-8b', 'llama-8b').
        encode_batch_size (int): Batch size for embedding generation.
        device (str): Device to run the model on ('cuda' or 'cpu').
    """

    def __init__(
        self,
        corpus_path: str,
        model_id: str,      
        corpus_format: str = "jsonl",          
        text_field: str = "text",             
        title_field: str = "title",           
        id_field: str = "id",                
        cache_dir: str = "./cache",
        encode_batch_size: int = 8, 

        checkpoint: Optional[str] = None,
        device: Optional[str] = None,
        **kwargs,
    ):
        super().__init__(**kwargs)

        if corpus_path is None:
            raise ValueError("ReasonEmbedRetriever requires `corpus_path`")

        if model_id is None:
            raise ValueError("ReasonEmbedRetriever requires `model_id`")

        self.corpus_path = corpus_path
        self.corpus_format = corpus_format.lower()
        self.text_field = text_field
        self.title_field = title_field
        self.id_field = id_field
        self.tokenizer_simple = SimpleTokenizer()

        self.model_id = model_id
        self.checkpoint = checkpoint

        self.cache_dir = cache_dir
        self.encode_batch_size = encode_batch_size

        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")

        # Load corpus
        self.doc_ids, self.doc_texts, self.doc_titles = self._load_corpus()
        self._docid_to_idx = {str(did): i for i, did in enumerate(self.doc_ids)}

        # Load model
        self.model = None
        self.tokenizer = None
        self._load_model()

        self.doc_emb = self._load_or_build_doc_embeddings()

    def _load_model(self):
        # ReasonEmbed SentenceTransformer models
        if self.model_id == "qwen3-8b":
            self.model = SentenceTransformer(self.checkpoint or "hanhainebula/reason-embed-qwen3-8b-0928", trust_remote_code=True, device=self.device)
            self.model = self.model.to(self.device)
        elif self.model_id == "qwen3-4b":
            self.model = SentenceTransformer(self.checkpoint or "hanhainebula/reason-embed-qwen3-4b-0928", trust_remote_code=True, device=self.device)
            self.model = self.model.to(self.device)
        elif self.model_id == "llama-8b":
            self.model = SentenceTransformer(self.checkpoint or "hanhainebula/reason-embed-llama-3.1-8b-0928", trust_remote_code=True, device=self.device)
            self.model = self.model.to(self.device)
        else:
            raise ValueError(f"The model {self.model_id} is not supported")

    def _load_corpus(self):
        if not os.path.exists(self.corpus_path):
            raise FileNotFoundError(f"Corpus not found: {self.corpus_path}")

        doc_ids: List[str] = []
        doc_texts: List[str] = []
        doc_titles: List[str] = []

        if self.corpus_format == "jsonl":
            with open(self.corpus_path, "r", encoding="utf-8") as f:
                for line in f:
                    if not line.strip():
                        continue
                    obj = json.loads(line)
                    doc_id = str(obj.get(self.id_field, len(doc_ids)))
                    text = obj.get(self.text_field, "") or ""
                    title = obj.get(self.title_field, "") or ""
                    doc_ids.append(doc_id)
                    doc_texts.append(text)
                    doc_titles.append(title)
        else:
            raise ValueError("Corpus format must be 'jsonl'")

        return doc_ids, doc_texts, doc_titles

    def _corpus_id(self):
        return os.path.basename(os.path.dirname(self.corpus_path))

    def _cache_file(self):
        corpus_id = self._corpus_id()
        folder = os.path.join(self.cache_dir, "doc_emb", corpus_id, self.model_id)
        os.makedirs(folder, exist_ok=True)
        return os.path.join(folder, "0.npy")

    def _load_or_build_doc_embeddings(self):
        cache_file = self._cache_file()
        if os.path.isfile(cache_file):
            doc_emb = np.load(cache_file, allow_pickle=True)
            return doc_emb

        docs = self.doc_texts

        # Encoding logic based on model_id
        if self.model_id in ["qwen3-4b", "qwen3-8b", "llama-8b"]:
            doc_emb = self.model.encode(docs, show_progress_bar=True, batch_size=self.encode_batch_size, normalize_embeddings=True)
        else:
            raise ValueError(f"Encoding not implemented for {self.model_id}")

        np.save(cache_file, doc_emb)

        assert len(self.doc_ids) == doc_emb.shape[0], (
            f"Embedding mismatch: {len(self.doc_ids)} docs vs {doc_emb.shape[0]} embeddings"
        )

        return doc_emb

    def _initialize_searcher(self):
        return None

    def _encode_queries(self, queries: List[str]):
        # ReasonEmbed models require a specific 'Instruct' prefix to trigger
        # reasoning capabilities during the embedding process.
        if self.model_id in ["qwen3-4b", "qwen3-8b", "llama-8b"]:
            task = "Given a search query, retrieve relevant passages that answer the query."
            instruct_queries = [f"Instruct: {task}\nQuery: {q}" for q in queries]
            return self.model.encode(instruct_queries, show_progress_bar=True, batch_size=self.encode_batch_size, normalize_embeddings=True)
        raise ValueError(f"Encoding not implemented for {self.model_id}")

    def retrieve(self, documents: List[Document]):
        queries = [d.question.question for d in documents]
        query_ids = [str(d.id) if d.id is not None else str(i) for i, d in enumerate(documents)]

        query_emb = self._encode_queries(queries)

        # Ensure doc_emb is loaded
        if self.doc_emb is None:
            self.doc_emb = self._load_or_build_doc_embeddings()

        # Use torchmetrics for cosine similarity
        scores = pairwise_cosine_similarity(torch.from_numpy(query_emb), torch.from_numpy(self.doc_emb)).tolist()

        results = get_scores(
            query_ids=query_ids,
            doc_ids=self.doc_ids,
            scores=scores,
            num_hits=self.n_docs,
        )

        # Fill contexts
        for d, query_id in zip(documents, query_ids):
            ctxs: List[Context] = []
            for doc_id, score in results[query_id].items():
                idx = self._docid_to_idx.get(str(doc_id))
                if idx is None:
                    continue

                text = self.doc_texts[idx]
                title = self.doc_titles[idx] if idx < len(self.doc_titles) else ""

                ctxs.append(
                    Context(
                        id=str(doc_id),
                        title=title,
                        text=text,
                        score=float(score),
                        has_answer=has_answers(text, d.answers.answers, self.tokenizer_simple),
                    )
                )
            d.contexts = ctxs

        return documents

BgeReasonerRetriever

Bases: BaseRetriever

Retriever for the BGE-Reasoner-Embed-Qwen3-8B model. - caches doc embeddings in cache_dir/doc_emb/corpus_id/bge-reasoner-embed/0.npy - encodes queries - cosine similarity + top-k

Source code in rankify/retrievers/bge_reasoner_retriever.py

class BgeReasonerRetriever(BaseRetriever):
    '''
    Retriever for the BGE-Reasoner-Embed-Qwen3-8B model.
    - caches doc embeddings in cache_dir/doc_emb/corpus_id/bge-reasoner-embed/0.npy
    - encodes queries
    - cosine similarity + top-k
    '''

    def __init__(
        self,
        corpus_path: str,
        model_id: str = "bge-reasoner-embed",      
        corpus_format: str = "jsonl",          
        text_field: str = "text",             
        title_field: str = "title",           
        id_field: str = "id",                

        cache_dir: str = "./cache",
        encode_batch_size: int = 8,

        checkpoint: Optional[str] = None, # Override default checkpoint
        device: Optional[str] = None,
        **kwargs,
    ):
        super().__init__(**kwargs)

        if corpus_path is None:
            raise ValueError("BgeReasonerEmbedRetriever requires `corpus_path`")

        if model_id is None:
            raise ValueError("BgeReasonerEmbedRetriever requires `model_id`")

        self.corpus_path = corpus_path
        self.corpus_format = corpus_format.lower()
        self.text_field = text_field
        self.title_field = title_field
        self.id_field = id_field
        self.tokenizer_simple = SimpleTokenizer()

        self.model_id = model_id
        self.checkpoint = checkpoint

        self.cache_dir = cache_dir
        self.encode_batch_size = encode_batch_size

        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")

        # Load corpus
        self.doc_ids, self.doc_texts, self.doc_titles = self._load_corpus()
        self._docid_to_idx = {str(did): i for i, did in enumerate(self.doc_ids)}

        # Load model
        self.model = None
        self.tokenizer = None
        self._load_model()

        self.doc_emb = self._load_or_build_doc_embeddings()

    def _load_model(self):
        # ReasonEmbed SentenceTransformer models
        if self.model_id == "bge-reasoner-embed":
            self.model = SentenceTransformer(self.checkpoint or "BAAI/bge-reasoner-embed-qwen3-8b-0923", trust_remote_code=True, device=self.device)
            self.model = self.model.to(self.device)
        else:
            raise ValueError(f"The model {self.model_id} is not supported")

    def _load_corpus(self):
        if not os.path.exists(self.corpus_path):
            raise FileNotFoundError(f"Corpus not found: {self.corpus_path}")

        doc_ids: List[str] = []
        doc_texts: List[str] = []
        doc_titles: List[str] = []

        if self.corpus_format == "jsonl":
            with open(self.corpus_path, "r", encoding="utf-8") as f:
                for line in f:
                    if not line.strip():
                        continue
                    obj = json.loads(line)
                    doc_id = str(obj.get(self.id_field, len(doc_ids)))
                    text = obj.get(self.text_field, "") or ""
                    title = obj.get(self.title_field, "") or ""
                    doc_ids.append(doc_id)
                    doc_texts.append(text)
                    doc_titles.append(title)
        else:
            raise ValueError("Corpus format must be 'jsonl'")

        return doc_ids, doc_texts, doc_titles

    def _corpus_id(self):
        return os.path.basename(os.path.dirname(self.corpus_path))

    def _cache_file(self):
        corpus_id = self._corpus_id()
        folder = os.path.join(self.cache_dir, "doc_emb", corpus_id, self.model_id)
        os.makedirs(folder, exist_ok=True)
        return os.path.join(folder, "0.npy")

    def _load_or_build_doc_embeddings(self):
        cache_file = self._cache_file()
        if os.path.isfile(cache_file):
            doc_emb = np.load(cache_file, allow_pickle=True)
            return doc_emb

        docs = self.doc_texts

        # Encoding logic based on model_id
        if self.model_id == "bge-reasoner-embed":
            doc_emb = self.model.encode(docs, show_progress_bar=True, batch_size=self.encode_batch_size, normalize_embeddings=True)
        else:
            raise ValueError(f"Encoding not implemented for {self.model_id}")

        np.save(cache_file, doc_emb)

        assert len(self.doc_ids) == doc_emb.shape[0], (
            f"Embedding mismatch: {len(self.doc_ids)} docs vs {doc_emb.shape[0]} embeddings"
        )

        return doc_emb

    def _initialize_searcher(self):
        return None

    def _encode_queries(self, queries: List[str]):
        if self.model_id == "bge-reasoner-embed":
            instruction = "Given a search query, retrieve relevant passages that answer the query."
            instruct_queries = [f"Instruct: {instruction}\nQuery: {q}" for q in queries]
            return self.model.encode(instruct_queries, show_progress_bar=True, batch_size=self.encode_batch_size, normalize_embeddings=True)
        raise ValueError(f"Encoding not implemented for {self.model_id}")

    def retrieve(self, documents: List[Document]):
        queries = [d.question.question for d in documents]
        query_ids = [str(d.id) if d.id is not None else str(i) for i, d in enumerate(documents)]

        query_emb = self._encode_queries(queries)

        # Ensure doc_emb is loaded
        if self.doc_emb is None:
            self.doc_emb = self._load_or_build_doc_embeddings()

        # Use torchmetrics for cosine similarity
        scores = pairwise_cosine_similarity(torch.from_numpy(query_emb), torch.from_numpy(self.doc_emb)).tolist()

        results = get_scores(
            query_ids=query_ids,
            doc_ids=self.doc_ids,
            scores=scores,
            num_hits=self.n_docs,
        )

        # Fill contexts
        for d, query_id in zip(documents, query_ids):
            ctxs: List[Context] = []
            for doc_id, score in results[query_id].items():
                idx = self._docid_to_idx.get(str(doc_id))
                if idx is None:
                    continue

                text = self.doc_texts[idx]
                title = self.doc_titles[idx] if idx < len(self.doc_titles) else ""

                ctxs.append(
                    Context(
                        id=str(doc_id),
                        title=title,
                        text=text,
                        score=float(score),
                        has_answer=has_answers(text, d.answers.answers, self.tokenizer_simple),
                    )
                )
            d.contexts = ctxs

        return documents

UniCOILRetriever

Bases: BaseRetriever

UniCOIL learned sparse retriever using Pyserini's LuceneImpactSearcher.

UniCOIL (Unified COmpact Index with Learned sparse representation) learns term weights via a BERT model with a linear projection to a scalar weight per token position.

References

• Lin & Ma (2021): "A Few Brief Notes on DeepImpact, COIL, and a Conceptual Framework for Information Retrieval Techniques." https://arxiv.org/abs/2106.14807
• Model: castorini/unicoil-msmarco-passage

Parameters:

Name	Type	Description	Default
index_type	str	'msmarco' to use the Pyserini prebuilt MS MARCO unicoil index, or 'msmarco-noexp' for the no-expansion variant. Supply index_folder for a custom Lucene impact index.	'msmarco'
index_folder	str	Path to a custom Lucene impact-index directory.	None
corpus_path	str	JSONL file for passage text lookup (fields: id, text, title). Required when using a custom index_folder because the Lucene index may not store raw text.	None
model_name	str	UniCOIL query encoder checkpoint.	'castorini/unicoil-msmarco-passage'
device	str	'cpu' or 'cuda'.	'cpu'

Example

from rankify.retrievers.retriever import Retriever

retriever = Retriever(
    method="unicoil",
    n_docs=10,
    index_type="msmarco",
    corpus_path="msmarco_passages.jsonl",
)

Source code in rankify/retrievers/unicoil_retriever.py

class UniCOILRetriever(BaseRetriever):
    """
    UniCOIL learned sparse retriever using Pyserini's LuceneImpactSearcher.

    UniCOIL (Unified COmpact Index with Learned sparse representation) learns
    term weights via a BERT model with a linear projection to a scalar weight
    per token position.

    References:
        - Lin & Ma (2021): "A Few Brief Notes on DeepImpact, COIL, and a Conceptual
          Framework for Information Retrieval Techniques."
          https://arxiv.org/abs/2106.14807
        - Model: castorini/unicoil-msmarco-passage

    Args:
        index_type (str): ``'msmarco'`` to use the Pyserini prebuilt MS MARCO unicoil
            index, or ``'msmarco-noexp'`` for the no-expansion variant.
            Supply ``index_folder`` for a custom Lucene impact index.
        index_folder (str, optional): Path to a custom Lucene impact-index directory.
        corpus_path (str, optional): JSONL file for passage text lookup
            (fields: ``id``, ``text``, ``title``).  Required when using a custom
            ``index_folder`` because the Lucene index may not store raw text.
        model_name (str): UniCOIL query encoder checkpoint.
        device (str): ``'cpu'`` or ``'cuda'``.

    Example:
        ```python
        from rankify.retrievers.retriever import Retriever

        retriever = Retriever(
            method="unicoil",
            n_docs=10,
            index_type="msmarco",
            corpus_path="msmarco_passages.jsonl",
        )
        ```
    """

    PREBUILT_INDEX_MAP = {
        "msmarco": "msmarco-v1-passage.unicoil",
        "msmarco-noexp": "msmarco-v1-passage.unicoil-noexp",
    }

    def __init__(
        self,
        index_type: str = "msmarco",
        index_folder: Optional[str] = None,
        corpus_path: Optional[str] = None,
        model_name: str = "castorini/unicoil-msmarco-passage",
        device: str = "cpu",
        **kwargs,
    ):
        super().__init__(**kwargs)
        if not PYSERINI_AVAILABLE:
            raise ImportError(
                "Pyserini is required for UniCOILRetriever. "
                "Install with: pip install pyserini"
            )
        self.index_type = index_type
        self.index_folder = index_folder
        self.corpus_path = corpus_path
        self.model_name = model_name
        self.device = device

        self.passage_store: dict = {}
        if corpus_path:
            self._load_corpus(corpus_path)

        self.searcher = self._initialize_searcher()

    def _load_corpus(self, corpus_path: str) -> None:
        """Load passage texts from a JSONL file into a dict keyed by id."""
        print(f"Loading corpus from {corpus_path}...")
        with open(corpus_path, "r", encoding="utf-8") as fh:
            for line in fh:
                if not line.strip():
                    continue
                obj = json.loads(line)
                pid = str(obj.get("id", obj.get("docid", "")))
                self.passage_store[pid] = {
                    "text": obj.get("text", obj.get("contents", "")),
                    "title": obj.get("title", ""),
                }
        print(f"Loaded {len(self.passage_store)} passages.")

    def _initialize_searcher(self) -> "LuceneImpactSearcher":
        """Build and return a LuceneImpactSearcher with a UniCOIL query encoder."""
        query_encoder = UniCoilQueryEncoder(self.model_name, device=self.device)
        if self.index_folder:
            return LuceneImpactSearcher(self.index_folder, query_encoder)
        prebuilt = self.PREBUILT_INDEX_MAP.get(self.index_type)
        if prebuilt is None:
            raise ValueError(
                f"Unknown index_type '{self.index_type}'. "
                f"Supported values: {list(self.PREBUILT_INDEX_MAP.keys())}, "
                "or pass index_folder."
            )
        return LuceneImpactSearcher.from_prebuilt_index(prebuilt, query_encoder)

    def _get_passage(self, docid: str):
        """Return (text, title) for *docid*, falling back to Lucene stored field."""
        passage = self.passage_store.get(docid)
        if passage is not None:
            return passage["text"], passage["title"]
        try:
            raw = self.searcher.doc(docid)
            if raw:
                parsed = json.loads(raw.raw())
                text = parsed.get("contents", parsed.get("text", ""))
                title = parsed.get("title", "")
                return text, title
        except Exception:
            pass
        return "Text not available", ""

    def retrieve(self, documents: List[Document]) -> List[Document]:
        """
        Retrieve passages for each document using the UniCOIL sparse index.

        Args:
            documents (List[Document]): Documents whose ``question`` field
                contains the query.

        Returns:
            List[Document]: Documents with populated ``contexts``.
        """
        for document in tqdm(documents, desc="UniCOIL retrieving"):
            query = document.question.question
            hits = self.searcher.search(query, k=self.n_docs)
            contexts: List[Context] = []
            for hit in hits:
                docid = str(hit.docid)
                score = float(hit.score)
                text, title = self._get_passage(docid)

                has_ans = False
                if document.answers and document.answers.answers:
                    try:
                        from pyserini.eval.evaluate_dpr_retrieval import (
                            has_answers,
                            SimpleTokenizer,
                        )
                        has_ans = has_answers(
                            text, document.answers.answers, SimpleTokenizer()
                        )
                    except Exception:
                        pass

                contexts.append(
                    Context(
                        id=docid,
                        title=title,
                        text=text,
                        score=score,
                        has_answer=has_ans,
                    )
                )
            document.contexts = contexts
        return documents

retrieve(documents)

Retrieve passages for each document using the UniCOIL sparse index.

Parameters:

Name	Type	Description	Default
documents	List[Document]	Documents whose question field contains the query.	required

Returns:

Type	Description
List[Document]	List[Document]: Documents with populated contexts.

Source code in rankify/retrievers/unicoil_retriever.py

def retrieve(self, documents: List[Document]) -> List[Document]:
    """
    Retrieve passages for each document using the UniCOIL sparse index.

    Args:
        documents (List[Document]): Documents whose ``question`` field
            contains the query.

    Returns:
        List[Document]: Documents with populated ``contexts``.
    """
    for document in tqdm(documents, desc="UniCOIL retrieving"):
        query = document.question.question
        hits = self.searcher.search(query, k=self.n_docs)
        contexts: List[Context] = []
        for hit in hits:
            docid = str(hit.docid)
            score = float(hit.score)
            text, title = self._get_passage(docid)

            has_ans = False
            if document.answers and document.answers.answers:
                try:
                    from pyserini.eval.evaluate_dpr_retrieval import (
                        has_answers,
                        SimpleTokenizer,
                    )
                    has_ans = has_answers(
                        text, document.answers.answers, SimpleTokenizer()
                    )
                except Exception:
                    pass

            contexts.append(
                Context(
                    id=docid,
                    title=title,
                    text=text,
                    score=score,
                    has_answer=has_ans,
                )
            )
        document.contexts = contexts
    return documents

SpladeV2Retriever

Bases: BaseRetriever

SPLADE-v2 learned sparse retriever using Pyserini's LuceneImpactSearcher.

SPLADE learns sparse, high-dimensional representations by predicting importance weights over the vocabulary for each token position, enabling efficient inverted-index retrieval with neural relevance signals.

References

• Formal et al. (2022): "From Distillation to Hard Negative Sampling" https://arxiv.org/abs/2205.04733
• Prebuilt index: msmarco-v1-passage.splade-pp-ed

Parameters:

Name	Type	Description	Default
index_type	str	Prebuilt index alias. One of: 'splade-pp-ed' – SPLADE++ EnsembleDistil (default) 'splade-pp-sd' – SPLADE++ SelfDistil Pass index_folder to use a custom Lucene impact index.	'splade-pp-ed'
index_folder	str	Path to a custom Lucene impact-index dir.	None
corpus_path	str	JSONL file for passage text lookup (fields: id, text, title).	None
model_name	str	SPLADE query encoder checkpoint.	None
device	str	'cpu' or 'cuda'.	'cpu'

Example

from rankify.retrievers.retriever import Retriever

retriever = Retriever(
    method="splade-v2",
    n_docs=10,
    index_type="splade-pp-ed",
    corpus_path="msmarco_passages.jsonl",
)

Source code in rankify/retrievers/splade_v2_retriever.py

class SpladeV2Retriever(BaseRetriever):
    """
    SPLADE-v2 learned sparse retriever using Pyserini's LuceneImpactSearcher.

    SPLADE learns sparse, high-dimensional representations by predicting
    importance weights over the vocabulary for each token position, enabling
    efficient inverted-index retrieval with neural relevance signals.

    References:
        - Formal et al. (2022): "From Distillation to Hard Negative Sampling"
          https://arxiv.org/abs/2205.04733
        - Prebuilt index: ``msmarco-v1-passage.splade-pp-ed``

    Args:
        index_type (str): Prebuilt index alias.  One of:

            * ``'splade-pp-ed'`` – SPLADE++ EnsembleDistil (default)
            * ``'splade-pp-sd'`` – SPLADE++ SelfDistil

            Pass ``index_folder`` to use a custom Lucene impact index.
        index_folder (str, optional): Path to a custom Lucene impact-index dir.
        corpus_path (str, optional): JSONL file for passage text lookup
            (fields: ``id``, ``text``, ``title``).
        model_name (str): SPLADE query encoder checkpoint.
        device (str): ``'cpu'`` or ``'cuda'``.

    Example:
        ```python
        from rankify.retrievers.retriever import Retriever

        retriever = Retriever(
            method="splade-v2",
            n_docs=10,
            index_type="splade-pp-ed",
            corpus_path="msmarco_passages.jsonl",
        )
        ```
    """

    PREBUILT_INDEX_MAP = {
        "splade-pp-ed": "msmarco-v1-passage.splade-pp-ed",
        "splade-pp-sd": "msmarco-v1-passage.splade-pp-sd",
        # Convenience aliases
        "msmarco": "msmarco-v1-passage.splade-pp-ed",
        "msmarco-splade-pp-ed": "msmarco-v1-passage.splade-pp-ed",
        "msmarco-splade-pp-sd": "msmarco-v1-passage.splade-pp-sd",
    }

    DEFAULT_MODEL_MAP = {
        "splade-pp-ed": "naver/splade-cocondenser-ensembledistil",
        "splade-pp-sd": "naver/splade-cocondenser-selfdistil",
        "msmarco": "naver/splade-cocondenser-ensembledistil",
        "msmarco-splade-pp-ed": "naver/splade-cocondenser-ensembledistil",
        "msmarco-splade-pp-sd": "naver/splade-cocondenser-selfdistil",
    }

    def __init__(
        self,
        index_type: str = "splade-pp-ed",
        index_folder: Optional[str] = None,
        corpus_path: Optional[str] = None,
        model_name: Optional[str] = None,
        device: str = "cpu",
        **kwargs,
    ):
        super().__init__(**kwargs)
        if not PYSERINI_AVAILABLE:
            raise ImportError(
                "Pyserini is required for SpladeV2Retriever. "
                "Install with: pip install pyserini"
            )
        self.index_type = index_type
        self.index_folder = index_folder
        self.corpus_path = corpus_path
        self.device = device
        # Choose a sensible default model for the requested index variant
        self.model_name = model_name or self.DEFAULT_MODEL_MAP.get(
            index_type, "naver/splade-cocondenser-ensembledistil"
        )

        self.passage_store: dict = {}
        if corpus_path:
            self._load_corpus(corpus_path)

        self.searcher = self._initialize_searcher()

    def _load_corpus(self, corpus_path: str) -> None:
        """Load passage texts from a JSONL file into a dict keyed by id."""
        print(f"Loading corpus from {corpus_path}...")
        with open(corpus_path, "r", encoding="utf-8") as fh:
            for line in fh:
                if not line.strip():
                    continue
                obj = json.loads(line)
                pid = str(obj.get("id", obj.get("docid", "")))
                self.passage_store[pid] = {
                    "text": obj.get("text", obj.get("contents", "")),
                    "title": obj.get("title", ""),
                }
        print(f"Loaded {len(self.passage_store)} passages.")

    def _initialize_searcher(self) -> "LuceneImpactSearcher":
        """Build and return a LuceneImpactSearcher with a SPLADE query encoder."""
        query_encoder = SpladeQueryEncoder(self.model_name, device=self.device)
        if self.index_folder:
            return LuceneImpactSearcher(self.index_folder, query_encoder)
        prebuilt = self.PREBUILT_INDEX_MAP.get(self.index_type)
        if prebuilt is None:
            raise ValueError(
                f"Unknown index_type '{self.index_type}'. "
                f"Supported values: {list(self.PREBUILT_INDEX_MAP.keys())}, "
                "or pass index_folder."
            )
        return LuceneImpactSearcher.from_prebuilt_index(prebuilt, query_encoder)

    def _get_passage(self, docid: str):
        """Return (text, title) for *docid*, falling back to Lucene stored field."""
        passage = self.passage_store.get(docid)
        if passage is not None:
            return passage["text"], passage["title"]
        try:
            raw = self.searcher.doc(docid)
            if raw:
                parsed = json.loads(raw.raw())
                text = parsed.get("contents", parsed.get("text", ""))
                title = parsed.get("title", "")
                return text, title
        except Exception:
            pass
        return "Text not available", ""

    def retrieve(self, documents: List[Document]) -> List[Document]:
        """
        Retrieve passages for each document using the SPLADE sparse index.

        Args:
            documents (List[Document]): Documents whose ``question`` field
                contains the query.

        Returns:
            List[Document]: Documents with populated ``contexts``.
        """
        for document in tqdm(documents, desc="SPLADE-v2 retrieving"):
            query = document.question.question
            hits = self.searcher.search(query, k=self.n_docs)
            contexts: List[Context] = []
            for hit in hits:
                docid = str(hit.docid)
                score = float(hit.score)
                text, title = self._get_passage(docid)

                has_ans = False
                if document.answers and document.answers.answers:
                    try:
                        from pyserini.eval.evaluate_dpr_retrieval import (
                            has_answers,
                            SimpleTokenizer,
                        )
                        has_ans = has_answers(
                            text, document.answers.answers, SimpleTokenizer()
                        )
                    except Exception:
                        pass

                contexts.append(
                    Context(
                        id=docid,
                        title=title,
                        text=text,
                        score=score,
                        has_answer=has_ans,
                    )
                )
            document.contexts = contexts
        return documents

retrieve(documents)

Retrieve passages for each document using the SPLADE sparse index.

Parameters:

Name	Type	Description	Default
documents	List[Document]	Documents whose question field contains the query.	required

Returns:

Type	Description
List[Document]	List[Document]: Documents with populated contexts.

Source code in rankify/retrievers/splade_v2_retriever.py

def retrieve(self, documents: List[Document]) -> List[Document]:
    """
    Retrieve passages for each document using the SPLADE sparse index.

    Args:
        documents (List[Document]): Documents whose ``question`` field
            contains the query.

    Returns:
        List[Document]: Documents with populated ``contexts``.
    """
    for document in tqdm(documents, desc="SPLADE-v2 retrieving"):
        query = document.question.question
        hits = self.searcher.search(query, k=self.n_docs)
        contexts: List[Context] = []
        for hit in hits:
            docid = str(hit.docid)
            score = float(hit.score)
            text, title = self._get_passage(docid)

            has_ans = False
            if document.answers and document.answers.answers:
                try:
                    from pyserini.eval.evaluate_dpr_retrieval import (
                        has_answers,
                        SimpleTokenizer,
                    )
                    has_ans = has_answers(
                        text, document.answers.answers, SimpleTokenizer()
                    )
                except Exception:
                    pass

            contexts.append(
                Context(
                    id=docid,
                    title=title,
                    text=text,
                    score=score,
                    has_answer=has_ans,
                )
            )
        document.contexts = contexts
    return documents

APIEmbeddingRetriever

Bases: BaseRetriever

Dense retriever backed by an external embedding API and a FAISS index.

Supported providers:

• 'openai' – uses openai.OpenAI client
• 'cohere' – uses cohere.Client
• 'voyage' – uses voyageai.Client

The corpus is embedded once and cached as a .npy file under cache_dir. Subsequent runs load the cache automatically.

Parameters:

Name	Type	Description	Default
provider	str	One of 'openai', 'cohere', or 'voyage'.	required
api_key	str	API key for the chosen provider.	required
corpus_path	str	Path to a JSONL corpus file. Each line must contain id, text, and optionally title.	required
model_name	str	Embedding model name. Defaults are text-embedding-3-small (OpenAI), embed-english-v3.0 (Cohere), and voyage-3 (Voyage).	None
cache_dir	str	Directory for cached embeddings.	'./cache'
embed_batch_size	int	Override the default API batch size.	None

Example

from rankify.retrievers.retriever import Retriever

retriever = Retriever(
    method="openai-embedding",
    n_docs=10,
    provider="openai",
    api_key="sk-...",
    corpus_path="corpus.jsonl",
)

Source code in rankify/retrievers/api_embedding_retriever.py

class APIEmbeddingRetriever(BaseRetriever):
    """
    Dense retriever backed by an external embedding API and a FAISS index.

    Supported providers:

    * ``'openai'``  – uses ``openai.OpenAI`` client
    * ``'cohere'``  – uses ``cohere.Client``
    * ``'voyage'``  – uses ``voyageai.Client``

    The corpus is embedded once and cached as a ``.npy`` file under
    ``cache_dir``.  Subsequent runs load the cache automatically.

    Args:
        provider (str): One of ``'openai'``, ``'cohere'``, or ``'voyage'``.
        api_key (str): API key for the chosen provider.
        corpus_path (str): Path to a JSONL corpus file.  Each line must
            contain ``id``, ``text``, and optionally ``title``.
        model_name (str, optional): Embedding model name.  Defaults are
            ``text-embedding-3-small`` (OpenAI), ``embed-english-v3.0``
            (Cohere), and ``voyage-3`` (Voyage).
        cache_dir (str): Directory for cached embeddings.
        embed_batch_size (int, optional): Override the default API batch size.

    Example:
        ```python
        from rankify.retrievers.retriever import Retriever

        retriever = Retriever(
            method="openai-embedding",
            n_docs=10,
            provider="openai",
            api_key="sk-...",
            corpus_path="corpus.jsonl",
        )
        ```
    """

    def __init__(
        self,
        provider: str,
        api_key: str,
        corpus_path: str,
        model_name: Optional[str] = None,
        cache_dir: str = "./cache",
        embed_batch_size: Optional[int] = None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        if not FAISS_AVAILABLE:
            raise ImportError(
                "faiss is required for APIEmbeddingRetriever. "
                "Install with: pip install faiss-cpu  (or faiss-gpu)"
            )

        self.provider = provider.lower()
        self.api_key = api_key
        self.corpus_path = corpus_path
        self.model_name = model_name or _DEFAULT_MODELS.get(self.provider, "")
        self.cache_dir = cache_dir
        self.embed_batch_size = embed_batch_size or _BATCH_SIZES.get(self.provider, 128)

        self._validate_provider()

        # Load corpus texts
        self.doc_ids, self.doc_texts, self.doc_titles = self._load_corpus()
        self._docid_to_idx = {did: i for i, did in enumerate(self.doc_ids)}

        # Build or load FAISS index
        self.doc_emb = self._load_or_build_embeddings()
        self.index = self._build_faiss_index(self.doc_emb)

    # ------------------------------------------------------------------
    # Validation
    # ------------------------------------------------------------------

    def _validate_provider(self) -> None:
        if self.provider == "openai" and not OPENAI_AVAILABLE:
            raise ImportError(
                "openai package is required. Install with: pip install openai"
            )
        if self.provider == "cohere" and not COHERE_AVAILABLE:
            raise ImportError(
                "cohere package is required. Install with: pip install cohere"
            )
        if self.provider == "voyage" and not VOYAGE_AVAILABLE:
            raise ImportError(
                "voyageai package is required. Install with: pip install voyageai"
            )
        if self.provider not in _DEFAULT_MODELS:
            raise ValueError(
                f"Unknown provider '{self.provider}'. "
                f"Supported: {list(_DEFAULT_MODELS.keys())}"
            )

    # ------------------------------------------------------------------
    # Corpus loading
    # ------------------------------------------------------------------

    def _load_corpus(self):
        if not os.path.exists(self.corpus_path):
            raise FileNotFoundError(f"Corpus not found: {self.corpus_path}")
        doc_ids, doc_texts, doc_titles = [], [], []
        with open(self.corpus_path, "r", encoding="utf-8") as fh:
            for line in fh:
                if not line.strip():
                    continue
                obj = json.loads(line)
                doc_ids.append(str(obj.get("id", obj.get("docid", len(doc_ids)))))
                doc_texts.append(obj.get("text", obj.get("contents", "")))
                doc_titles.append(obj.get("title", ""))
        return doc_ids, doc_texts, doc_titles

    # ------------------------------------------------------------------
    # Embedding helpers
    # ------------------------------------------------------------------

    def _embed_openai(self, texts: List[str], input_type: str = "document") -> np.ndarray:
        client = _openai_module.OpenAI(api_key=self.api_key)
        all_emb = []
        for i in tqdm(range(0, len(texts), self.embed_batch_size), desc="OpenAI embed"):
            batch = texts[i : i + self.embed_batch_size]
            response = client.embeddings.create(input=batch, model=self.model_name)
            batch_emb = [item.embedding for item in response.data]
            all_emb.extend(batch_emb)
        return np.array(all_emb, dtype=np.float32)

    def _embed_cohere(self, texts: List[str], input_type: str = "search_document") -> np.ndarray:
        client = _cohere_module.Client(self.api_key)
        all_emb = []
        for i in tqdm(range(0, len(texts), self.embed_batch_size), desc="Cohere embed"):
            batch = texts[i : i + self.embed_batch_size]
            response = client.embed(
                texts=batch, model=self.model_name, input_type=input_type
            )
            all_emb.extend(response.embeddings)
        return np.array(all_emb, dtype=np.float32)

    def _embed_voyage(self, texts: List[str], input_type: str = "document") -> np.ndarray:
        client = _voyage_module.Client(api_key=self.api_key)
        all_emb = []
        for i in tqdm(range(0, len(texts), self.embed_batch_size), desc="Voyage embed"):
            batch = texts[i : i + self.embed_batch_size]
            response = client.embed(batch, model=self.model_name, input_type=input_type)
            all_emb.extend(response.embeddings)
        return np.array(all_emb, dtype=np.float32)

    def _embed(self, texts: List[str], input_type: str = "document") -> np.ndarray:
        if self.provider == "openai":
            return self._embed_openai(texts, input_type)
        if self.provider == "cohere":
            cohere_type = "search_document" if input_type == "document" else "search_query"
            return self._embed_cohere(texts, cohere_type)
        if self.provider == "voyage":
            return self._embed_voyage(texts, input_type)
        raise ValueError(f"Unknown provider '{self.provider}'")

    # ------------------------------------------------------------------
    # Cache management
    # ------------------------------------------------------------------

    def _cache_path(self) -> str:
        corpus_id = os.path.splitext(os.path.basename(self.corpus_path))[0]
        folder = os.path.join(
            self.cache_dir, "api_emb", self.provider, self.model_name.replace("/", "_"), corpus_id
        )
        os.makedirs(folder, exist_ok=True)
        return os.path.join(folder, "doc_emb.npy")

    def _load_or_build_embeddings(self) -> np.ndarray:
        cache = self._cache_path()
        if os.path.isfile(cache):
            print(f"Loading cached embeddings from {cache}")
            return np.load(cache)
        print(f"Embedding {len(self.doc_texts)} documents via {self.provider} ({self.model_name})…")
        emb = self._embed(self.doc_texts, input_type="document")
        np.save(cache, emb)
        print(f"Saved embeddings to {cache}")
        return emb

    # ------------------------------------------------------------------
    # FAISS index
    # ------------------------------------------------------------------

    def _build_faiss_index(self, emb: np.ndarray) -> "faiss.Index":
        dim = emb.shape[1]
        # L2-normalize for cosine similarity via inner product
        norms = np.linalg.norm(emb, axis=1, keepdims=True)
        norms = np.where(norms == 0, 1.0, norms)
        emb_norm = (emb / norms).astype(np.float32)
        index = faiss.IndexFlatIP(dim)
        index.add(emb_norm)
        return index

    # ------------------------------------------------------------------
    # BaseRetriever interface
    # ------------------------------------------------------------------

    def _initialize_searcher(self):
        return None  # FAISS index is built in __init__

    def retrieve(self, documents: List[Document]) -> List[Document]:
        """
        Retrieve passages for each document by embedding the query and
        searching the FAISS index.

        Args:
            documents (List[Document]): Documents with query in ``question``.

        Returns:
            List[Document]: Documents with populated ``contexts``.
        """
        queries = [d.question.question for d in documents]
        query_emb = self._embed(queries, input_type="query")

        # L2-normalize for inner-product cosine search
        norms = np.linalg.norm(query_emb, axis=1, keepdims=True)
        norms = np.where(norms == 0, 1.0, norms)
        query_emb_norm = (query_emb / norms).astype(np.float32)

        scores_matrix, indices_matrix = self.index.search(query_emb_norm, self.n_docs)

        for document, scores_row, idx_row in zip(documents, scores_matrix, indices_matrix):
            contexts: List[Context] = []
            for score, idx in zip(scores_row, idx_row):
                if idx < 0 or idx >= len(self.doc_ids):
                    continue
                doc_id = self.doc_ids[idx]
                text = self.doc_texts[idx]
                title = self.doc_titles[idx]

                has_ans = False
                if document.answers and document.answers.answers:
                    try:
                        from pyserini.eval.evaluate_dpr_retrieval import (
                            has_answers,
                            SimpleTokenizer,
                        )
                        has_ans = has_answers(
                            text, document.answers.answers, SimpleTokenizer()
                        )
                    except Exception:
                        pass

                contexts.append(
                    Context(
                        id=doc_id,
                        title=title,
                        text=text,
                        score=float(score),
                        has_answer=has_ans,
                    )
                )
            document.contexts = contexts
        return documents

retrieve(documents)

Retrieve passages for each document by embedding the query and searching the FAISS index.

Parameters:

Name	Type	Description	Default
documents	List[Document]	Documents with query in question.	required

Returns:

Type	Description
List[Document]	List[Document]: Documents with populated contexts.

Source code in rankify/retrievers/api_embedding_retriever.py

def retrieve(self, documents: List[Document]) -> List[Document]:
    """
    Retrieve passages for each document by embedding the query and
    searching the FAISS index.

    Args:
        documents (List[Document]): Documents with query in ``question``.

    Returns:
        List[Document]: Documents with populated ``contexts``.
    """
    queries = [d.question.question for d in documents]
    query_emb = self._embed(queries, input_type="query")

    # L2-normalize for inner-product cosine search
    norms = np.linalg.norm(query_emb, axis=1, keepdims=True)
    norms = np.where(norms == 0, 1.0, norms)
    query_emb_norm = (query_emb / norms).astype(np.float32)

    scores_matrix, indices_matrix = self.index.search(query_emb_norm, self.n_docs)

    for document, scores_row, idx_row in zip(documents, scores_matrix, indices_matrix):
        contexts: List[Context] = []
        for score, idx in zip(scores_row, idx_row):
            if idx < 0 or idx >= len(self.doc_ids):
                continue
            doc_id = self.doc_ids[idx]
            text = self.doc_texts[idx]
            title = self.doc_titles[idx]

            has_ans = False
            if document.answers and document.answers.answers:
                try:
                    from pyserini.eval.evaluate_dpr_retrieval import (
                        has_answers,
                        SimpleTokenizer,
                    )
                    has_ans = has_answers(
                        text, document.answers.answers, SimpleTokenizer()
                    )
                except Exception:
                    pass

            contexts.append(
                Context(
                    id=doc_id,
                    title=title,
                    text=text,
                    score=float(score),
                    has_answer=has_ans,
                )
            )
        document.contexts = contexts
    return documents

Retriever

Unified retriever interface for the rankify framework.

Provides a simple interface to access different retrieval methods (BM25, DPR, ANCE, BPR) with consistent parameters.

Example

# Initialize with BM25
retriever = Retriever(method="bm25", n_docs=10, index_type="wiki")

# Initialize with DPR
retriever = Retriever(method="dpr-multi", n_docs=5, index_type="msmarco")

# Initialize with ANCE (UPDATED - now works with index_type)
retriever = Retriever(method="ance", n_docs=10, index_type="wiki")

# Initialize with ANCE-Multi (uses prebuilt Wikipedia indices)
retriever = Retriever(method="ance-multi", n_docs=10, index_type="wiki")

# Initialize with custom index folder (works for all methods)
retriever = Retriever(method="ance", n_docs=10, index_folder="/path/to/index")

# Retrieve documents
retrieved_documents = retriever.retrieve(documents)

Source code in rankify/retrievers/retriever.py

class Retriever:
    """
    Unified retriever interface for the rankify framework.

    Provides a simple interface to access different retrieval methods
    (BM25, DPR, ANCE, BPR) with consistent parameters.

    Example:
        ```python
        # Initialize with BM25
        retriever = Retriever(method="bm25", n_docs=10, index_type="wiki")

        # Initialize with DPR
        retriever = Retriever(method="dpr-multi", n_docs=5, index_type="msmarco")

        # Initialize with ANCE (UPDATED - now works with index_type)
        retriever = Retriever(method="ance", n_docs=10, index_type="wiki")

        # Initialize with ANCE-Multi (uses prebuilt Wikipedia indices)
        retriever = Retriever(method="ance-multi", n_docs=10, index_type="wiki")

        # Initialize with custom index folder (works for all methods)
        retriever = Retriever(method="ance", n_docs=10, index_folder="/path/to/index")

        # Retrieve documents
        retrieved_documents = retriever.retrieve(documents)
        ```
    """

    def __init__(self, method: str, n_docs: int = 10, index_type: str = "wiki", 
                 index_folder: str = None, encoder_name: str = None, **kwargs):
        """
        Initialize the retriever.

        Args:
            method (str): Retrieval method ('bm25', 'dpr-multi', 'dpr-single', 'ance', 'ance-multi', 'bpr-single', etc.)
            n_docs (int): Number of documents to retrieve per query
            index_type (str): Index type ('wiki', 'msmarco') - ignored if index_folder is provided
            index_folder (str): Path to custom index folder (optional)
            encoder_name (str): Model name for encoding (method-specific)
            **kwargs: Additional parameters passed to the specific retriever
        """
        self.method = method.lower()
        self.n_docs = n_docs
        self.index_type = index_type.lower()
        self.index_folder = index_folder
        self.encoder_name = encoder_name
        self.kwargs = kwargs

        # Initialize the specific retriever
        self.retriever = self._initialize_retriever()

    def _initialize_retriever(self) -> BaseRetriever:
        """Initialize the specific retriever based on the method."""
        if self.method not in METHOD_MAP:
            supported_methods = ", ".join(METHOD_MAP.keys())
            raise ValueError(f"Unsupported method '{self.method}'. "
                           f"Supported methods: {supported_methods}")

        retriever_class = METHOD_MAP[self.method]

        # Prepare initialization parameters
        init_params = {
            "n_docs": self.n_docs,
            **self.kwargs
        }

        # UPDATED: Handles no-index retrievers
        NO_INDEX_METHODS = {
            "diver-dense",
            "diver-bm25",
            "reasonir",
            "reason-embed",
            "bge-reasoner-embed",
            "unicoil",
            "unicoil-noexp",
            "splade-v2",
            "openai-embedding",
            "cohere-embedding",
            "voyage-embedding",
        }
        # No index retrievers do NOT use index_type or index_folder
            # They require corpus_path + model_id via kwargs
        if self.method in NO_INDEX_METHODS:
            return retriever_class(**init_params)

        # UPDATED: Handle all ANCE variants the same way (with index_type support)
        if self.method in ["ance", "ance-msmarco", "ance-multi"]:
            init_params["index_type"] = self.index_type

            # Add index_folder if provided
            if self.index_folder:
                init_params["index_folder"] = self.index_folder

            # Add encoder_name if provided
            if self.encoder_name:
                init_params["encoder_name"] = self.encoder_name

        # Handle other retrieval methods (same as before)
        else:
            init_params["index_type"] = self.index_type

            # Add index_folder if provided
            if self.index_folder:
                init_params["index_folder"] = self.index_folder

            # Add method parameter for dense retrievers
            if self.method in ["dpr-multi", "dpr-single", "bpr-single"]:
                init_params["method"] = self.method

        return retriever_class(**init_params)

    def retrieve(self, documents: List[Document]) -> List[Document]:
        """
        Retrieve relevant contexts for the given documents.

        Args:
            documents (List[Document]): List of documents containing queries

        Returns:
            List[Document]: Documents updated with retrieved contexts
        """
        return self.retriever.retrieve(documents)

    @classmethod
    def supported_methods(cls) -> List[str]:
        """Get list of supported retrieval methods."""
        return list(METHOD_MAP.keys())

    def __repr__(self) -> str:
        index_info = f"index_folder='{self.index_folder}'" if self.index_folder else f"index_type='{self.index_type}'"
        return (f"Retriever(method='{self.method}', n_docs={self.n_docs}, "
                f"{index_info})")

__init__(method, n_docs=10, index_type='wiki', index_folder=None, encoder_name=None, **kwargs)

Initialize the retriever.

Parameters:

Name	Type	Description	Default
method	str	Retrieval method ('bm25', 'dpr-multi', 'dpr-single', 'ance', 'ance-multi', 'bpr-single', etc.)	required
n_docs	int	Number of documents to retrieve per query	10
index_type	str	Index type ('wiki', 'msmarco') - ignored if index_folder is provided	'wiki'
index_folder	str	Path to custom index folder (optional)	None
encoder_name	str	Model name for encoding (method-specific)	None
**kwargs		Additional parameters passed to the specific retriever	{}

Source code in rankify/retrievers/retriever.py

def __init__(self, method: str, n_docs: int = 10, index_type: str = "wiki", 
             index_folder: str = None, encoder_name: str = None, **kwargs):
    """
    Initialize the retriever.

    Args:
        method (str): Retrieval method ('bm25', 'dpr-multi', 'dpr-single', 'ance', 'ance-multi', 'bpr-single', etc.)
        n_docs (int): Number of documents to retrieve per query
        index_type (str): Index type ('wiki', 'msmarco') - ignored if index_folder is provided
        index_folder (str): Path to custom index folder (optional)
        encoder_name (str): Model name for encoding (method-specific)
        **kwargs: Additional parameters passed to the specific retriever
    """
    self.method = method.lower()
    self.n_docs = n_docs
    self.index_type = index_type.lower()
    self.index_folder = index_folder
    self.encoder_name = encoder_name
    self.kwargs = kwargs

    # Initialize the specific retriever
    self.retriever = self._initialize_retriever()

retrieve(documents)

Retrieve relevant contexts for the given documents.

Parameters:

Name	Type	Description	Default
documents	List[Document]	List of documents containing queries	required

Returns:

Type	Description
List[Document]	List[Document]: Documents updated with retrieved contexts

Source code in rankify/retrievers/retriever.py

def retrieve(self, documents: List[Document]) -> List[Document]:
    """
    Retrieve relevant contexts for the given documents.

    Args:
        documents (List[Document]): List of documents containing queries

    Returns:
        List[Document]: Documents updated with retrieved contexts
    """
    return self.retriever.retrieve(documents)

supported_methods() classmethod

Get list of supported retrieval methods.

Source code in rankify/retrievers/retriever.py

@classmethod
def supported_methods(cls) -> List[str]:
    """Get list of supported retrieval methods."""
    return list(METHOD_MAP.keys())