Indexing Module

The indexing module provides tools for creating custom search indices for various retrieval methods.

Overview

Rankify supports building indices for: - BM25 (Lucene-based sparse retrieval) - DPR (Dense Passage Retrieval) - ANCE (Approximate Nearest Neighbor Negative Contrastive Estimation) - BGE (BAAI General Embedding) - ColBERT (Contextualized Late Interaction over BERT) - Contriever (Contrastive Retriever)

CLI Usage

Use the rankify-index command to build indices:

# BM25 Index
rankify-index index data/corpus.jsonl --retriever bm25 --output ./indices

# Dense Retrievers
rankify-index index data/corpus.jsonl --retriever dpr --device cuda --batch_size 16

# See all options
rankify-index index --help

API Reference

rankify.indexing

__all__ = ['LuceneIndexer', 'DPRIndexer', 'ContrieverIndexer', 'ColBERTIndexer', 'BGEIndexer'] module-attribute

LuceneIndexer

Bases: BaseIndexer

Lucene Indexer for creating and loading Lucene indices using Pyserini. This class handles the preparation of the corpus, building the index, and loading the index.

Source code in rankify/indexing/lucene_indexer.py

class LuceneIndexer(BaseIndexer):
    """
    Lucene Indexer for creating and loading Lucene indices using Pyserini.
    This class handles the preparation of the corpus, building the index, and loading the index.
    """
    def __init__(self, corpus_path, output_dir="rankify_indices", chunk_size=1024, threads=32, index_type="wiki",
                 retriever_name="bm25",**kwags):
        super().__init__(corpus_path, output_dir, chunk_size, threads, index_type, retriever_name)
        self.id_mapping = {}  # string_id -> integer_id mapping
        self.mapping_file = self.output_dir / "id_mapping.json"

        if index_type =="wiki" and retriever_name == "bm25":
            self.index_dir = self.output_dir / f"bm25_index"
        else:
            self.index_dir = self.output_dir / f"{retriever_name}_index_{index_type}"

    def _create_id_mapping(self):
        """
        Pre-scan the corpus to create string ID -> integer ID mapping.
        """
        import json

        logging.info("Creating ID mapping...")
        self.id_mapping = {}
        next_id = 0

        with open(self.corpus_path, "r", encoding="utf-8") as f:
            for line in f:
                try:
                    doc = json.loads(line.strip())
                    original_id = doc.get("id")
                    if original_id and original_id not in self.id_mapping:
                        self.id_mapping[str(original_id)] = next_id
                        next_id += 1
                except:
                    continue

        logging.info(f"Created mapping for {len(self.id_mapping)} unique IDs")
    def _save_id_mapping(self):
        """Save the ID mapping to a JSON file."""
        import json
        with open(self.mapping_file, "w", encoding="utf-8") as f:
            json.dump(self.id_mapping, f, ensure_ascii=False, indent=2)
        logging.info(f"ID mapping saved to {self.mapping_file}")
    def _save_pyserini_corpus(self):
        """
        Convert the corpus to the Pyserini JSONL format.
        This method uses the `to_pyserini_jsonl` function to convert the corpus file
        :return: str - The path to the converted Pyserini JSONL corpus file.
        """
        return to_pyserini_jsonl(self.corpus_path, self.output_dir, self.chunk_size, self.threads)

    def build_index(self):
        """
        Build the Lucene index from the corpus.
        """
        # 1) Create & save mapping BEFORE converting corpus (so converter can load it)
        self._create_id_mapping()
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self._save_id_mapping()

        # 2) Convert corpus (may return str or Path depending on your converter)
        corpus_path = Path(self._save_pyserini_corpus())

        # 3) Prep temp dir
        temp_corpus_dir = self.output_dir / "temp_corpus"
        if temp_corpus_dir.exists():
            shutil.rmtree(temp_corpus_dir)
        temp_corpus_dir.mkdir(parents=True)

        corpus_file_path = temp_corpus_dir / "corpus.json"

        # 4) Move the file into the temp dir (works across filesystems)
        shutil.move(str(corpus_path), str(corpus_file_path))

        # 5) Reset index dir
        if self.index_dir.exists():
            shutil.rmtree(self.index_dir)
        self.index_dir.mkdir(parents=True)

        # 6) Run Pyserini
        cmd = [
            "python", "-m", "pyserini.index.lucene",
            "-collection", "JsonCollection",
            "-generator", "DefaultLuceneDocumentGenerator",
            "-threads", str(self.threads),
            "-input", str(temp_corpus_dir),
            "-index", str(self.index_dir),
            "-storePositions", "-storeDocvectors", "-storeRaw"
        ]
        logging.info(f"Running Pyserini indexer:\n{' '.join(cmd)}")
        subprocess.run(cmd, check=True)

        # 7) Cleanup & finish
        shutil.rmtree(temp_corpus_dir)
        self._save_title_map()
        logging.info(f"Indexing complete. Index stored at {self.index_dir}")

    def load_index(self):
        """
        Load the Lucene index from the specified directory.

        This method checks if the index directory exists and is not empty,
        and raises an error if the index is locked or not found.
        """
        if not self.index_dir.exists() or not any(self.index_dir.iterdir()):
            raise FileNotFoundError(f"Index directory {self.index_dir} does not exist or is empty.")
        logging.info(f"Index loaded from {self.index_dir}")

build_index()

Build the Lucene index from the corpus.

Source code in rankify/indexing/lucene_indexer.py

def build_index(self):
    """
    Build the Lucene index from the corpus.
    """
    # 1) Create & save mapping BEFORE converting corpus (so converter can load it)
    self._create_id_mapping()
    self.output_dir.mkdir(parents=True, exist_ok=True)
    self._save_id_mapping()

    # 2) Convert corpus (may return str or Path depending on your converter)
    corpus_path = Path(self._save_pyserini_corpus())

    # 3) Prep temp dir
    temp_corpus_dir = self.output_dir / "temp_corpus"
    if temp_corpus_dir.exists():
        shutil.rmtree(temp_corpus_dir)
    temp_corpus_dir.mkdir(parents=True)

    corpus_file_path = temp_corpus_dir / "corpus.json"

    # 4) Move the file into the temp dir (works across filesystems)
    shutil.move(str(corpus_path), str(corpus_file_path))

    # 5) Reset index dir
    if self.index_dir.exists():
        shutil.rmtree(self.index_dir)
    self.index_dir.mkdir(parents=True)

    # 6) Run Pyserini
    cmd = [
        "python", "-m", "pyserini.index.lucene",
        "-collection", "JsonCollection",
        "-generator", "DefaultLuceneDocumentGenerator",
        "-threads", str(self.threads),
        "-input", str(temp_corpus_dir),
        "-index", str(self.index_dir),
        "-storePositions", "-storeDocvectors", "-storeRaw"
    ]
    logging.info(f"Running Pyserini indexer:\n{' '.join(cmd)}")
    subprocess.run(cmd, check=True)

    # 7) Cleanup & finish
    shutil.rmtree(temp_corpus_dir)
    self._save_title_map()
    logging.info(f"Indexing complete. Index stored at {self.index_dir}")

load_index()

Load the Lucene index from the specified directory.

This method checks if the index directory exists and is not empty, and raises an error if the index is locked or not found.

Source code in rankify/indexing/lucene_indexer.py

def load_index(self):
    """
    Load the Lucene index from the specified directory.

    This method checks if the index directory exists and is not empty,
    and raises an error if the index is locked or not found.
    """
    if not self.index_dir.exists() or not any(self.index_dir.iterdir()):
        raise FileNotFoundError(f"Index directory {self.index_dir} does not exist or is empty.")
    logging.info(f"Index loaded from {self.index_dir}")

