HyDE Retriever

rankify.retrievers.hyde_retriever

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

IndexManager

Manages downloading, caching, and loading of retrieval indexes.

Handles both prebuilt indexes (wiki, msmarco) and custom user indexes.

Source code in rankify/retrievers/index_manager.py

class IndexManager:
    """
    Manages downloading, caching, and loading of retrieval indexes.

    Handles both prebuilt indexes (wiki, msmarco) and custom user indexes.
    """

    def __init__(self, cache_dir: str = None):
        self.cache_dir = cache_dir or os.environ.get("RERANKING_CACHE_DIR", "./cache")
        self.index_configs = self._load_index_configs()

    def _load_index_configs(self) -> Dict:
        """Load index configuration mappings."""
        return {
            "bm25": {
                "wiki": {
                    "url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/bm25_wiki.zip",
                    "prebuilt": "wikipedia-dpr-100w"
                },
                "msmarco": {
                    "url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/bm25_index_msmarco.zip?download=true",
                    "prebuilt": None
                }
            },
            "dpr-multi": {
                "wiki": {
                    "prebuilt": "wikipedia-dpr-100w.dpr-multi",
                    "encoder": "facebook/dpr-question_encoder-multiset-base"
                },
                "msmarco": {
                    "url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/msmarco_passage_dpr_multi.zip",
                    "encoder": "facebook/dpr-question_encoder-multiset-base"
                }
            },
            "dpr-single": {
                "wiki": {
                    "prebuilt": "wikipedia-dpr-100w.dpr-single-nq",
                    "encoder": "facebook/dpr-question_encoder-single-nq-base"
                },
                "msmarco": {
                    "url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/msmarco_passage_dpr_single.zip",
                    "encoder": "facebook/dpr-question_encoder-single-nq-base"
                }
            },
            "ance-multi": {
                "wiki": {
                    "prebuilt": "wikipedia-dpr-100w.ance-multi",
                    "encoder": "castorini/ance-dpr-question-multi"
                },
                "msmarco": {
                    "prebuilt": "msmarco-v1-passage.ance", 
                    "encoder": "castorini/ance-msmarco-passage"
                }
            },

            "bpr-single": {
                "wiki": {
                    "prebuilt": "wikipedia-dpr-100w.bpr-single-nq",
                    "encoder": "castorini/bpr-nq-question-encoder"
                },
                "msmarco": {
                    "url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/msmarco_passage_bpr.zip",
                    "encoder": "castorini/bpr-nq-question-encoder"
                }
            },
            "bge": {
                "wiki": {
                    "urls": [
                        "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/bgb_index.tar.gz.part1?download=true",
                        "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/bgb_index.tar.gz.part2?download=true",
                    ],
                    "passages_url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/psgs_w100/psgs_w100.tsv?download=true",
                    "encoder": "BAAI/bge-large-en-v1.5"
                },
                "msmarco": {
                    "urls": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/msmarco_embeddings_bgb.zip?download=true",
                    "passages_url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/msmarco-passage-corpus/msmarco-passage-corpus.tsv?download=true",
                    "encoder": "BAAI/bge-large-en-v1.5"
                }
            },
            "colbert": {
                "wiki": {
                    "urls": [
                        "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/wikipedia_embeddings_colbert/wiki.zip.001?download=true",
                        "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/wikipedia_embeddings_colbert/wiki.zip.002?download=true",
                        "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/wikipedia_embeddings_colbert/wiki.zip.003?download=true",
                        "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/wikipedia_embeddings_colbert/wiki.zip.004?download=true",
                        "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/wikipedia_embeddings_colbert/wiki.zip.005?download=true"
                    ],
                    "passages_url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/psgs_w100/psgs_w100.tsv?download=true",
                    "model": "colbert-ir/colbertv2.0"
                },
                "msmarco": {
                    "urls": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/msmarco_embeddings_colbert.zip?download=true",
                    "passages_url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/msmarco-passage-corpus/msmarco-passage-corpus.tsv?download=true",
                    "model": "colbert-ir/colbertv2.0"
                }
            },
            "contriever": {
                "wiki": {
                    "url": "https://dl.fbaipublicfiles.com/contriever/embeddings/contriever-msmarco/wikipedia_embeddings.tar",
                    "passages_url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/psgs_w100/psgs_w100.tsv?download=true",
                    "model": "facebook/contriever-msmarco",
                    "vector_size": 768
                },
                "msmarco": {
                    "url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/index/msmarco_embeddings_contriever.zip?download=true",
                    "passages_url": "https://huggingface.co/datasets/abdoelsayed/reranking-datasets/resolve/main/msmarco-passage-corpus/msmarco-passage-corpus.tsv?download=true",
                    "model": "facebook/contriever-msmarco",
                    "vector_size": 768
                }
            },
            "online": {
                "web": {
                    "search_provider": "web",
                    "chunk_size": 500,
                    "chunk_overlap": 50,
                    "requires_api_key": True
                }
            },
            "hyde": {
                "wiki": {
                    "base_retriever": "contriever",
                    "base_model": "facebook/contriever-msmarco",
                    "base_index_type": "wiki",
                    "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
                },
                "msmarco": {
                    "base_retriever": "contriever", 
                    "base_model": "facebook/contriever-msmarco",
                    "base_index_type": "msmarco",
                    "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 get_index_path(self, method: str, index_type: str, custom_path: str = None) -> str:
        """
        Get the path to the index for a given method and index type.

        Args:
            method (str): Retrieval method (e.g., 'bm25', 'dpr-multi', 'ance')
            index_type (str): Index type ('wiki', 'msmarco', or 'custom')
            custom_path (str): Path to custom index if index_type is 'custom'

        Returns:
            str: Path to the index
        """
        if custom_path:
            return custom_path

        if method not in self.index_configs:
            raise ValueError(f"Unsupported method: {method}")

        if index_type not in self.index_configs[method]:
            raise ValueError(f"Unsupported index type '{index_type}' for method '{method}'")

        config = self.index_configs[method][index_type]

        # If it's a prebuilt index, return the identifier
        if "prebuilt" in config and config["prebuilt"]:
            return config["prebuilt"]

        # Otherwise, download and return local path
        return self._ensure_index_downloaded(method, index_type)

    def _ensure_index_downloaded(self, method: str, index_type: str) -> str:
        """Download and extract index if not already available."""
        config = self.index_configs[method][index_type]
        url = config["url"]

        # Create local directory path
        index_name = f"{method}_{index_type}"
        local_dir = os.path.join(self.cache_dir, "index", index_name)

        if not os.path.exists(local_dir):
            print(f"Downloading {method} index for {index_type}...")
            self._download_and_extract(url, local_dir)

        return local_dir

    def _download_and_extract(self, url: str, destination: str):
        """Download and extract a ZIP file."""
        os.makedirs(destination, exist_ok=True)

        zip_name = os.path.basename(url).split("?")[0]
        zip_path = os.path.join(self.cache_dir, "temp", zip_name)

        os.makedirs(os.path.dirname(zip_path), exist_ok=True)

        if not os.path.exists(zip_path):
            response = requests.get(url, stream=True)
            response.raise_for_status()

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

        print(f"Extracting {zip_name}...")
        with zipfile.ZipFile(zip_path, "r") as zip_ref:
            zip_ref.extractall(destination)

        # Clean up ZIP file
        os.remove(zip_path)
        print("Extraction complete.")

    def load_id_mapping(self, index_path: str) -> Optional[Dict]:
        """Load ID mapping if available."""
        mapping_path = os.path.join(index_path, "id_mapping.json")
        if os.path.exists(mapping_path):
            with open(mapping_path, "r", encoding="utf-8") as f:
                mapping = json.load(f)
                return {v: k for k, v in mapping.items()}  # Reverse mapping
        return None

    # def load_corpus(self, index_path: str) -> Dict:
    #     """Load corpus data from index path."""
    #     corpus_file = os.path.join(index_path, "corpus.jsonl")
    #     if not os.path.exists(corpus_file):
    #         return {}

    #     corpus = {}
    #     with open(corpus_file, "r", encoding="utf-8") as f:
    #         for line in f:
    #             doc = json.loads(line.strip())
    #             doc_id = doc.get("docid") or doc.get("id")
    #             contents = doc.get("contents", "")
    #             title = doc.get("title", contents[:100] if contents else "No Title")
    #             corpus[str(doc_id)] = {"contents": contents, "title": title}

    #     return corpus
    def load_corpus(self, index_path: str) -> Dict:
        """Load corpus data from index path (supports multiple formats)."""
        # Try different corpus file formats in order of preference
        corpus_files = [
            "corpus_metadata.json",  # Your custom format
            "corpus.jsonl",          # Standard JSONL format
            "corpus.json"            # Alternative JSON format
        ]

        for corpus_filename in corpus_files:
            corpus_file = os.path.join(index_path, corpus_filename)
            if os.path.exists(corpus_file):
                print(f"📚 Loading corpus from {corpus_filename}")

                if corpus_filename.endswith('.jsonl'):
                    # JSONL format (one JSON object per line)
                    return self._load_corpus_jsonl(corpus_file)
                else:
                    # JSON format (single JSON object)
                    return self._load_corpus_json(corpus_file)

        print(f"❌ No corpus file found in {index_path}")
        return {}

    def _load_corpus_jsonl(self, corpus_file: str) -> Dict:
        """Load corpus from JSONL format (one JSON object per line)."""
        corpus = {}
        with open(corpus_file, "r", encoding="utf-8") as f:
            for line in f:
                doc = json.loads(line.strip())
                doc_id = doc.get("docid") or doc.get("id")
                contents = doc.get("contents", "")
                title = doc.get("title", contents[:100] if contents else "No Title")
                corpus[str(doc_id)] = {"contents": contents, "title": title}
        return corpus

    def _load_corpus_json(self, corpus_file: str) -> Dict:
        """Load corpus from JSON format (single JSON object)."""
        corpus = {}
        with open(corpus_file, "r", encoding="utf-8") as f:
            data = json.load(f)

        # Handle different JSON structures
        if isinstance(data, dict):
            # If it's a dict, iterate through it
            for doc_id, doc_data in data.items():
                if isinstance(doc_data, dict):
                    contents = doc_data.get("contents") or doc_data.get("text") or ""
                    title = doc_data.get("title", contents[:100] if contents else "No Title")
                else:
                    # If doc_data is a string, use it as contents
                    contents = str(doc_data)
                    title = contents[:100] if contents else "No Title"
                corpus[str(doc_id)] = {"contents": contents, "title": title}

        elif isinstance(data, list):
            # If it's a list, iterate through documents
            for doc in data:
                doc_id = doc.get("docid") or doc.get("id")
                contents = doc.get("contents", "")
                title = doc.get("title", contents[:100] if contents else "No Title")
                corpus[str(doc_id)] = {"contents": contents, "title": title}

        return corpus
    def download_and_extract_index(self, url: str) -> str:
        """Download and extract index from URL, return local path."""
        # This is used by DenseRetriever pattern
        import tempfile
        import zipfile

        # Create temporary file
        with tempfile.NamedTemporaryFile(suffix='.zip', delete=False) as tmp_file:
            tmp_path = tmp_file.name

        try:
            # Download
            print(f"Downloading from {url}...")
            response = requests.get(url, stream=True)
            response.raise_for_status()

            with open(tmp_path, 'wb') as f:
                for chunk in tqdm(response.iter_content(chunk_size=1024)):
                    f.write(chunk)

            # Extract to cache directory
            extract_dir = os.path.join(self.cache_dir, "downloaded_index")
            os.makedirs(extract_dir, exist_ok=True)

            with zipfile.ZipFile(tmp_path, 'r') as zip_ref:
                zip_ref.extractall(extract_dir)

            return extract_dir

        finally:
            # Clean up temporary file
            if os.path.exists(tmp_path):
                os.remove(tmp_path)

get_index_path(method, index_type, custom_path=None)

Get the path to the index for a given method and index type.

Parameters:

Name	Type	Description	Default
method	str	Retrieval method (e.g., 'bm25', 'dpr-multi', 'ance')	required
index_type	str	Index type ('wiki', 'msmarco', or 'custom')	required
custom_path	str	Path to custom index if index_type is 'custom'	None

Returns:

Name	Type	Description
str	str	Path to the index

Source code in rankify/retrievers/index_manager.py

def get_index_path(self, method: str, index_type: str, custom_path: str = None) -> str:
    """
    Get the path to the index for a given method and index type.

    Args:
        method (str): Retrieval method (e.g., 'bm25', 'dpr-multi', 'ance')
        index_type (str): Index type ('wiki', 'msmarco', or 'custom')
        custom_path (str): Path to custom index if index_type is 'custom'

    Returns:
        str: Path to the index
    """
    if custom_path:
        return custom_path

    if method not in self.index_configs:
        raise ValueError(f"Unsupported method: {method}")

    if index_type not in self.index_configs[method]:
        raise ValueError(f"Unsupported index type '{index_type}' for method '{method}'")

    config = self.index_configs[method][index_type]

    # If it's a prebuilt index, return the identifier
    if "prebuilt" in config and config["prebuilt"]:
        return config["prebuilt"]

    # Otherwise, download and return local path
    return self._ensure_index_downloaded(method, index_type)

load_id_mapping(index_path)

Load ID mapping if available.

Source code in rankify/retrievers/index_manager.py

def load_id_mapping(self, index_path: str) -> Optional[Dict]:
    """Load ID mapping if available."""
    mapping_path = os.path.join(index_path, "id_mapping.json")
    if os.path.exists(mapping_path):
        with open(mapping_path, "r", encoding="utf-8") as f:
            mapping = json.load(f)
            return {v: k for k, v in mapping.items()}  # Reverse mapping
    return None

load_corpus(index_path)

Load corpus data from index path (supports multiple formats).

Source code in rankify/retrievers/index_manager.py

def load_corpus(self, index_path: str) -> Dict:
    """Load corpus data from index path (supports multiple formats)."""
    # Try different corpus file formats in order of preference
    corpus_files = [
        "corpus_metadata.json",  # Your custom format
        "corpus.jsonl",          # Standard JSONL format
        "corpus.json"            # Alternative JSON format
    ]

    for corpus_filename in corpus_files:
        corpus_file = os.path.join(index_path, corpus_filename)
        if os.path.exists(corpus_file):
            print(f"📚 Loading corpus from {corpus_filename}")

            if corpus_filename.endswith('.jsonl'):
                # JSONL format (one JSON object per line)
                return self._load_corpus_jsonl(corpus_file)
            else:
                # JSON format (single JSON object)
                return self._load_corpus_json(corpus_file)

    print(f"❌ No corpus file found in {index_path}")
    return {}

download_and_extract_index(url)

Download and extract index from URL, return local path.

Source code in rankify/retrievers/index_manager.py

def download_and_extract_index(self, url: str) -> str:
    """Download and extract index from URL, return local path."""
    # This is used by DenseRetriever pattern
    import tempfile
    import zipfile

    # Create temporary file
    with tempfile.NamedTemporaryFile(suffix='.zip', delete=False) as tmp_file:
        tmp_path = tmp_file.name

    try:
        # Download
        print(f"Downloading from {url}...")
        response = requests.get(url, stream=True)
        response.raise_for_status()

        with open(tmp_path, 'wb') as f:
            for chunk in tqdm(response.iter_content(chunk_size=1024)):
                f.write(chunk)

        # Extract to cache directory
        extract_dir = os.path.join(self.cache_dir, "downloaded_index")
        os.makedirs(extract_dir, exist_ok=True)

        with zipfile.ZipFile(tmp_path, 'r') as zip_ref:
            zip_ref.extractall(extract_dir)

        return extract_dir

    finally:
        # Clean up temporary file
        if os.path.exists(tmp_path):
            os.remove(tmp_path)

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

Promptor

Source code in rankify/utils/retrievers/hyde/promptor.py

class Promptor:
    def __init__(self, task: str, language: str = 'en'):
        self.task = task
        self.language = language

    def build_prompt(self, query: str):
        if self.task == 'web search':
            return WEB_SEARCH.format(query)
        elif self.task == 'scifact':
            return SCIFACT.format(query)
        elif self.task == 'arguana':
            return ARGUANA.format(query)
        elif self.task == 'trec-covid':
            return TREC_COVID.format(query)
        elif self.task == 'fiqa':
            return FIQA.format(query)
        elif self.task == 'dbpedia-entity':
            return DBPEDIA_ENTITY.format(query)
        elif self.task == 'trec-news':
            return TREC_NEWS.format(query)
        elif self.task == 'mr-tydi':
            return MR_TYDI.format(self.language, query)
        else:
            raise ValueError('Task not supported')

OpenAIGenerator

Bases: Generator

Source code in rankify/utils/retrievers/hyde/generator.py

class OpenAIGenerator(Generator):
    def __init__(self, model_name, api_key, base_url=None, n=8, max_tokens=512, temperature=0.7, top_p=1, frequency_penalty=0.0, presence_penalty=0.0, stop=['\n\n\n'], wait_till_success=False):
        super().__init__(model_name, api_key)
        self.n = n
        self.max_tokens = max_tokens
        self.temperature = temperature
        self.top_p = top_p
        self.frequency_penalty = frequency_penalty
        self.presence_penalty = presence_penalty
        self.stop = stop
        self.wait_till_success = wait_till_success
        self.base_url = base_url
        self._client_init()


    """@staticmethod
    def parse_response(response):
        to_return = []
        for _, g in enumerate(response['choices']):
            text = g['text']
            logprob = sum(g['logprobs']['token_logprobs'])
            to_return.append((text, logprob))
        texts = [r[0] for r in sorted(to_return, key=lambda tup: tup[1], reverse=True)]
        return texts"""
    @staticmethod
    def parse_response(response):
        to_return = []
        for choice in response.choices:
            text = choice.message.content
            if hasattr(choice, 'logprobs') and choice.logprobs and hasattr(choice.logprobs, 'content') and choice.logprobs.content:
                logprob = sum(token_logprob.logprob for token_logprob in choice.logprobs.content)
                #print(f"Logprobs content: {choice.logprobs.content}")  # Should print your list
            else:
                logprob = None
                print("No valid logprobs content found")
            to_return.append((text, logprob))
        to_return.sort(key=lambda x: x[1] if x[1] is not None else float('-inf'), reverse=True)
        texts = [r[0] for r in to_return]
        return texts

    def _client_init(self):
        self.client = openai.OpenAI(
            base_url=self.base_url,
            api_key=self.api_key,
        )
        self.client

    def generate(self, prompt):
        get_results = False
        while not get_results:
            try:
                result = self.client.chat.completions.create(
                    messages=[{"role":"user", "content": prompt}],
                    model=self.model_name,
                    max_completion_tokens=self.max_tokens,
                    temperature=self.temperature,
                    frequency_penalty=self.frequency_penalty,
                    presence_penalty=self.presence_penalty,
                    top_p=self.top_p,
                    n=self.n, # some models only support n=1
                    stop=self.stop,
                    logprobs=True # some models are not compatible with this setting
                )
                get_results = True
            except Exception as e:
                if self.wait_till_success:
                    time.sleep(1)
                else:
                    raise e
        return self.parse_response(result)

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}"

Context

Represents a context with metadata such as score and title.

Attributes:

Name	Type	Description
score	float	The relevance score of the context.
has_answer	bool	Whether the context contains an answer.
id	int	The identifier of the context.
title	str	The title of the context.
text	str	The text of the context.

Source code in rankify/dataset/dataset.py

class Context:
    """
    Represents a context with metadata such as score and title.

    Attributes:
        score (float, optional): The relevance score of the context.
        has_answer (bool, optional): Whether the context contains an answer.
        id (int, optional): The identifier of the context.
        title (str, optional): The title of the context.
        text (str, optional): The text of the context.
    """
    def __init__(self, score: float=None, has_answer: bool=None, id: str=None, title: str=None, text: str=None)-> None:
        """
        Initializes a Context instance.

        Args:
            score (float, optional): The relevance score.
            has_answer (bool, optional): Whether the context contains an answer.
            id (int, optional): The identifier of the context.
            title (str, optional): The title of the context.
            text (str, optional): The text of the context.

        Example:
            ```python
            c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
            print(c)
            ```
        """
        self.score: Optional[float] = score
        self.has_answer: Optional[bool] = has_answer
        self.id: Optional[str] = id
        self.title: Optional[str] = title
        self.text: Optional[str] = text

    def to_dict(self, save_text: bool=False) -> Dict[str, Optional[object]]:

        """
        Converts the Context instance to a dictionary.

        Args:
            save_text (bool): Whether to include text in the output dictionary.

        Returns:
            dict: The context data.

        Example:
            ```python
            c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
            print(c.to_dict())
            ```
        """
        context_dict = {
            "score": float(self.score) if self.score is not None else None,
            "has_answer": self.has_answer,
            "id": self.id,
            }

        # Include 'text' only if save_text is True
        if save_text:
            context_dict["text"] = self.text
            context_dict["title"] =  self.title

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

        Returns:
            str: The formatted context.

        Example:
            ```python
            c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
            print(str(c))
            ```
        """
        return f"ID: {self.id}\nHas Answer: {self.has_answer}\nTitle: {self.title}\nText: {self.text}\nScore: {self.score}"

__init__(score=None, has_answer=None, id=None, title=None, text=None)

Initializes a Context instance.

Parameters:

Name	Type	Description	Default
score	float	The relevance score.	None
has_answer	bool	Whether the context contains an answer.	None
id	int	The identifier of the context.	None
title	str	The title of the context.	None
text	str	The text of the context.	None

Example

c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
print(c)

Source code in rankify/dataset/dataset.py

def __init__(self, score: float=None, has_answer: bool=None, id: str=None, title: str=None, text: str=None)-> None:
    """
    Initializes a Context instance.

    Args:
        score (float, optional): The relevance score.
        has_answer (bool, optional): Whether the context contains an answer.
        id (int, optional): The identifier of the context.
        title (str, optional): The title of the context.
        text (str, optional): The text of the context.

    Example:
        ```python
        c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
        print(c)
        ```
    """
    self.score: Optional[float] = score
    self.has_answer: Optional[bool] = has_answer
    self.id: Optional[str] = id
    self.title: Optional[str] = title
    self.text: Optional[str] = text

to_dict(save_text=False)

Converts the Context instance to a dictionary.

Parameters:

Name	Type	Description	Default
save_text	bool	Whether to include text in the output dictionary.	False

Returns:

Name	Type	Description
dict	Dict[str, Optional[object]]	The context data.

Example

c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
print(c.to_dict())

Source code in rankify/dataset/dataset.py

def to_dict(self, save_text: bool=False) -> Dict[str, Optional[object]]:

    """
    Converts the Context instance to a dictionary.

    Args:
        save_text (bool): Whether to include text in the output dictionary.

    Returns:
        dict: The context data.

    Example:
        ```python
        c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
        print(c.to_dict())
        ```
    """
    context_dict = {
        "score": float(self.score) if self.score is not None else None,
        "has_answer": self.has_answer,
        "id": self.id,
        }

    # Include 'text' only if save_text is True
    if save_text:
        context_dict["text"] = self.text
        context_dict["title"] =  self.title

    return context_dict

__str__()

Returns a string representation of the Context instance.

Returns:

Name	Type	Description
str	str	The formatted context.

Example

c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
print(str(c))

Source code in rankify/dataset/dataset.py

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

    Returns:
        str: The formatted context.

    Example:
        ```python
        c = Context(score=0.9, has_answer=True, id=1, title="Paris", text="The capital of France is Paris.")
        print(str(c))
        ```
    """
    return f"ID: {self.id}\nHas Answer: {self.has_answer}\nTitle: {self.title}\nText: {self.text}\nScore: {self.score}"

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