DPRIndexer

Bases: BaseIndexer

DPR Indexer that builds dense FAISS-based indexes using Pyserini.

Supports indexing using models like DPR, ANCE, or BPR via HuggingFace.

Parameters:

Name	Type	Description	Default
corpus_path	str	Path to the corpus file.	required
encoder_name	str	HuggingFace model name for the DPR encoder.	'facebook/dpr-ctx_encoder-single-nq-base'
output_dir	str	Directory to save the index.	'rankify_indices'
chunk_size	int	Size of chunks to process the corpus.	100
threads	int	Number of threads to use for processing.	32
index_type	str	Type of index to build (default: "wiki").	'wiki'
batch_size	int	Batch size for encoding passages.	16
device	str	Device to use for encoding ("cpu" or "cuda").	'cuda'

Source code in rankify/indexing/dpr_indexer.py

class DPRIndexer(BaseIndexer):
    """
    DPR Indexer that builds dense FAISS-based indexes using Pyserini.

    Supports indexing using models like DPR, ANCE, or BPR via HuggingFace.

    Args:
        corpus_path (str): Path to the corpus file.
        encoder_name (str): HuggingFace model name for the DPR encoder.
        output_dir (str): Directory to save the index.
        chunk_size (int): Size of chunks to process the corpus.
        threads (int): Number of threads to use for processing.
        index_type (str): Type of index to build (default: "wiki").
        batch_size (int): Batch size for encoding passages.
        device (str): Device to use for encoding ("cpu" or "cuda").
    """

    def __init__(self,
                 corpus_path,
                 encoder_name="facebook/dpr-ctx_encoder-single-nq-base",
                 output_dir="rankify_indices",
                 chunk_size=100,
                 threads=32,
                 index_type="wiki",
                 retriever_name="dpr",
                 batch_size=16,
                 device="cuda"):
        super().__init__(corpus_path, output_dir, chunk_size, threads, index_type)
        self.encoder_name = encoder_name
        self.index_dir = self.output_dir / f"dpr_index_{index_type}"
        self.device = device
        self.batch_size = batch_size

    def _save_dense_corpus(self):
        """
        Convert the corpus to dense-compatible Pyserini format.
        Each passage must have a `text` and optionally a `title`.
        """
        return to_pyserini_jsonl_dense(self.corpus_path, self.output_dir, self.chunk_size, self.threads)

    def _create_id_mapping(self, corpus_file):
        """
        Create mapping between FAISS sequential IDs and original document IDs.

        Args:
            corpus_file (Path): Path to the corpus JSONL file

        Returns:
            dict: Mapping from FAISS index (0, 1, 2, ...) to original doc IDs
        """
        logging.info("Creating ID mapping for FAISS index...")

        faiss_to_docid = {}
        docid_to_faiss = {}

        with open(corpus_file, 'r', encoding='utf-8') as f:
            for faiss_idx, line in enumerate(f):
                doc = json.loads(line.strip())
                original_docid = doc.get('docid') or doc.get('id', str(faiss_idx))

                faiss_to_docid[faiss_idx] = str(original_docid)
                docid_to_faiss[str(original_docid)] = faiss_idx

        logging.info(f"Created ID mapping for {len(faiss_to_docid)} documents")
        return faiss_to_docid, docid_to_faiss

    def _save_id_mapping(self, faiss_to_docid, docid_to_faiss):
        """
        Save ID mappings to files.

        Args:
            faiss_to_docid (dict): Mapping from FAISS index to original doc ID
            docid_to_faiss (dict): Mapping from original doc ID to FAISS index
        """
        # Save FAISS index to original docid mapping
        faiss_to_docid_file = self.index_dir / "faiss_to_docid.json"
        with open(faiss_to_docid_file, 'w', encoding='utf-8') as f:
            json.dump(faiss_to_docid, f, indent=2)

        # Save original docid to FAISS index mapping  
        docid_to_faiss_file = self.index_dir / "docid_to_faiss.json"
        with open(docid_to_faiss_file, 'w', encoding='utf-8') as f:
            json.dump(docid_to_faiss, f, indent=2)

        logging.info(f"ID mappings saved to {self.index_dir}")

    def _save_corpus_with_metadata(self, corpus_file):
        """
        Save corpus with metadata for easy retrieval.

        Args:
            corpus_file (Path): Path to the corpus JSONL file
        """
        logging.info("Saving corpus metadata...")

        corpus_metadata = {}

        with open(corpus_file, 'r', encoding='utf-8') as f:
            for line in f:
                doc = json.loads(line.strip())
                docid = doc.get('docid') or doc.get('id')

                # Extract title and contents
                contents = doc.get('contents', '')
                title = doc.get('title', '')

                # If no explicit title, try to extract from contents
                if not title and contents:
                    lines = contents.split('\n')
                    title = lines[0] if lines else "No Title"

                corpus_metadata[str(docid)] = {
                    'title': title,
                    'contents': contents,
                    'text': doc.get('text', contents)  # Some formats use 'text' instead of 'contents'
                }

        # Save corpus metadata
        corpus_metadata_file = self.index_dir / "corpus_metadata.json"
        with open(corpus_metadata_file, 'w', encoding='utf-8') as f:
            json.dump(corpus_metadata, f, indent=2, ensure_ascii=False)

        logging.info(f"Corpus metadata saved to {corpus_metadata_file}")

    def build_index(self):
        """
        Build the DPR dense index using FAISS.
        Steps:
        1. Converts corpus to Pyserini JSONL format.
        2. Creates ID mappings between FAISS and original document IDs.
        3. Runs DPR indexing command using the HuggingFace encoder.
        4. Saves all necessary mapping files.
        """
        corpus_path = self._save_dense_corpus()

        temp_corpus_dir = self.output_dir / "temp_corpus"

        if temp_corpus_dir.exists():
            shutil.rmtree(temp_corpus_dir)

        temp_corpus_dir.mkdir(parents=True)

        # Move corpus into temporary directory
        dense_file = temp_corpus_dir / "corpus.jsonl"
        corpus_path.rename(dense_file)

        # Create ID mappings before indexing
        faiss_to_docid, docid_to_faiss = self._create_id_mapping(dense_file)

        if self.index_dir.exists():
            shutil.rmtree(self.index_dir)
        self.index_dir.mkdir(parents=True)

        print("Start indexing ............................")
        cmd = [
            "python", "-m", "pyserini.encode",

            "input",
            "--corpus", str(dense_file),

            "output",
            "--embeddings", str(self.index_dir / "embeddings"),

            "encoder",
            "--encoder", self.encoder_name,
            "--batch-size", str(self.batch_size),
            "--max-length", "512",
        ]

        if self.device == "cuda":
            if torch.cuda.is_available():
                cmd.extend(["--device", "cuda"])
            else:
                logging.warning("CUDA is not available. Using CPU for encoding.")
                cmd.extend(["--device", "cpu"])
        else:
            cmd.extend(["--device", self.device])

        logging.info(f"Encoding dense vectors with DPR: {' '.join(cmd)}")
        subprocess.run(cmd, check=True)

        # Now index the dense vectors into FAISS
        index_cmd = [
            "python", "-m", "pyserini.index.faiss",
            "--input", str(self.index_dir / "embeddings"),
            "--output", str(self.index_dir),
        ]

        logging.info(f"Building FAISS index: {' '.join(index_cmd)}")
        subprocess.run(index_cmd, check=True)

        # Save ID mappings
        self._save_id_mapping(faiss_to_docid, docid_to_faiss)

        # Save corpus metadata for retrieval
        self._save_corpus_with_metadata(dense_file)

        # Save the original title map (inherited from base class)
        self._save_title_map()

        # Move corpus.jsonl to final location
        if dense_file.exists():
            dest_file = self.index_dir / "corpus.jsonl"
            if dest_file.exists():
                dest_file.unlink()
            shutil.move(str(dense_file), str(dest_file))

        if temp_corpus_dir.exists():
            shutil.rmtree(temp_corpus_dir)

        logging.info(f"Dense indexing complete. Index stored at {self.index_dir}")
        logging.info(f"Index files created:")
        logging.info(f"  - FAISS index files")
        logging.info(f"  - corpus.jsonl")
        logging.info(f"  - faiss_to_docid.json")
        logging.info(f"  - docid_to_faiss.json") 
        logging.info(f"  - corpus_metadata.json")

    def load_index(self):
        """
        Load the DPR index from the specified directory.

        This method checks if the index directory exists and is not empty,
        and raises an error if the index is locked or not found.
        """
        if not self.index_dir.exists() or not any(self.index_dir.iterdir()):
            raise FileNotFoundError(f"Index directory {self.index_dir} does not exist or is empty.")

        # Check for required mapping files
        required_files = [
            "faiss_to_docid.json",
            "docid_to_faiss.json", 
            "corpus_metadata.json"
        ]

        for file_name in required_files:
            file_path = self.index_dir / file_name
            if not file_path.exists():
                logging.warning(f"Missing mapping file: {file_path}")

        logging.info(f"Index loaded from {self.index_dir}")

build_index()

Build the DPR dense index using FAISS. Steps: 1. Converts corpus to Pyserini JSONL format. 2. Creates ID mappings between FAISS and original document IDs. 3. Runs DPR indexing command using the HuggingFace encoder. 4. Saves all necessary mapping files.

Source code in rankify/indexing/dpr_indexer.py

def build_index(self):
    """
    Build the DPR dense index using FAISS.
    Steps:
    1. Converts corpus to Pyserini JSONL format.
    2. Creates ID mappings between FAISS and original document IDs.
    3. Runs DPR indexing command using the HuggingFace encoder.
    4. Saves all necessary mapping files.
    """
    corpus_path = self._save_dense_corpus()

    temp_corpus_dir = self.output_dir / "temp_corpus"

    if temp_corpus_dir.exists():
        shutil.rmtree(temp_corpus_dir)

    temp_corpus_dir.mkdir(parents=True)

    # Move corpus into temporary directory
    dense_file = temp_corpus_dir / "corpus.jsonl"
    corpus_path.rename(dense_file)

    # Create ID mappings before indexing
    faiss_to_docid, docid_to_faiss = self._create_id_mapping(dense_file)

    if self.index_dir.exists():
        shutil.rmtree(self.index_dir)
    self.index_dir.mkdir(parents=True)

    print("Start indexing ............................")
    cmd = [
        "python", "-m", "pyserini.encode",

        "input",
        "--corpus", str(dense_file),

        "output",
        "--embeddings", str(self.index_dir / "embeddings"),

        "encoder",
        "--encoder", self.encoder_name,
        "--batch-size", str(self.batch_size),
        "--max-length", "512",
    ]

    if self.device == "cuda":
        if torch.cuda.is_available():
            cmd.extend(["--device", "cuda"])
        else:
            logging.warning("CUDA is not available. Using CPU for encoding.")
            cmd.extend(["--device", "cpu"])
    else:
        cmd.extend(["--device", self.device])

    logging.info(f"Encoding dense vectors with DPR: {' '.join(cmd)}")
    subprocess.run(cmd, check=True)

    # Now index the dense vectors into FAISS
    index_cmd = [
        "python", "-m", "pyserini.index.faiss",
        "--input", str(self.index_dir / "embeddings"),
        "--output", str(self.index_dir),
    ]

    logging.info(f"Building FAISS index: {' '.join(index_cmd)}")
    subprocess.run(index_cmd, check=True)

    # Save ID mappings
    self._save_id_mapping(faiss_to_docid, docid_to_faiss)

    # Save corpus metadata for retrieval
    self._save_corpus_with_metadata(dense_file)

    # Save the original title map (inherited from base class)
    self._save_title_map()

    # Move corpus.jsonl to final location
    if dense_file.exists():
        dest_file = self.index_dir / "corpus.jsonl"
        if dest_file.exists():
            dest_file.unlink()
        shutil.move(str(dense_file), str(dest_file))

    if temp_corpus_dir.exists():
        shutil.rmtree(temp_corpus_dir)

    logging.info(f"Dense indexing complete. Index stored at {self.index_dir}")
    logging.info(f"Index files created:")
    logging.info(f"  - FAISS index files")
    logging.info(f"  - corpus.jsonl")
    logging.info(f"  - faiss_to_docid.json")
    logging.info(f"  - docid_to_faiss.json") 
    logging.info(f"  - corpus_metadata.json")

load_index()

Load the DPR index from the specified directory.

This method checks if the index directory exists and is not empty, and raises an error if the index is locked or not found.

Source code in rankify/indexing/dpr_indexer.py

def load_index(self):
    """
    Load the DPR index from the specified directory.

    This method checks if the index directory exists and is not empty,
    and raises an error if the index is locked or not found.
    """
    if not self.index_dir.exists() or not any(self.index_dir.iterdir()):
        raise FileNotFoundError(f"Index directory {self.index_dir} does not exist or is empty.")

    # Check for required mapping files
    required_files = [
        "faiss_to_docid.json",
        "docid_to_faiss.json", 
        "corpus_metadata.json"
    ]

    for file_name in required_files:
        file_path = self.index_dir / file_name
        if not file_path.exists():
            logging.warning(f"Missing mapping file: {file_path}")

    logging.info(f"Index loaded from {self.index_dir}")

ContrieverIndexer

Bases: BaseIndexer

Corrected Contriever Indexer that efficiently builds dense FAISS-based indexes.

Key improvements: - Memory-efficient batch processing - Proper error handling - Streamlined embedding indexing - Better resource management

Source code in rankify/indexing/contriever_indexer.py

class ContrieverIndexer(BaseIndexer):
    """
    Corrected Contriever Indexer that efficiently builds dense FAISS-based indexes.

    Key improvements:
    - Memory-efficient batch processing
    - Proper error handling
    - Streamlined embedding indexing
    - Better resource management
    """

    def __init__(self,
                 corpus_path,
                 encoder_name="facebook/contriever",
                 output_dir="rankify_indices",
                 chunk_size=5000000,
                 threads=32,
                 index_type="wiki",
                 batch_size=1000000,
                 retriever_name="contriever",
                 device="cuda",
                 embedding_batch_size=32):
        super().__init__(corpus_path, output_dir, chunk_size, threads, index_type, retriever_name)
        self.encoder_name = encoder_name
        self.device = device
        self.batch_size = batch_size
        self.embedding_batch_size = embedding_batch_size  # For embedding generation
        self.index_dir = self.output_dir / f"contriever_index_{index_type}"

        # Validate parameters
        if not Path(corpus_path).exists():
            raise FileNotFoundError(f"Corpus file not found: {corpus_path}")

    def _save_dense_embeddings(self) -> List[Path]:
        """
        Convert the corpus into Contriever embeddings stored in shard .pkl files.
        Returns list of paths to embedding files.
        """
        logging.info("Generating Contriever embeddings...")

        try:
            embedding_files = to_contriever_embedding_chunks(
                self.corpus_path,
                self.index_dir,
                self.chunk_size,
                self.encoder_name,
                self.embedding_batch_size,  # Use smaller batch for embeddings
                self.device
            )

            if not embedding_files:
                raise ValueError("No embedding files were generated")

            logging.info(f"Generated {len(embedding_files)} embedding files")
            return embedding_files

        except Exception as e:
            logging.error(f"Error generating embeddings: {e}")
            raise

    def _save_corpus(self) -> str:
        """Convert the corpus to dense-compatible Pyserini format."""
        logging.info("Converting corpus to Pyserini JSONL format...")
        return to_tsv(self.corpus_path, self.index_dir, self.chunk_size, self.threads)

    def build_index(self):
        """Build the Contriever index with improved error handling and efficiency."""
        try:
            # Clean and create index directory
            if self.index_dir.exists():
                shutil.rmtree(self.index_dir)
            self.index_dir.mkdir(parents=True)

            # Convert corpus
            self._save_corpus()

            # Generate embeddings
            embedding_files = self._save_dense_embeddings()

            # Create index
            logging.info("Creating Contriever index...")
            index = Indexer(vector_sz=768, n_subquantizers=0, n_bits=8)

            # Index embeddings efficiently
            self._index_encoded_data_efficient(index, embedding_files)

            # Serialize index
            logging.info("Serializing index...")
            index.serialize(self.index_dir)

            # Save title map
            self._save_title_map()

            # Clean up embedding files to save space
            self._cleanup_embedding_files(embedding_files)

            logging.info(f"✅ Contriever index successfully stored at {self.index_dir}")

        except Exception as e:
            logging.error(f"Error building index: {e}")
            # Clean up on failure
            if self.index_dir.exists():
                shutil.rmtree(self.index_dir)
            raise

    def _index_encoded_data_efficient(self, index: Indexer, embedding_files: List[Path]):
        """
        Efficiently loads and indexes encoded passage embeddings into FAISS.

        This version processes files one at a time and batches within each file,
        avoiding memory issues with large corpora.
        """
        total_indexed = 0

        for file_path in tqdm(embedding_files, desc="Indexing embedding files"):
            try:
                logging.info(f"Processing embedding file: {file_path}")

                # Load embeddings from file
                with open(file_path, "rb") as fin:
                    ids, embeddings = pickle.load(fin)

                if len(ids) != len(embeddings):
                    raise ValueError(f"Mismatch between IDs ({len(ids)}) and embeddings ({len(embeddings)}) in {file_path}")

                # Process this file's embeddings in batches
                num_embeddings = len(embeddings)
                for start_idx in tqdm(range(0, num_embeddings, self.batch_size), 
                                    desc=f"Batching {file_path.name}", leave=False):
                    end_idx = min(start_idx + self.batch_size, num_embeddings)

                    batch_ids = ids[start_idx:end_idx]
                    batch_embeddings = embeddings[start_idx:end_idx]

                    # Validate embeddings
                    if not isinstance(batch_embeddings, np.ndarray):
                        batch_embeddings = np.array(batch_embeddings)

                    if batch_embeddings.dtype != np.float32:
                        batch_embeddings = batch_embeddings.astype(np.float32)

                    # Add to index
                    index.index_data(batch_ids, batch_embeddings)
                    total_indexed += len(batch_ids)

                logging.info(f"Processed {num_embeddings} embeddings from {file_path.name}")

            except Exception as e:
                logging.error(f"Error processing file {file_path}: {e}")
                raise

        logging.info(f"✅ Successfully indexed {total_indexed} embeddings")

    def _cleanup_embedding_files(self, embedding_files: List[Path]):
        """Clean up embedding files after indexing to save disk space."""
        try:
            for file_path in embedding_files:
                if file_path.exists():
                    file_path.unlink()
            logging.info(f"Cleaned up {len(embedding_files)} embedding files")
        except Exception as e:
            logging.warning(f"Error cleaning up embedding files: {e}")

    def load_index(self) -> Indexer:
        """Load the Contriever index from the specified directory."""
        if not self.index_dir.exists():
            raise FileNotFoundError(f"Index directory {self.index_dir} does not exist")

        if not any(self.index_dir.iterdir()):
            raise FileNotFoundError(f"Index directory {self.index_dir} is empty")

        try:
            index = Indexer(vector_sz=768, n_subquantizers=0, n_bits=8)
            index.deserialize_from(self.index_dir)
            logging.info(f"✅ Successfully loaded index from {self.index_dir}")
            return index

        except Exception as e:
            logging.error(f"Error loading index from {self.index_dir}: {e}")
            raise

    def validate_index(self) -> bool:
        """Validate that the index was built correctly."""
        try:
            index = self.load_index()

            # Check if index has embeddings
            if hasattr(index, 'index') and hasattr(index.index, 'ntotal'):
                num_vectors = index.index.ntotal
                logging.info(f"Index contains {num_vectors} vectors")
                return num_vectors > 0

            return True

        except Exception as e:
            logging.error(f"Index validation failed: {e}")
            return False

    def get_index_stats(self) -> dict:
        """Get statistics about the built index."""
        try:
            index = self.load_index()
            stats = {
                "index_type": "contriever",
                "vector_dimension": 768,
                "index_directory": str(self.index_dir),
                "encoder_model": self.encoder_name
            }

            if hasattr(index, 'index') and hasattr(index.index, 'ntotal'):
                stats["num_vectors"] = index.index.ntotal

            # Check index directory size
            total_size = sum(f.stat().st_size for f in self.index_dir.rglob('*') if f.is_file())
            stats["index_size_mb"] = total_size / (1024 * 1024)

            return stats

        except Exception as e:
            logging.error(f"Error getting index stats: {e}")
            return {"error": str(e)}

build_index()

Build the Contriever index with improved error handling and efficiency.

Source code in rankify/indexing/contriever_indexer.py

def build_index(self):
    """Build the Contriever index with improved error handling and efficiency."""
    try:
        # Clean and create index directory
        if self.index_dir.exists():
            shutil.rmtree(self.index_dir)
        self.index_dir.mkdir(parents=True)

        # Convert corpus
        self._save_corpus()

        # Generate embeddings
        embedding_files = self._save_dense_embeddings()

        # Create index
        logging.info("Creating Contriever index...")
        index = Indexer(vector_sz=768, n_subquantizers=0, n_bits=8)

        # Index embeddings efficiently
        self._index_encoded_data_efficient(index, embedding_files)

        # Serialize index
        logging.info("Serializing index...")
        index.serialize(self.index_dir)

        # Save title map
        self._save_title_map()

        # Clean up embedding files to save space
        self._cleanup_embedding_files(embedding_files)

        logging.info(f"✅ Contriever index successfully stored at {self.index_dir}")

    except Exception as e:
        logging.error(f"Error building index: {e}")
        # Clean up on failure
        if self.index_dir.exists():
            shutil.rmtree(self.index_dir)
        raise

load_index()

Load the Contriever index from the specified directory.

Source code in rankify/indexing/contriever_indexer.py

def load_index(self) -> Indexer:
    """Load the Contriever index from the specified directory."""
    if not self.index_dir.exists():
        raise FileNotFoundError(f"Index directory {self.index_dir} does not exist")

    if not any(self.index_dir.iterdir()):
        raise FileNotFoundError(f"Index directory {self.index_dir} is empty")

    try:
        index = Indexer(vector_sz=768, n_subquantizers=0, n_bits=8)
        index.deserialize_from(self.index_dir)
        logging.info(f"✅ Successfully loaded index from {self.index_dir}")
        return index

    except Exception as e:
        logging.error(f"Error loading index from {self.index_dir}: {e}")
        raise

validate_index()

Validate that the index was built correctly.

Source code in rankify/indexing/contriever_indexer.py

def validate_index(self) -> bool:
    """Validate that the index was built correctly."""
    try:
        index = self.load_index()

        # Check if index has embeddings
        if hasattr(index, 'index') and hasattr(index.index, 'ntotal'):
            num_vectors = index.index.ntotal
            logging.info(f"Index contains {num_vectors} vectors")
            return num_vectors > 0

        return True

    except Exception as e:
        logging.error(f"Index validation failed: {e}")
        return False

get_index_stats()

Get statistics about the built index.

Source code in rankify/indexing/contriever_indexer.py

def get_index_stats(self) -> dict:
    """Get statistics about the built index."""
    try:
        index = self.load_index()
        stats = {
            "index_type": "contriever",
            "vector_dimension": 768,
            "index_directory": str(self.index_dir),
            "encoder_model": self.encoder_name
        }

        if hasattr(index, 'index') and hasattr(index.index, 'ntotal'):
            stats["num_vectors"] = index.index.ntotal

        # Check index directory size
        total_size = sum(f.stat().st_size for f in self.index_dir.rglob('*') if f.is_file())
        stats["index_size_mb"] = total_size / (1024 * 1024)

        return stats

    except Exception as e:
        logging.error(f"Error getting index stats: {e}")
        return {"error": str(e)}

ColBERTIndexer

Bases: BaseIndexer

ColBERT Indexer that properly handles string IDs by converting them to sequential integers.

This indexer: 1. Converts string IDs to sequential integers (required by ColBERT) 2. Maintains ID mappings for retrieval 3. Creates proper TSV header format 4. Generates all necessary files for the retriever

Source code in rankify/indexing/colbert_indexer.py

class ColBERTIndexer(BaseIndexer):
    """
    ColBERT Indexer that properly handles string IDs by converting them to sequential integers.

    This indexer:
    1. Converts string IDs to sequential integers (required by ColBERT)
    2. Maintains ID mappings for retrieval
    3. Creates proper TSV header format
    4. Generates all necessary files for the retriever
    """

    def __init__(self,
                 corpus_path,
                 encoder_name="colbert-ir/colbertv2.0",
                 output_dir="rankify_indices",
                 chunk_size=100,
                 threads=32,
                 index_type="wiki",
                 batch_size=64,
                 retriever_name="colbert",
                 device="cuda"):
        super().__init__(corpus_path, output_dir, chunk_size, threads, index_type, retriever_name)
        self.encoder_name = encoder_name
        self.index_dir = self.output_dir / f"colbert_index_{index_type}"
        self.device = device
        self.batch_size = batch_size
        self.id_mapping_file = self.index_dir / "id_mapping.json"
        self.passages_file = self.index_dir / "passages.tsv"  # For retriever compatibility

    def _create_colbert_tsv_from_jsonl(self) -> tuple[str, dict, dict]:
        """
        Create ColBERT-compatible TSV directly from JSONL corpus.

        Returns:
            tuple: (tsv_path, original_to_sequential_mapping, sequential_to_original_mapping)
        """
        logging.info("Creating ColBERT TSV directly from JSONL corpus...")

        colbert_tsv_path = self.index_dir / "collection.tsv"  # ColBERT collection file
        passages_tsv_path = self.passages_file  # For retriever

        original_to_sequential = {}
        sequential_to_original = {}

        # Create both files simultaneously
        with open(self.corpus_path, 'r', encoding='utf-8') as infile, \
             open(colbert_tsv_path, 'w', encoding='utf-8', newline='') as collection_file, \
             open(passages_tsv_path, 'w', encoding='utf-8', newline='') as passages_file:

            collection_writer = csv.writer(collection_file, delimiter='\t')
            passages_writer = csv.writer(passages_file, delimiter='\t')

            # Write headers
            collection_writer.writerow(['id', 'text', 'title'])  # ColBERT format
            passages_writer.writerow(['id', 'text', 'title'])    # Retriever format

            sequential_id = 1  # Start from 1 to match line indices
            processed_docs = 0

            for line_idx, line in enumerate(infile):
                line = line.strip()
                if not line:
                    continue

                try:
                    doc = json.loads(line)

                    # Extract fields with fallbacks
                    original_id = doc.get("id", f"doc_{line_idx}")
                    text = doc.get("contents", doc.get("text", doc.get("passage", ""))).strip()
                    title = doc.get("title", "").strip()

                    if not text:  # Skip empty documents
                        logging.warning(f"Skipping document {original_id} with empty text")
                        continue

                    # Clean text and title for TSV (remove problematic characters)
                    text = self._clean_tsv_field(text)
                    title = self._clean_tsv_field(title)

                    # Create mappings
                    original_to_sequential[original_id] = sequential_id
                    sequential_to_original[sequential_id] = original_id

                    # Write to both files
                    collection_writer.writerow([sequential_id, text, title])  # For ColBERT
                    passages_writer.writerow([original_id, text, title])      # For retriever (keeps original IDs)

                    sequential_id += 1
                    processed_docs += 1

                    if processed_docs % 10000 == 0:
                        logging.info(f"Processed {processed_docs} documents...")

                except json.JSONDecodeError as e:
                    logging.warning(f"Skipping malformed JSON at line {line_idx}: {e}")
                    continue
                except Exception as e:
                    logging.warning(f"Error processing document at line {line_idx}: {e}")
                    continue

        logging.info(f"✅ Created ColBERT TSV with {processed_docs} documents")
        logging.info(f"   Collection file: {colbert_tsv_path}")
        logging.info(f"   Passages file: {passages_tsv_path}")

        return str(colbert_tsv_path), original_to_sequential, sequential_to_original

    def _clean_tsv_field(self, field: str) -> str:
        """Clean field for TSV format by removing problematic characters."""
        if not field:
            return ""

        # Replace tabs, newlines, and carriage returns with spaces
        field = field.replace('\t', ' ').replace('\n', ' ').replace('\r', ' ')

        # Remove multiple spaces
        field = ' '.join(field.split())

        # Truncate if too long (prevent memory issues)
        if len(field) > 10000:
            field = field[:10000] + "..."

        return field

    def _save_id_mappings(self, original_to_sequential: dict, sequential_to_original: dict):
        """Save ID mappings for retriever use."""
        mappings = {
            "original_to_sequential": original_to_sequential,
            "sequential_to_original": {str(k): v for k, v in sequential_to_original.items()},
            "total_documents": len(original_to_sequential),
            "encoder_name": self.encoder_name,
            "index_type": self.index_type
        }

        with open(self.id_mapping_file, 'w', encoding='utf-8') as f:
            json.dump(mappings, f, indent=2, ensure_ascii=False)

        logging.info(f"✅ Saved ID mappings to {self.id_mapping_file}")
        logging.info(f"   Mapped {len(original_to_sequential)} original IDs to sequential IDs")

    def _verify_colbert_tsv(self, tsv_path: str):
        """Verify that the TSV format is correct for ColBERT."""
        logging.info(f"Verifying ColBERT TSV format: {tsv_path}")

        issues_found = []

        with open(tsv_path, 'r', encoding='utf-8') as f:
            lines = f.readlines()

        if not lines:
            raise ValueError("TSV file is empty")

        # Check header
        header = lines[0].strip().split('\t')
        expected_header = ['id', 'text', 'title']
        if header != expected_header:
            issues_found.append(f"Header mismatch: expected {expected_header}, got {header}")

        # Check first 10 data lines
        for i in range(1, min(11, len(lines))):
            line = lines[i].strip()
            if not line:
                continue

            parts = line.split('\t')
            if len(parts) < 2:
                issues_found.append(f"Line {i}: Insufficient columns: {len(parts)}")
                continue

            pid = parts[0]
            text = parts[1]

            # Verify PID is integer and matches expected sequence
            try:
                pid_int = int(pid)
                expected_pid = i  # Line index should match PID for ColBERT
                if pid_int != expected_pid:
                    issues_found.append(f"Line {i}: PID {pid} doesn't match expected {expected_pid}")
            except ValueError:
                issues_found.append(f"Line {i}: PID '{pid}' is not an integer")

            # Check for empty text
            if not text.strip():
                issues_found.append(f"Line {i}: Empty text field")

        if issues_found:
            logging.error("❌ TSV verification failed:")
            for issue in issues_found:
                logging.error(f"   {issue}")
            raise ValueError(f"TSV format issues found: {len(issues_found)} problems")

        logging.info(f"✅ TSV verification passed. Total lines: {len(lines)}")

    def build_index(self):
        """Build the ColBERT index with proper string ID handling."""
        logging.info("Building ColBERT index with string ID support...")

        # Clean up existing index
        if self.index_dir.exists():
            logging.info(f"Removing existing index: {self.index_dir}")
            shutil.rmtree(self.index_dir)

        self.index_dir.mkdir(parents=True, exist_ok=True)

        try:
            # Create ColBERT-compatible TSV and ID mappings
            collection_path, original_to_sequential, sequential_to_original = self._create_colbert_tsv_from_jsonl()

            # Save ID mappings for retriever
            self._save_id_mappings(original_to_sequential, sequential_to_original)

            # Verify TSV format
            self._verify_colbert_tsv(collection_path)

        except Exception as e:
            logging.error(f"Failed to create ColBERT TSV: {e}")
            raise

        # Configure ColBERT
        config = ColBERTConfig(
            root=str(self.output_dir),
            index_root=str(self.output_dir),
            index_name=f"colbert_index_{self.index_type}",
            index_path=str(self.index_dir),
            checkpoint=self.encoder_name,
            collection=collection_path,
            doc_maxlen=180,        # Max document length
            nbits=2,               # Compression bits
            avoid_fork_if_possible=True,
            kmeans_niters=4,       # K-means iterations
            nranks=1,              # Single GPU/process
            bsize=self.batch_size,
            index_bsize=self.batch_size,
        )

        # Create indexer and build index
        logging.info(f"Initializing ColBERT indexer with model: {self.encoder_name}")
        indexer = Indexer(checkpoint=self.encoder_name, config=config, verbose=2)

        try:
            logging.info("Starting ColBERT indexing process...")
            index_path = indexer.index(
                name=config.index_name, 
                collection=Collection.cast(collection_path), 
                overwrite=True,
            )

            logging.info(f"✅ ColBERT index successfully built!")
            logging.info(f"   Index path: {index_path}")
            logging.info(f"   Collection: {collection_path}")
            logging.info(f"   Passages file: {self.passages_file}")
            logging.info(f"   ID mappings: {self.id_mapping_file}")

            # Save title map for compatibility
            self._save_title_map()

            return index_path

        except Exception as e:
            logging.error(f"❌ ColBERT indexing failed: {e}")

            # Provide debugging info
            self._debug_indexing_failure(collection_path)
            raise

    def _debug_indexing_failure(self, collection_path: str):
        """Debug helper for indexing failures."""
        logging.info("=== DEBUGGING INDEXING FAILURE ===")

        # Check collection file
        if not os.path.exists(collection_path):
            logging.error(f"Collection file missing: {collection_path}")
            return

        # Check file size and format
        file_size = os.path.getsize(collection_path) / (1024 * 1024)  # MB
        logging.info(f"Collection file size: {file_size:.2f} MB")

        with open(collection_path, 'r', encoding='utf-8') as f:
            lines = f.readlines()

        logging.info(f"Collection lines: {len(lines)}")

        if len(lines) > 0:
            logging.info(f"First line (header): {lines[0].strip()}")
        if len(lines) > 1:
            logging.info(f"Second line (data): {lines[1].strip()}")

        # Check for common issues
        if len(lines) < 2:
            logging.error("Collection has no data lines")
        elif len(lines) > 1000000:
            logging.warning(f"Large collection ({len(lines)} lines) - may need more memory")

    def load_index(self):
        """Load the ColBERT index."""
        logging.info(f"Loading ColBERT index from {self.index_dir}...")

        try:
            # Check if index files exist
            index_files = list(self.index_dir.glob("*.pt"))  # PyTorch files
            if not index_files:
                raise FileNotFoundError(f"No index files found in {self.index_dir}")

            # Load using IndexLoader
            index = IndexLoader(index_path=str(self.index_dir), use_gpu=(self.device != "cpu"))
            logging.info(f"✅ Successfully loaded ColBERT index")
            return index

        except Exception as e:
            logging.error(f"❌ Failed to load ColBERT index: {e}")
            return None

    def load_id_mappings(self):
        """Load ID mappings for retrieval."""
        if not self.id_mapping_file.exists():
            logging.warning(f"ID mapping file not found: {self.id_mapping_file}")
            return {}, {}

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

            original_to_sequential = mappings.get("original_to_sequential", {})
            sequential_to_original = {int(k): v for k, v in mappings.get("sequential_to_original", {}).items()}

            logging.info(f"Loaded ID mappings: {len(original_to_sequential)} entries")
            return original_to_sequential, sequential_to_original

        except Exception as e:
            logging.error(f"Error loading ID mappings: {e}")
            return {}, {}

    def get_original_id(self, sequential_id: int) -> str:
        """Convert sequential ID back to original ID."""
        _, sequential_to_original = self.load_id_mappings()
        return sequential_to_original.get(sequential_id, str(sequential_id))

    def get_sequential_id(self, original_id: str) -> int:
        """Convert original ID to sequential ID."""
        original_to_sequential, _ = self.load_id_mappings()
        return original_to_sequential.get(original_id, -1)

build_index()

Build the ColBERT index with proper string ID handling.

Source code in rankify/indexing/colbert_indexer.py

def build_index(self):
    """Build the ColBERT index with proper string ID handling."""
    logging.info("Building ColBERT index with string ID support...")

    # Clean up existing index
    if self.index_dir.exists():
        logging.info(f"Removing existing index: {self.index_dir}")
        shutil.rmtree(self.index_dir)

    self.index_dir.mkdir(parents=True, exist_ok=True)

    try:
        # Create ColBERT-compatible TSV and ID mappings
        collection_path, original_to_sequential, sequential_to_original = self._create_colbert_tsv_from_jsonl()

        # Save ID mappings for retriever
        self._save_id_mappings(original_to_sequential, sequential_to_original)

        # Verify TSV format
        self._verify_colbert_tsv(collection_path)

    except Exception as e:
        logging.error(f"Failed to create ColBERT TSV: {e}")
        raise

    # Configure ColBERT
    config = ColBERTConfig(
        root=str(self.output_dir),
        index_root=str(self.output_dir),
        index_name=f"colbert_index_{self.index_type}",
        index_path=str(self.index_dir),
        checkpoint=self.encoder_name,
        collection=collection_path,
        doc_maxlen=180,        # Max document length
        nbits=2,               # Compression bits
        avoid_fork_if_possible=True,
        kmeans_niters=4,       # K-means iterations
        nranks=1,              # Single GPU/process
        bsize=self.batch_size,
        index_bsize=self.batch_size,
    )

    # Create indexer and build index
    logging.info(f"Initializing ColBERT indexer with model: {self.encoder_name}")
    indexer = Indexer(checkpoint=self.encoder_name, config=config, verbose=2)

    try:
        logging.info("Starting ColBERT indexing process...")
        index_path = indexer.index(
            name=config.index_name, 
            collection=Collection.cast(collection_path), 
            overwrite=True,
        )

        logging.info(f"✅ ColBERT index successfully built!")
        logging.info(f"   Index path: {index_path}")
        logging.info(f"   Collection: {collection_path}")
        logging.info(f"   Passages file: {self.passages_file}")
        logging.info(f"   ID mappings: {self.id_mapping_file}")

        # Save title map for compatibility
        self._save_title_map()

        return index_path

    except Exception as e:
        logging.error(f"❌ ColBERT indexing failed: {e}")

        # Provide debugging info
        self._debug_indexing_failure(collection_path)
        raise

load_index()

Load the ColBERT index.

Source code in rankify/indexing/colbert_indexer.py

def load_index(self):
    """Load the ColBERT index."""
    logging.info(f"Loading ColBERT index from {self.index_dir}...")

    try:
        # Check if index files exist
        index_files = list(self.index_dir.glob("*.pt"))  # PyTorch files
        if not index_files:
            raise FileNotFoundError(f"No index files found in {self.index_dir}")

        # Load using IndexLoader
        index = IndexLoader(index_path=str(self.index_dir), use_gpu=(self.device != "cpu"))
        logging.info(f"✅ Successfully loaded ColBERT index")
        return index

    except Exception as e:
        logging.error(f"❌ Failed to load ColBERT index: {e}")
        return None

load_id_mappings()

Load ID mappings for retrieval.

Source code in rankify/indexing/colbert_indexer.py

def load_id_mappings(self):
    """Load ID mappings for retrieval."""
    if not self.id_mapping_file.exists():
        logging.warning(f"ID mapping file not found: {self.id_mapping_file}")
        return {}, {}

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

        original_to_sequential = mappings.get("original_to_sequential", {})
        sequential_to_original = {int(k): v for k, v in mappings.get("sequential_to_original", {}).items()}

        logging.info(f"Loaded ID mappings: {len(original_to_sequential)} entries")
        return original_to_sequential, sequential_to_original

    except Exception as e:
        logging.error(f"Error loading ID mappings: {e}")
        return {}, {}

get_original_id(sequential_id)

Convert sequential ID back to original ID.

Source code in rankify/indexing/colbert_indexer.py

def get_original_id(self, sequential_id: int) -> str:
    """Convert sequential ID back to original ID."""
    _, sequential_to_original = self.load_id_mappings()
    return sequential_to_original.get(sequential_id, str(sequential_id))

get_sequential_id(original_id)

Convert original ID to sequential ID.

Source code in rankify/indexing/colbert_indexer.py

def get_sequential_id(self, original_id: str) -> int:
    """Convert original ID to sequential ID."""
    original_to_sequential, _ = self.load_id_mappings()
    return original_to_sequential.get(original_id, -1)

BGEIndexer

Bases: BaseIndexer

Corrected BGE Indexer that builds dense FAISS-based indexes with proper normalization.

Key improvements: - L2 normalization of embeddings - Proper CLS token extraction for BGE models - Cosine similarity via normalized embeddings - Better error handling and validation

Source code in rankify/indexing/bge_indexer.py

class BGEIndexer(BaseIndexer):
    """
    Corrected BGE Indexer that builds dense FAISS-based indexes with proper normalization.

    Key improvements:
    - L2 normalization of embeddings
    - Proper CLS token extraction for BGE models
    - Cosine similarity via normalized embeddings
    - Better error handling and validation
    """

    def __init__(self,
                 corpus_path,
                 encoder_name="BAAI/bge-large-en-v1.5",
                 output_dir="rankify_indices",
                 index_dir=None,
                 chunk_size=50000,
                 threads=32,
                 index_type="wiki",
                 batch_size=32,
                 retriever_name="bge",
                 device="cuda"):
        super().__init__(corpus_path, output_dir, chunk_size, threads, index_type, retriever_name)
        self.encoder_name = encoder_name
        self.device = device
        self.batch_size = batch_size
        if not index_dir:
            self.index_dir = self.output_dir / f"bge_index_{index_type}"
        else:
            self.index_dir = Path(index_dir)

    def _normalize_embeddings(self, embeddings):
        """L2 normalize embeddings for cosine similarity."""
        return embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)

    def _extract_bge_embeddings(self, texts, tokenizer, model):
        """
        Extract embeddings using BGE model with proper CLS token handling.
        """
        # Tokenize with proper parameters for BGE
        inputs = tokenizer(
            texts, 
            padding=True, 
            truncation=True, 
            max_length=512,
            return_tensors="pt"
        ).to(self.device)

        with torch.no_grad():
            outputs = model(**inputs)
            # BGE uses CLS token (first token) representation
            embeddings = outputs.last_hidden_state[:, 0, :]  # CLS token

        return embeddings.cpu().numpy()

    def build_faiss_index(self):
        """Build FAISS index with normalized embeddings for cosine similarity."""
        embeddings_path = self.index_dir / "bge_embeddings.h5"
        logging.info(f"Loading embeddings from {embeddings_path}...")

        with h5py.File(embeddings_path, "r") as f:
            embeddings = f["embeddings"][:].astype(np.float32)

        # Normalize embeddings for cosine similarity
        logging.info("Normalizing embeddings for cosine similarity...")
        embeddings = self._normalize_embeddings(embeddings)

        embedding_dim = embeddings.shape[1]
        logging.info(f"Building FAISS index with embedding dimension {embedding_dim}...")

        # Use IndexFlatIP with normalized embeddings = cosine similarity
        index = faiss.IndexFlatIP(embedding_dim)

        # Add embeddings in batches for memory efficiency
        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])
            index.add(embeddings[start:end])

        index_path = self.index_dir / "faiss_index.bin"
        logging.info(f"Saving FAISS index to {index_path}...")
        faiss.write_index(index, str(index_path))

    def _merge_embedding_chunks(self):
        """Merge embedding chunks with proper validation."""
        chunk_files = list(self.index_dir.glob("embeddings_chunk_*.pkl"))
        logging.info(f"Found {len(chunk_files)} embedding chunk files to merge.")

        if not chunk_files:
            raise ValueError("No embedding chunk files found to merge.")

        all_embeddings = []
        all_ids = []
        dims = []

        for file in tqdm(chunk_files, desc="Loading embedding chunks"):
            try:
                with open(file, "rb") as fin:
                    chunk = pickle.load(fin)

                if not isinstance(chunk, list):
                    raise ValueError(f"Unexpected format in {file.name}: {type(chunk)}")

                for entry in chunk:
                    embedding = entry.get("embedding")
                    doc_id = entry.get("id")

                    if embedding is None or doc_id is None:
                        logging.warning(f"Missing keys in chunk entry from {file.name}")
                        continue

                    # Ensure embedding is numpy array
                    if isinstance(embedding, list):
                        embedding = np.array(embedding)

                    all_embeddings.append(embedding)
                    all_ids.append(doc_id)
                    dims.append(len(embedding))

            except Exception as e:
                logging.error(f"Error processing file {file}: {e}")
                continue

        if not all_embeddings:
            raise ValueError("No valid embeddings found in chunk files.")

        if len(set(dims)) != 1:
            raise ValueError(f"Embedding dimension mismatch: found dimensions {set(dims)}")

        # Stack embeddings and normalize
        all_embeddings = np.vstack(all_embeddings).astype(np.float32)
        logging.info(f"Loaded {len(all_embeddings)} embeddings with dimension {all_embeddings.shape[1]}")

        # Save merged embeddings
        merged_embedding_path = self.index_dir / "bge_embeddings.h5"
        with h5py.File(merged_embedding_path, "w") as hf:
            hf.create_dataset("embeddings", data=all_embeddings)

        # Save merged IDs
        merged_ids_path = self.index_dir / "bge_doc_ids.pkl"
        with open(merged_ids_path, "wb") as f:
            pickle.dump(all_ids, f)

        # Clean up chunk files
        for file in chunk_files:
            file.unlink()

        logging.info(f"Merged embeddings saved to {merged_embedding_path}")
        logging.info(f"Merged IDs saved to {merged_ids_path}")

        return merged_embedding_path, merged_ids_path

    def _save_dense_embeddings(self):
        """Save dense embeddings with proper BGE handling."""
        logging.info(f"Loading BGE model: {self.encoder_name}")
        tokenizer = AutoTokenizer.from_pretrained(self.encoder_name)
        model = AutoModel.from_pretrained(self.encoder_name).to(self.device)
        model.eval()

        chunk_dir = self.index_dir
        chunk_dir.mkdir(parents=True, exist_ok=True)

        chunk_paths = []

        with open(self.corpus_path, "r", encoding="utf-8") as f:
            for chunk_lines, start_idx in tqdm(self._generate_chunks(f, self.chunk_size), desc="Embedding chunks"):
                chunk_data = self._bge_embed_chunk(chunk_lines, tokenizer, model)

                chunk_path = chunk_dir / f"embeddings_chunk_{start_idx}.pkl"
                with open(chunk_path, "wb") as fout:
                    pickle.dump(chunk_data, fout)

                chunk_paths.append(chunk_path)

        logging.info(f"Saved {len(chunk_paths)} embedding chunk files.")
        return chunk_paths

    def _generate_chunks(self, file_obj, chunk_size):
        """Generate chunks from file with line counting."""
        chunk = []
        start_idx = 0

        for i, line in enumerate(file_obj):
            chunk.append(line.strip())

            if len(chunk) >= chunk_size:
                yield chunk, start_idx
                chunk = []
                start_idx = i + 1

        if chunk:
            yield chunk, start_idx

    def _bge_embed_chunk(self, chunk_lines, tokenizer, model):
        """
        Corrected BGE embedding function with proper text handling.
        """
        chunk_data = []
        texts = []
        doc_ids = []

        # Parse documents from chunk
        for line in chunk_lines:
            if not line.strip():
                continue

            try:
                doc = json.loads(line)
                text = doc.get("text", doc.get("contents", "")).strip()
                doc_id = doc.get("id", f"doc_{len(doc_ids)}")

                if text:  # Only process non-empty texts
                    texts.append(text)
                    doc_ids.append(doc_id)

            except json.JSONDecodeError:
                logging.warning(f"Failed to parse JSON line: {line[:100]}...")
                continue

        if not texts:
            return []

        # Process in batches
        for i in range(0, len(texts), self.batch_size):
            batch_texts = texts[i:i + self.batch_size]
            batch_ids = doc_ids[i:i + self.batch_size]

            # Extract embeddings using corrected method
            embeddings = self._extract_bge_embeddings(batch_texts, tokenizer, model)

            # Store embeddings with IDs
            for doc_id, embedding in zip(batch_ids, embeddings):
                chunk_data.append({
                    "id": doc_id,
                    "embedding": embedding
                })

        return chunk_data

    def _save_corpus(self):
        """Convert the corpus to TSV format for indexing."""
        return to_tsv(self.corpus_path, self.index_dir, self.chunk_size, self.threads)

    def build_index(self):
        """Build the BGE index from the corpus."""
        if self.index_dir.exists():
            shutil.rmtree(self.index_dir)
        self.index_dir.mkdir(parents=True)

        self._save_corpus()
        self._save_dense_embeddings()
        self._merge_embedding_chunks()
        self.build_faiss_index()

        logging.info(f"BGE index stored at {self.index_dir}")

    def load_index(self) -> tuple:
        """Load the FAISS index and document IDs from the index directory."""
        faiss_index_path = self.index_dir / "faiss_index.bin"
        merged_ids_path = self.index_dir / "bge_doc_ids.pkl"

        if not faiss_index_path.exists():
            raise FileNotFoundError(f"FAISS index file not found at {faiss_index_path}")
        if not merged_ids_path.exists():
            raise FileNotFoundError(f"Merged IDs file not found at {merged_ids_path}")

        logging.info(f"Loading FAISS index from {faiss_index_path}...")
        index = faiss.read_index(str(faiss_index_path))

        logging.info(f"Loading document IDs from {merged_ids_path}...")
        with open(merged_ids_path, "rb") as f:
            doc_ids = pickle.load(f)

        logging.info(f"Loaded FAISS index and {len(doc_ids)} document IDs from {self.index_dir}")

        return index, doc_ids

build_faiss_index()

Build FAISS index with normalized embeddings for cosine similarity.

Source code in rankify/indexing/bge_indexer.py

def build_faiss_index(self):
    """Build FAISS index with normalized embeddings for cosine similarity."""
    embeddings_path = self.index_dir / "bge_embeddings.h5"
    logging.info(f"Loading embeddings from {embeddings_path}...")

    with h5py.File(embeddings_path, "r") as f:
        embeddings = f["embeddings"][:].astype(np.float32)

    # Normalize embeddings for cosine similarity
    logging.info("Normalizing embeddings for cosine similarity...")
    embeddings = self._normalize_embeddings(embeddings)

    embedding_dim = embeddings.shape[1]
    logging.info(f"Building FAISS index with embedding dimension {embedding_dim}...")

    # Use IndexFlatIP with normalized embeddings = cosine similarity
    index = faiss.IndexFlatIP(embedding_dim)

    # Add embeddings in batches for memory efficiency
    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])
        index.add(embeddings[start:end])

    index_path = self.index_dir / "faiss_index.bin"
    logging.info(f"Saving FAISS index to {index_path}...")
    faiss.write_index(index, str(index_path))

build_index()

Build the BGE index from the corpus.

Source code in rankify/indexing/bge_indexer.py

def build_index(self):
    """Build the BGE index from the corpus."""
    if self.index_dir.exists():
        shutil.rmtree(self.index_dir)
    self.index_dir.mkdir(parents=True)

    self._save_corpus()
    self._save_dense_embeddings()
    self._merge_embedding_chunks()
    self.build_faiss_index()

    logging.info(f"BGE index stored at {self.index_dir}")

load_index()

Load the FAISS index and document IDs from the index directory.

Source code in rankify/indexing/bge_indexer.py

def load_index(self) -> tuple:
    """Load the FAISS index and document IDs from the index directory."""
    faiss_index_path = self.index_dir / "faiss_index.bin"
    merged_ids_path = self.index_dir / "bge_doc_ids.pkl"

    if not faiss_index_path.exists():
        raise FileNotFoundError(f"FAISS index file not found at {faiss_index_path}")
    if not merged_ids_path.exists():
        raise FileNotFoundError(f"Merged IDs file not found at {merged_ids_path}")

    logging.info(f"Loading FAISS index from {faiss_index_path}...")
    index = faiss.read_index(str(faiss_index_path))

    logging.info(f"Loading document IDs from {merged_ids_path}...")
    with open(merged_ids_path, "rb") as f:
        doc_ids = pickle.load(f)

    logging.info(f"Loaded FAISS index and {len(doc_ids)} document IDs from {self.index_dir}")

    return index, doc_ids