From 82132ea350190367c6f838b71c0a0cfdf8a4e60a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kacper=20=C5=81ukawski?= Date: Tue, 1 Jul 2025 12:20:29 +0200 Subject: [PATCH 01/13] Implement MuveraEmbedding --- fastembed/__init__.py | 2 + fastembed/text/muvera_embedding.py | 464 +++++++++++++++++++++++++++++ fastembed/text/text_embedding.py | 2 + 3 files changed, 468 insertions(+) create mode 100644 fastembed/text/muvera_embedding.py diff --git a/fastembed/__init__.py b/fastembed/__init__.py index 7a2e41340..64b6831e3 100644 --- a/fastembed/__init__.py +++ b/fastembed/__init__.py @@ -4,6 +4,7 @@ from fastembed.late_interaction import LateInteractionTextEmbedding from fastembed.late_interaction_multimodal import LateInteractionMultimodalEmbedding from fastembed.sparse import SparseEmbedding, SparseTextEmbedding +from fastembed.text.muvera_embedding import MuveraEmbedding from fastembed.text import TextEmbedding try: @@ -19,4 +20,5 @@ "ImageEmbedding", "LateInteractionTextEmbedding", "LateInteractionMultimodalEmbedding", + "MuveraEmbedding", ] diff --git a/fastembed/text/muvera_embedding.py b/fastembed/text/muvera_embedding.py new file mode 100644 index 000000000..1a6642cec --- /dev/null +++ b/fastembed/text/muvera_embedding.py @@ -0,0 +1,464 @@ +""" +MUVERA (Multi-Vector Retrieval Architecture) embeddings implementation. + +This module provides a text embedding class that combines late interaction models +(like ColBERT) with the MUVERA algorithm to create fixed-dimensional embeddings +from variable-length token sequences. +""" + +import numpy as np +from typing import Any, Iterable, Optional, Union, Type +from dataclasses import asdict + +from fastembed.common.types import NumpyArray +from fastembed.common.model_description import DenseModelDescription +from fastembed.text.text_embedding_base import TextEmbeddingBase + + +class SimHashProjection: + """ + SimHash projection component for MUVERA clustering. + + This class implements locality-sensitive hashing using random hyperplanes + to partition the vector space into 2^k_sim clusters. Each vector is assigned + to a cluster based on which side of k_sim random hyperplanes it falls on. + + Attributes: + k_sim (int): Number of SimHash functions (hyperplanes) + d (int): Dimensionality of input vectors + simhash_vectors (np.ndarray): Random hyperplane normal vectors of shape (d, k_sim) + """ + + def __init__(self, k_sim: int, d: int): + """ + Initialize SimHash projection with random hyperplanes. + + Args: + k_sim (int): Number of SimHash functions, determines 2^k_sim clusters + d (int): Dimensionality of input vectors + """ + self.k_sim = k_sim + self.d = d + # Generate k_sim random hyperplanes (normal vectors) from standard normal distribution + self.simhash_vectors = np.random.normal(size=(d, k_sim)) + + def get_cluster_id(self, vector: np.ndarray) -> int: + """ + Compute the cluster ID for a given vector using SimHash. + + The cluster ID is determined by computing the dot product of the vector + with each hyperplane normal vector, taking the sign, and interpreting + the resulting binary string as an integer. + + Args: + vector (np.ndarray): Input vector of shape (d,) + + Returns: + int: Cluster ID in range [0, 2^k_sim - 1] + + Raises: + AssertionError: If vector shape doesn't match expected dimensionality + """ + assert vector.shape == (self.d,), f"Expected vector of shape ({self.d},), got {vector.shape}" + + # Project vector onto each hyperplane normal vector + dot_product = np.dot(vector, self.simhash_vectors) + + # Apply sign function to get binary values (1 if positive, 0 if negative) + binary_values = (dot_product > 0).astype(int) + + # Convert binary representation to decimal cluster ID + # Each bit position i contributes bit_value * 2^i to the final ID + cluster_id = 0 + for i, bit in enumerate(binary_values): + cluster_id += bit * (2 ** i) + return cluster_id + + +class MuveraAlgorithm: + """ + MUVERA (Multi-Vector Retrieval Architecture) algorithm implementation. + + This class creates Fixed Dimensional Encodings (FDEs) from variable-length + sequences of vectors by using SimHash clustering and random projections. + The process involves: + 1. Clustering vectors using multiple SimHash projections + 2. Computing cluster centers (with different strategies for docs vs queries) + 3. Applying random projections for dimensionality reduction + 4. Concatenating results from all projections + + Attributes: + k_sim (int): Number of SimHash functions per projection + d (int): Input vector dimensionality + d_proj (int): Output dimensionality after random projection + R_reps (int): Number of random projection repetitions + simhash_projections (List[SimHashProjection]): SimHash instances for clustering + S_projections (np.ndarray): Random projection matrices of shape (R_reps, d, d_proj) + """ + + def __init__(self, k_sim: int, d: int, d_proj: int, R_reps: int): + """ + Initialize MUVERA algorithm with specified parameters. + + Args: + k_sim (int): Number of SimHash functions (creates 2^k_sim clusters) + d (int): Dimensionality of input vectors + d_proj (int): Dimensionality after random projection (must be <= d) + R_reps (int): Number of random projection repetitions for robustness + + Raises: + ValueError: If d_proj > d (cannot project to higher dimensionality) + """ + if d_proj > d: + raise ValueError(f"Cannot project to a higher dimensionality (d_proj={d_proj} > d={d})") + + self.k_sim = k_sim + self.d = d + self.d_proj = d_proj + self.R_reps = R_reps + # Create R_reps independent SimHash projections for robustness + self.simhash_projections = [SimHashProjection(k_sim=self.k_sim, d=self.d) for _ in range(R_reps)] + # Random projection matrices with entries from {-1, +1} for each repetition + self.S_projections = np.random.choice([-1, 1], size=(R_reps, d, d_proj)) + + def get_output_dimension(self) -> int: + """ + Get the output dimension of the MUVERA algorithm. + + Returns: + int: Output dimension (R_reps * B * d_proj) where B = 2^k_sim + """ + B = 2 ** self.k_sim + return self.R_reps * B * self.d_proj + + def encode_document(self, vectors: np.ndarray) -> np.ndarray: + """ + Encode a document's vectors into a Fixed Dimensional Encoding (FDE). + + Uses document-specific settings: normalizes cluster centers by vector count + and fills empty clusters using Hamming distance-based selection. + + Args: + vectors (np.ndarray): Document vectors of shape (n_tokens, d) + + Returns: + np.ndarray: Fixed dimensional encoding of shape (R_reps * B * d_proj,) + """ + return self.encode(vectors, fill_empty_clusters=True, normalize_by_count=True) + + def encode_query(self, vectors: np.ndarray) -> np.ndarray: + """ + Encode a query's vectors into a Fixed Dimensional Encoding (FDE). + + Uses query-specific settings: no normalization by count and no empty + cluster filling to preserve query vector magnitudes. + + Args: + vectors (np.ndarray): Query vectors of shape (n_tokens, d) + + Returns: + np.ndarray: Fixed dimensional encoding of shape (R_reps * B * d_proj,) + """ + return self.encode(vectors, fill_empty_clusters=False, normalize_by_count=False) + + def encode(self, vectors: np.ndarray, fill_empty_clusters: bool = True, normalize_by_count: bool = True) -> np.ndarray: + """ + Core encoding method that transforms variable-length vector sequences into FDEs. + + The encoding process: + 1. For each of R_reps random projections: + a. Assign vectors to clusters using SimHash + b. Compute cluster centers (sum of vectors in each cluster) + c. Optionally normalize by cluster size + d. Fill empty clusters using Hamming distance if requested + e. Apply random projection for dimensionality reduction + f. Flatten cluster centers into a vector + 2. Concatenate all projection results + + Args: + vectors (np.ndarray): Input vectors of shape (n_vectors, d) + fill_empty_clusters (bool): Whether to fill empty clusters using nearest + vectors based on Hamming distance of cluster IDs + normalize_by_count (bool): Whether to normalize cluster centers by the + number of vectors assigned to each cluster + + Returns: + np.ndarray: Fixed dimensional encoding of shape (R_reps * B * d_proj,) + where B = 2^k_sim is the number of clusters + + Raises: + AssertionError: If input vectors don't have expected dimensionality + """ + assert vectors.shape[1] == self.d, f"Expected vectors of shape (n, {self.d}), got {vectors.shape}" + + # Store results from each random projection + output_vectors = [] + + # B is the number of clusters (2^k_sim) + B = 2 ** self.k_sim + for projection_index, simhash in enumerate(self.simhash_projections): + # Initialize cluster centers and count vectors assigned to each cluster + cluster_centers = np.zeros((B, self.d)) + cluster_vector_counts = np.zeros(B) + + # Assign each vector to its cluster and accumulate cluster centers + for vector in vectors: + cluster_id = simhash.get_cluster_id(vector) + cluster_centers[cluster_id] += vector + cluster_vector_counts[cluster_id] += 1 + + # Normalize cluster centers by the number of vectors (for documents) + if normalize_by_count: + for i in range(B): + if cluster_vector_counts[i] == 0: + continue # Skip empty clusters + cluster_centers[i] /= cluster_vector_counts[i] + + # Fill empty clusters using vectors with minimum Hamming distance + if fill_empty_clusters: + for i in range(B): + if cluster_vector_counts[i] == 0: # Empty cluster found + min_hamming = float("inf") + best_vector = None + # Find vector whose cluster ID has minimum Hamming distance to i + for vector in vectors: + vector_cluster_id = simhash.get_cluster_id(vector) + # Hamming distance = number of differing bits in binary representation + hamming_dist = bin(i ^ vector_cluster_id).count("1") + if hamming_dist < min_hamming: + min_hamming = hamming_dist + best_vector = vector + # Assign the best matching vector to the empty cluster + if best_vector is not None: + cluster_centers[i] = best_vector + + # Apply random projection for dimensionality reduction if needed + if self.d_proj < self.d: + S = self.S_projections[projection_index] # Get projection matrix for this repetition + projected_centers = (1 / np.sqrt(self.d_proj)) * np.dot(cluster_centers, S) + + # Flatten cluster centers into a single vector and add to output + output_vectors.append(projected_centers.flatten()) + continue + + # If no projection needed (d_proj == d), use original cluster centers + output_vectors.append(cluster_centers.flatten()) + + # Concatenate results from all R_reps projections into final FDE + return np.concatenate(output_vectors) + + +class MuveraEmbedding(TextEmbeddingBase): + """ + MUVERA (Multi-Vector Retrieval Architecture) text embedding class. + + This class combines late interaction models (like ColBERT) with the MUVERA algorithm + to create fixed-dimensional embeddings from variable-length token sequences. + It's compatible with the fastembed TextEmbedding interface while supporting + multivector capabilities through late interaction models. + + The MUVERA algorithm transforms variable-length token embeddings into fixed-dimensional + embeddings using SimHash clustering and random projections, making it suitable for + traditional dense retrieval systems while preserving the benefits of late interaction. + + Attributes: + late_interaction_model (LateInteractionTextEmbedding): The underlying late interaction model + muvera_algorithm (MuveraAlgorithm): The MUVERA algorithm instance + k_sim (int): Number of SimHash functions (controls clustering) + d_proj (int): Dimensionality after random projection + R_reps (int): Number of random projection repetitions + """ + + def __init__( + self, + model_name: str = "colbert-ir/colbertv2.0", + cache_dir: Optional[str] = None, + threads: Optional[int] = None, + k_sim: int = 4, + d_proj: int = 32, + R_reps: int = 10, + **kwargs: Any, + ): + """ + Initialize MuveraEmbedding with a late interaction model and MUVERA parameters. + + Args: + model_name (str): Name of the late interaction model to use. + Must be a supported late interaction model (ColBERT, JinaColBERT, etc.) + cache_dir (Optional[str]): Cache directory for model files + threads (Optional[int]): Number of threads for model inference + k_sim (int): Number of SimHash functions (creates 2^k_sim clusters). Default: 4 + d_proj (int): Dimensionality after random projection. Default: 32 + R_reps (int): Number of random projection repetitions for robustness. Default: 10 + **kwargs: Additional arguments passed to the late interaction model + + Raises: + ValueError: If the model_name is not a supported late interaction model + """ + super().__init__(model_name, cache_dir, threads, **kwargs) + + # Initialize the late interaction model (import locally to avoid circular imports) + try: + from fastembed.late_interaction.late_interaction_text_embedding import LateInteractionTextEmbedding + self.late_interaction_model = LateInteractionTextEmbedding( + model_name=model_name, + cache_dir=cache_dir, + threads=threads, + **kwargs + ) + except ValueError as e: + raise ValueError( + f"Model {model_name} is not supported as a late interaction model. " + f"Please use a supported late interaction model like 'colbert-ir/colbertv2.0' or 'jinaai/jina-colbert-v2'. " + f"Original error: {e}" + ) + + # Store MUVERA parameters + self.k_sim = k_sim + self.d_proj = d_proj + self.R_reps = R_reps + + # Get the token embedding dimension from the late interaction model + self.token_dim = self.late_interaction_model.embedding_size + + # Initialize MUVERA algorithm + self.muvera_algorithm = MuveraAlgorithm( + k_sim=k_sim, + d=self.token_dim, + d_proj=d_proj, + R_reps=R_reps + ) + + # Cache the output embedding size + self._embedding_size = self.muvera_algorithm.get_output_dimension() + + @property + def embedding_size(self) -> int: + """Get the embedding size of the MUVERA output.""" + return self._embedding_size + + @classmethod + def get_embedding_size(cls, model_name: str, k_sim: int = 4, d_proj: int = 32, R_reps: int = 10) -> int: + """ + Get the embedding size for a given model and MUVERA parameters. + + Args: + model_name (str): Name of the late interaction model + k_sim (int): Number of SimHash functions. Default: 4 + d_proj (int): Dimensionality after random projection. Default: 32 + R_reps (int): Number of random projection repetitions. Default: 10 + + Returns: + int: The size of the MUVERA embedding (R_reps * 2^k_sim * d_proj) + + Raises: + ValueError: If the model name is not found in supported models + """ + # Calculate MUVERA output dimension + B = 2 ** k_sim + return R_reps * B * d_proj + + @classmethod + def _list_supported_models(cls) -> list[DenseModelDescription]: + """ + List supported models (same as late interaction models). + + Returns: + list[DenseModelDescription]: List of supported late interaction models + """ + from fastembed.late_interaction.late_interaction_text_embedding import LateInteractionTextEmbedding + return LateInteractionTextEmbedding._list_supported_models() + + @classmethod + def list_supported_models(cls) -> list[dict[str, Any]]: + """ + Lists the supported models. + + Returns: + list[dict[str, Any]]: A list of dictionaries containing the model information. + """ + return [asdict(model) for model in cls._list_supported_models()] + + def embed( + self, + documents: Union[str, Iterable[str]], + batch_size: int = 256, + parallel: Optional[int] = None, + **kwargs: Any, + ) -> Iterable[NumpyArray]: + """ + Encode a list of documents into MUVERA embeddings. + + This method uses the late interaction model to get token-level embeddings, + then applies the MUVERA algorithm to create fixed-dimensional embeddings. + Documents are encoded with normalization and empty cluster filling. + + Args: + documents: Iterator of documents or single document to embed + batch_size: Batch size for encoding -- higher values will use more memory, but be faster + parallel: If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets. + If 0, use all available cores. + If None, don't use data-parallel processing, use default onnxruntime threading instead. + + Returns: + Iterable[NumpyArray]: List of MUVERA embeddings, one per document + """ + # Handle single string input + if isinstance(documents, str): + documents = [documents] + + # Get token-level embeddings from the late interaction model + token_embeddings = self.late_interaction_model.embed( + documents=documents, + batch_size=batch_size, + parallel=parallel, + **kwargs + ) + + # Apply MUVERA algorithm to each document's token embeddings + for token_embedding in token_embeddings: + # token_embedding shape: (n_tokens, token_dim) + muvera_embedding = self.muvera_algorithm.encode_document(token_embedding) + yield muvera_embedding.astype(np.float32) + + def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]: + """ + Embeds queries using MUVERA algorithm. + + This method uses the late interaction model to get token-level embeddings, + then applies the MUVERA algorithm with query-specific settings (no normalization + and no empty cluster filling) to preserve query vector magnitudes. + + Args: + query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries. + + Returns: + Iterable[NumpyArray]: The MUVERA query embeddings. + """ + # Handle single string input + if isinstance(query, str): + query = [query] + + # Get token-level embeddings from the late interaction model + token_embeddings = self.late_interaction_model.query_embed(query, **kwargs) + + # Apply MUVERA algorithm to each query's token embeddings + for token_embedding in token_embeddings: + # token_embedding shape: (n_tokens, token_dim) + muvera_embedding = self.muvera_algorithm.encode_query(token_embedding) + yield muvera_embedding.astype(np.float32) + + def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]: + """ + Embeds a list of text passages into MUVERA embeddings. + + This is an alias for the embed method, following the fastembed interface. + + Args: + texts (Iterable[str]): The list of texts to embed. + **kwargs: Additional keyword arguments to pass to the embed method. + + Yields: + Iterable[NumpyArray]: The MUVERA embeddings. + """ + yield from self.embed(texts, **kwargs) \ No newline at end of file diff --git a/fastembed/text/text_embedding.py b/fastembed/text/text_embedding.py index 117f5af79..de60d7e6b 100644 --- a/fastembed/text/text_embedding.py +++ b/fastembed/text/text_embedding.py @@ -8,6 +8,7 @@ from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding from fastembed.text.pooled_embedding import PooledEmbedding from fastembed.text.multitask_embedding import JinaEmbeddingV3 +from fastembed.text.muvera_embedding import MuveraEmbedding from fastembed.text.onnx_embedding import OnnxTextEmbedding from fastembed.text.text_embedding_base import TextEmbeddingBase from fastembed.common.model_description import DenseModelDescription, ModelSource, PoolingType @@ -20,6 +21,7 @@ class TextEmbedding(TextEmbeddingBase): PooledNormalizedEmbedding, PooledEmbedding, JinaEmbeddingV3, + MuveraEmbedding, CustomTextEmbedding, ] From 8734bc498b4a72559ae4e5b1236d45282355cb8a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kacper=20=C5=81ukawski?= Date: Mon, 7 Jul 2025 16:22:00 +0200 Subject: [PATCH 02/13] Add random generator parameter for reproducibility in MuveraEmbedding --- fastembed/text/muvera_embedding.py | 29 +++++++++++++++++++++++------ 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/fastembed/text/muvera_embedding.py b/fastembed/text/muvera_embedding.py index 1a6642cec..e0b7f4449 100644 --- a/fastembed/text/muvera_embedding.py +++ b/fastembed/text/muvera_embedding.py @@ -29,18 +29,24 @@ class SimHashProjection: simhash_vectors (np.ndarray): Random hyperplane normal vectors of shape (d, k_sim) """ - def __init__(self, k_sim: int, d: int): + def __init__( + self, + k_sim: int, + d: int, + random_generator: np.random.Generator = np.random + ): """ Initialize SimHash projection with random hyperplanes. Args: k_sim (int): Number of SimHash functions, determines 2^k_sim clusters d (int): Dimensionality of input vectors + random_generator (np.random.Generator): Random number random_generator for reproducibility """ self.k_sim = k_sim self.d = d # Generate k_sim random hyperplanes (normal vectors) from standard normal distribution - self.simhash_vectors = np.random.normal(size=(d, k_sim)) + self.simhash_vectors = random_generator.normal(size=(d, k_sim)) def get_cluster_id(self, vector: np.ndarray) -> int: """ @@ -96,7 +102,14 @@ class MuveraAlgorithm: S_projections (np.ndarray): Random projection matrices of shape (R_reps, d, d_proj) """ - def __init__(self, k_sim: int, d: int, d_proj: int, R_reps: int): + def __init__( + self, + k_sim: int, + d: int, + d_proj: int, + R_reps: int, + random_generator: np.random.Generator = np.random + ): """ Initialize MUVERA algorithm with specified parameters. @@ -105,6 +118,7 @@ def __init__(self, k_sim: int, d: int, d_proj: int, R_reps: int): d (int): Dimensionality of input vectors d_proj (int): Dimensionality after random projection (must be <= d) R_reps (int): Number of random projection repetitions for robustness + random_generator (np.random.Generator): Random number random_generator for reproducibility Raises: ValueError: If d_proj > d (cannot project to higher dimensionality) @@ -117,9 +131,9 @@ def __init__(self, k_sim: int, d: int, d_proj: int, R_reps: int): self.d_proj = d_proj self.R_reps = R_reps # Create R_reps independent SimHash projections for robustness - self.simhash_projections = [SimHashProjection(k_sim=self.k_sim, d=self.d) for _ in range(R_reps)] + self.simhash_projections = [SimHashProjection(k_sim=self.k_sim, d=self.d, random_generator=random_generator) for _ in range(R_reps)] # Random projection matrices with entries from {-1, +1} for each repetition - self.S_projections = np.random.choice([-1, 1], size=(R_reps, d, d_proj)) + self.S_projections = random_generator.choice([-1, 1], size=(R_reps, d, d_proj)) def get_output_dimension(self) -> int: """ @@ -277,6 +291,7 @@ def __init__( k_sim: int = 4, d_proj: int = 32, R_reps: int = 10, + random_seed: Optional[int] = None, **kwargs: Any, ): """ @@ -322,11 +337,13 @@ def __init__( self.token_dim = self.late_interaction_model.embedding_size # Initialize MUVERA algorithm + generator = np.random.default_rng(random_seed) if random_seed is not None else np.random self.muvera_algorithm = MuveraAlgorithm( k_sim=k_sim, d=self.token_dim, d_proj=d_proj, - R_reps=R_reps + R_reps=R_reps, + random_generator=generator, ) # Cache the output embedding size From 6d3569d3e3132b31280b24b6e2c962591b663a53 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kacper=20=C5=81ukawski?= Date: Mon, 7 Jul 2025 16:23:43 +0200 Subject: [PATCH 03/13] Document random_seed parameter --- fastembed/text/muvera_embedding.py | 1 + 1 file changed, 1 insertion(+) diff --git a/fastembed/text/muvera_embedding.py b/fastembed/text/muvera_embedding.py index e0b7f4449..67074346b 100644 --- a/fastembed/text/muvera_embedding.py +++ b/fastembed/text/muvera_embedding.py @@ -305,6 +305,7 @@ def __init__( k_sim (int): Number of SimHash functions (creates 2^k_sim clusters). Default: 4 d_proj (int): Dimensionality after random projection. Default: 32 R_reps (int): Number of random projection repetitions for robustness. Default: 10 + random_seed (Optional[int]): Random seed for reproducibility. Default: None **kwargs: Additional arguments passed to the late interaction model Raises: From 1e4ac4d7dd72c41607e0e512ef36bc190e926c01 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kacper=20=C5=81ukawski?= Date: Mon, 7 Jul 2025 16:52:30 +0200 Subject: [PATCH 04/13] Remove unnecessary module docstring from muvera_embedding.py --- fastembed/text/muvera_embedding.py | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/fastembed/text/muvera_embedding.py b/fastembed/text/muvera_embedding.py index 67074346b..d83d7b6d2 100644 --- a/fastembed/text/muvera_embedding.py +++ b/fastembed/text/muvera_embedding.py @@ -1,13 +1,5 @@ -""" -MUVERA (Multi-Vector Retrieval Architecture) embeddings implementation. - -This module provides a text embedding class that combines late interaction models -(like ColBERT) with the MUVERA algorithm to create fixed-dimensional embeddings -from variable-length token sequences. -""" - import numpy as np -from typing import Any, Iterable, Optional, Union, Type +from typing import Any, Iterable, Optional, Union from dataclasses import asdict from fastembed.common.types import NumpyArray From bb913e5320e1379a1b92aeb5399a3d2945347dc0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kacper=20=C5=81ukawski?= Date: Wed, 9 Jul 2025 14:02:38 +0200 Subject: [PATCH 05/13] refactor: clean up constructor parameters and improve formatting in MuveraEmbedding --- fastembed/text/muvera_embedding.py | 75 ++++++++++++++++++------------ 1 file changed, 44 insertions(+), 31 deletions(-) diff --git a/fastembed/text/muvera_embedding.py b/fastembed/text/muvera_embedding.py index d83d7b6d2..c167edcc6 100644 --- a/fastembed/text/muvera_embedding.py +++ b/fastembed/text/muvera_embedding.py @@ -21,12 +21,7 @@ class SimHashProjection: simhash_vectors (np.ndarray): Random hyperplane normal vectors of shape (d, k_sim) """ - def __init__( - self, - k_sim: int, - d: int, - random_generator: np.random.Generator = np.random - ): + def __init__(self, k_sim: int, d: int, random_generator: np.random.Generator): """ Initialize SimHash projection with random hyperplanes. @@ -57,7 +52,9 @@ def get_cluster_id(self, vector: np.ndarray) -> int: Raises: AssertionError: If vector shape doesn't match expected dimensionality """ - assert vector.shape == (self.d,), f"Expected vector of shape ({self.d},), got {vector.shape}" + assert vector.shape == ( + self.d, + ), f"Expected vector of shape ({self.d},), got {vector.shape}" # Project vector onto each hyperplane normal vector dot_product = np.dot(vector, self.simhash_vectors) @@ -69,7 +66,7 @@ def get_cluster_id(self, vector: np.ndarray) -> int: # Each bit position i contributes bit_value * 2^i to the final ID cluster_id = 0 for i, bit in enumerate(binary_values): - cluster_id += bit * (2 ** i) + cluster_id += bit * (2**i) return cluster_id @@ -100,7 +97,7 @@ def __init__( d: int, d_proj: int, R_reps: int, - random_generator: np.random.Generator = np.random + random_generator: np.random.Generator, ): """ Initialize MUVERA algorithm with specified parameters. @@ -116,14 +113,19 @@ def __init__( ValueError: If d_proj > d (cannot project to higher dimensionality) """ if d_proj > d: - raise ValueError(f"Cannot project to a higher dimensionality (d_proj={d_proj} > d={d})") + raise ValueError( + f"Cannot project to a higher dimensionality (d_proj={d_proj} > d={d})" + ) self.k_sim = k_sim self.d = d self.d_proj = d_proj self.R_reps = R_reps # Create R_reps independent SimHash projections for robustness - self.simhash_projections = [SimHashProjection(k_sim=self.k_sim, d=self.d, random_generator=random_generator) for _ in range(R_reps)] + self.simhash_projections = [ + SimHashProjection(k_sim=self.k_sim, d=self.d, random_generator=random_generator) + for _ in range(R_reps) + ] # Random projection matrices with entries from {-1, +1} for each repetition self.S_projections = random_generator.choice([-1, 1], size=(R_reps, d, d_proj)) @@ -134,7 +136,7 @@ def get_output_dimension(self) -> int: Returns: int: Output dimension (R_reps * B * d_proj) where B = 2^k_sim """ - B = 2 ** self.k_sim + B = 2**self.k_sim return self.R_reps * B * self.d_proj def encode_document(self, vectors: np.ndarray) -> np.ndarray: @@ -167,7 +169,12 @@ def encode_query(self, vectors: np.ndarray) -> np.ndarray: """ return self.encode(vectors, fill_empty_clusters=False, normalize_by_count=False) - def encode(self, vectors: np.ndarray, fill_empty_clusters: bool = True, normalize_by_count: bool = True) -> np.ndarray: + def encode( + self, + vectors: np.ndarray, + fill_empty_clusters: bool = True, + normalize_by_count: bool = True, + ) -> np.ndarray: """ Core encoding method that transforms variable-length vector sequences into FDEs. @@ -195,13 +202,15 @@ def encode(self, vectors: np.ndarray, fill_empty_clusters: bool = True, normaliz Raises: AssertionError: If input vectors don't have expected dimensionality """ - assert vectors.shape[1] == self.d, f"Expected vectors of shape (n, {self.d}), got {vectors.shape}" + assert ( + vectors.shape[1] == self.d + ), f"Expected vectors of shape (n, {self.d}), got {vectors.shape}" # Store results from each random projection output_vectors = [] # B is the number of clusters (2^k_sim) - B = 2 ** self.k_sim + B = 2**self.k_sim for projection_index, simhash in enumerate(self.simhash_projections): # Initialize cluster centers and count vectors assigned to each cluster cluster_centers = np.zeros((B, self.d)) @@ -240,7 +249,9 @@ def encode(self, vectors: np.ndarray, fill_empty_clusters: bool = True, normaliz # Apply random projection for dimensionality reduction if needed if self.d_proj < self.d: - S = self.S_projections[projection_index] # Get projection matrix for this repetition + S = self.S_projections[ + projection_index + ] # Get projection matrix for this repetition projected_centers = (1 / np.sqrt(self.d_proj)) * np.dot(cluster_centers, S) # Flatten cluster centers into a single vector and add to output @@ -307,12 +318,12 @@ def __init__( # Initialize the late interaction model (import locally to avoid circular imports) try: - from fastembed.late_interaction.late_interaction_text_embedding import LateInteractionTextEmbedding + from fastembed.late_interaction.late_interaction_text_embedding import ( + LateInteractionTextEmbedding, + ) + self.late_interaction_model = LateInteractionTextEmbedding( - model_name=model_name, - cache_dir=cache_dir, - threads=threads, - **kwargs + model_name=model_name, cache_dir=cache_dir, threads=threads, **kwargs ) except ValueError as e: raise ValueError( @@ -330,7 +341,7 @@ def __init__( self.token_dim = self.late_interaction_model.embedding_size # Initialize MUVERA algorithm - generator = np.random.default_rng(random_seed) if random_seed is not None else np.random + generator = np.random.default_rng(random_seed) self.muvera_algorithm = MuveraAlgorithm( k_sim=k_sim, d=self.token_dim, @@ -340,7 +351,7 @@ def __init__( ) # Cache the output embedding size - self._embedding_size = self.muvera_algorithm.get_output_dimension() + self._embedding_size: int = self.muvera_algorithm.get_output_dimension() @property def embedding_size(self) -> int: @@ -348,7 +359,9 @@ def embedding_size(self) -> int: return self._embedding_size @classmethod - def get_embedding_size(cls, model_name: str, k_sim: int = 4, d_proj: int = 32, R_reps: int = 10) -> int: + def get_embedding_size( + cls, model_name: str, k_sim: int = 4, d_proj: int = 32, R_reps: int = 10 + ) -> int: """ Get the embedding size for a given model and MUVERA parameters. @@ -365,7 +378,7 @@ def get_embedding_size(cls, model_name: str, k_sim: int = 4, d_proj: int = 32, R ValueError: If the model name is not found in supported models """ # Calculate MUVERA output dimension - B = 2 ** k_sim + B = 2**k_sim return R_reps * B * d_proj @classmethod @@ -376,7 +389,10 @@ def _list_supported_models(cls) -> list[DenseModelDescription]: Returns: list[DenseModelDescription]: List of supported late interaction models """ - from fastembed.late_interaction.late_interaction_text_embedding import LateInteractionTextEmbedding + from fastembed.late_interaction.late_interaction_text_embedding import ( + LateInteractionTextEmbedding, + ) + return LateInteractionTextEmbedding._list_supported_models() @classmethod @@ -419,10 +435,7 @@ def embed( # Get token-level embeddings from the late interaction model token_embeddings = self.late_interaction_model.embed( - documents=documents, - batch_size=batch_size, - parallel=parallel, - **kwargs + documents=documents, batch_size=batch_size, parallel=parallel, **kwargs ) # Apply MUVERA algorithm to each document's token embeddings @@ -471,4 +484,4 @@ def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyAr Yields: Iterable[NumpyArray]: The MUVERA embeddings. """ - yield from self.embed(texts, **kwargs) \ No newline at end of file + yield from self.embed(texts, **kwargs) From f4735cfa0eba1f5757923af13b990b4aeee0809a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kacper=20=C5=81ukawski?= Date: Fri, 11 Jul 2025 12:06:19 +0200 Subject: [PATCH 06/13] refactor: rename muvera_embedding.py to muvera.py and update related references --- fastembed/__init__.py | 2 - fastembed/postprocess/__init__.py | 0 fastembed/postprocess/muvera.py | 262 ++++++++++++++++ fastembed/text/muvera_embedding.py | 487 ----------------------------- fastembed/text/text_embedding.py | 2 - tests/test_postprocess.py | 0 6 files changed, 262 insertions(+), 491 deletions(-) create mode 100644 fastembed/postprocess/__init__.py create mode 100644 fastembed/postprocess/muvera.py delete mode 100644 fastembed/text/muvera_embedding.py create mode 100644 tests/test_postprocess.py diff --git a/fastembed/__init__.py b/fastembed/__init__.py index 64b6831e3..7a2e41340 100644 --- a/fastembed/__init__.py +++ b/fastembed/__init__.py @@ -4,7 +4,6 @@ from fastembed.late_interaction import LateInteractionTextEmbedding from fastembed.late_interaction_multimodal import LateInteractionMultimodalEmbedding from fastembed.sparse import SparseEmbedding, SparseTextEmbedding -from fastembed.text.muvera_embedding import MuveraEmbedding from fastembed.text import TextEmbedding try: @@ -20,5 +19,4 @@ "ImageEmbedding", "LateInteractionTextEmbedding", "LateInteractionMultimodalEmbedding", - "MuveraEmbedding", ] diff --git a/fastembed/postprocess/__init__.py b/fastembed/postprocess/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/fastembed/postprocess/muvera.py b/fastembed/postprocess/muvera.py new file mode 100644 index 000000000..05ee52fe8 --- /dev/null +++ b/fastembed/postprocess/muvera.py @@ -0,0 +1,262 @@ +import numpy as np + +from fastembed.common.types import NumpyArray + + +class SimHashProjection: + """ + SimHash projection component for MUVERA clustering. + + This class implements locality-sensitive hashing using random hyperplanes + to partition the vector space into 2^k_sim clusters. Each vector is assigned + to a cluster based on which side of k_sim random hyperplanes it falls on. + + Attributes: + k_sim (int): Number of SimHash functions (hyperplanes) + d (int): Dimensionality of input vectors + simhash_vectors (np.ndarray): Random hyperplane normal vectors of shape (d, k_sim) + """ + + def __init__(self, k_sim: int, d: int, random_generator: np.random.Generator): + """ + Initialize SimHash projection with random hyperplanes. + + Args: + k_sim (int): Number of SimHash functions, determines 2^k_sim clusters + d (int): Dimensionality of input vectors + random_generator (np.random.Generator): Random number generator for reproducibility + """ + self.k_sim = k_sim + self.d = d + # Generate k_sim random hyperplanes (normal vectors) from standard normal distribution + self.simhash_vectors = random_generator.normal(size=(d, k_sim)) + + def get_cluster_id(self, vector: np.ndarray) -> int: + """ + Compute the cluster ID for a given vector using SimHash. + + The cluster ID is determined by computing the dot product of the vector + with each hyperplane normal vector, taking the sign, and interpreting + the resulting binary string as an integer. + + Args: + vector (np.ndarray): Input vector of shape (d,) + + Returns: + int: Cluster ID in range [0, 2^k_sim - 1] + + Raises: + AssertionError: If a vector shape doesn't match expected dimensionality + """ + assert vector.shape == ( + self.d, + ), f"Expected vector of shape ({self.d},), got {vector.shape}" + + # Project vector onto each hyperplane normal vector + dot_product = np.dot(vector, self.simhash_vectors) + + # Apply sign function to get binary values (1 if positive, 0 if negative) + binary_values = (dot_product > 0).astype(int) + + # Convert binary representation to decimal cluster ID + # Each bit position i contributes bit_value * 2^i to the final ID + cluster_id = 0 + for i, bit in enumerate(binary_values): + cluster_id += bit * (2**i) + return cluster_id + + +class MuveraPostprocessor: + """ + MUVERA (Multi-Vector Retrieval Architecture) algorithm implementation. + + This class creates Fixed Dimensional Encodings (FDEs) from variable-length + sequences of vectors by using SimHash clustering and random projections. + The process involves: + 1. Clustering vectors using multiple SimHash projections + 2. Computing cluster centers (with different strategies for docs vs queries) + 3. Applying random projections for dimensionality reduction + 4. Concatenating results from all projections + + Attributes: + k_sim (int): Number of SimHash functions per projection + d (int): Input vector dimensionality + d_proj (int): Output dimensionality after random projection + R_reps (int): Number of random projection repetitions + simhash_projections (List[SimHashProjection]): SimHash instances for clustering + S_projections (np.ndarray): Random projection matrices of shape (R_reps, d, d_proj) + """ + + def __init__( + self, + k_sim: int, + d: int, + d_proj: int, + R_reps: int, + random_seed: np.random.Generator = 42, + ): + """ + Initialize MUVERA algorithm with specified parameters. + + Args: + k_sim (int): Number of SimHash functions (creates 2^k_sim clusters) + d (int): Dimensionality of input vectors + d_proj (int): Dimensionality after random projection (must be <= d) + R_reps (int): Number of random projection repetitions for robustness + random_seed (np.random.Generator): Seed for random generator + + Raises: + ValueError: If d_proj > d (cannot project to higher dimensionality) + """ + if d_proj > d: + raise ValueError( + f"Cannot project to a higher dimensionality (d_proj={d_proj} > d={d})" + ) + + self.k_sim = k_sim + self.d = d + self.d_proj = d_proj + self.R_reps = R_reps + # Create R_reps independent SimHash projections for robustness + generator = np.random.default_rng(random_seed) + self.simhash_projections = [ + SimHashProjection(k_sim=self.k_sim, d=self.d, random_generator=generator) + for _ in range(R_reps) + ] + # Random projection matrices with entries from {-1, +1} for each repetition + self.S_projections = random_seed.choice([-1, 1], size=(R_reps, d, d_proj)) + + def get_output_dimension(self) -> int: + """ + Get the output dimension of the MUVERA algorithm. + + Returns: + int: Output dimension (R_reps * B * d_proj) where B = 2^k_sim + """ + B = 2**self.k_sim + return self.R_reps * B * self.d_proj + + def process_document(self, vectors: np.ndarray) -> np.ndarray: + """ + Encode a document's vectors into a Fixed Dimensional Encoding (FDE). + + Uses document-specific settings: normalizes cluster centers by vector count + and fills empty clusters using Hamming distance-based selection. + + Args: + vectors (np.ndarray): Document vectors of shape (n_tokens, d) + + Returns: + np.ndarray: Fixed dimensional encoding of shape (R_reps * B * d_proj,) + """ + return self.process(vectors, fill_empty_clusters=True, normalize_by_count=True) + + def process_query(self, vectors: np.ndarray) -> np.ndarray: + """ + Encode a query's vectors into a Fixed Dimensional Encoding (FDE). + + Uses query-specific settings: no normalization by count and no empty + cluster filling to preserve query vector magnitudes. + + Args: + vectors (np.ndarray): Query vectors of shape (n_tokens, d) + + Returns: + np.ndarray: Fixed dimensional encoding of shape (R_reps * B * d_proj,) + """ + return self.process(vectors, fill_empty_clusters=False, normalize_by_count=False) + + def process( + self, + vectors: NumpyArray, + fill_empty_clusters: bool = True, + normalize_by_count: bool = True, + ) -> NumpyArray: + """ + Core encoding method that transforms variable-length vector sequences into FDEs. + + The encoding process: + 1. For each of R_reps random projections: + a. Assign vectors to clusters using SimHash + b. Compute cluster centers (sum of vectors in each cluster) + c. Optionally normalize by cluster size + d. Fill empty clusters using Hamming distance if requested + e. Apply random projection for dimensionality reduction + f. Flatten cluster centers into a vector + 2. Concatenate all projection results + + Args: + vectors (np.ndarray): Input vectors of shape (n_vectors, d) + fill_empty_clusters (bool): Whether to fill empty clusters using nearest + vectors based on Hamming distance of cluster IDs + normalize_by_count (bool): Whether to normalize cluster centers by the + number of vectors assigned to each cluster + + Returns: + np.ndarray: Fixed dimensional encoding of shape (R_reps * B * d_proj) + where B = 2^k_sim is the number of clusters + + Raises: + AssertionError: If input vectors don't have expected dimensionality + """ + assert ( + vectors.shape[1] == self.d + ), f"Expected vectors of shape (n, {self.d}), got {vectors.shape}" + + # Store results from each random projection + output_vectors = [] + + # B is the number of clusters (2^k_sim) + B = 2**self.k_sim + for projection_index, simhash in enumerate(self.simhash_projections): + # Initialize cluster centers and count vectors assigned to each cluster + cluster_centers = np.zeros((B, self.d)) + cluster_vector_counts = np.zeros(B) + + # Assign each vector to its cluster and accumulate cluster centers + for vector in vectors: + cluster_id = simhash.get_cluster_id(vector) + cluster_centers[cluster_id] += vector + cluster_vector_counts[cluster_id] += 1 + + # Normalize cluster centers by the number of vectors (for documents) + if normalize_by_count: + for i in range(B): + if cluster_vector_counts[i] == 0: + continue # Skip empty clusters + cluster_centers[i] /= cluster_vector_counts[i] + + # Fill empty clusters using vectors with minimum Hamming distance + if fill_empty_clusters: + for i in range(B): + if cluster_vector_counts[i] == 0: # Empty cluster found + min_hamming = float("inf") + best_vector = None + # Find vector whose cluster ID has minimum Hamming distance to i + for vector in vectors: + vector_cluster_id = simhash.get_cluster_id(vector) + # Hamming distance = number of differing bits in binary representation + hamming_dist = bin(i ^ vector_cluster_id).count("1") + if hamming_dist < min_hamming: + min_hamming = hamming_dist + best_vector = vector + # Assign the best matching vector to the empty cluster + if best_vector is not None: + cluster_centers[i] = best_vector + + # Apply random projection for dimensionality reduction if needed + if self.d_proj < self.d: + S = self.S_projections[ + projection_index + ] # Get projection matrix for this repetition + projected_centers = (1 / np.sqrt(self.d_proj)) * np.dot(cluster_centers, S) + + # Flatten cluster centers into a single vector and add to output + output_vectors.append(projected_centers.flatten()) + continue + + # If no projection needed (d_proj == d), use original cluster centers + output_vectors.append(cluster_centers.flatten()) + + # Concatenate results from all R_reps projections into final FDE + return np.concatenate(output_vectors) diff --git a/fastembed/text/muvera_embedding.py b/fastembed/text/muvera_embedding.py deleted file mode 100644 index c167edcc6..000000000 --- a/fastembed/text/muvera_embedding.py +++ /dev/null @@ -1,487 +0,0 @@ -import numpy as np -from typing import Any, Iterable, Optional, Union -from dataclasses import asdict - -from fastembed.common.types import NumpyArray -from fastembed.common.model_description import DenseModelDescription -from fastembed.text.text_embedding_base import TextEmbeddingBase - - -class SimHashProjection: - """ - SimHash projection component for MUVERA clustering. - - This class implements locality-sensitive hashing using random hyperplanes - to partition the vector space into 2^k_sim clusters. Each vector is assigned - to a cluster based on which side of k_sim random hyperplanes it falls on. - - Attributes: - k_sim (int): Number of SimHash functions (hyperplanes) - d (int): Dimensionality of input vectors - simhash_vectors (np.ndarray): Random hyperplane normal vectors of shape (d, k_sim) - """ - - def __init__(self, k_sim: int, d: int, random_generator: np.random.Generator): - """ - Initialize SimHash projection with random hyperplanes. - - Args: - k_sim (int): Number of SimHash functions, determines 2^k_sim clusters - d (int): Dimensionality of input vectors - random_generator (np.random.Generator): Random number random_generator for reproducibility - """ - self.k_sim = k_sim - self.d = d - # Generate k_sim random hyperplanes (normal vectors) from standard normal distribution - self.simhash_vectors = random_generator.normal(size=(d, k_sim)) - - def get_cluster_id(self, vector: np.ndarray) -> int: - """ - Compute the cluster ID for a given vector using SimHash. - - The cluster ID is determined by computing the dot product of the vector - with each hyperplane normal vector, taking the sign, and interpreting - the resulting binary string as an integer. - - Args: - vector (np.ndarray): Input vector of shape (d,) - - Returns: - int: Cluster ID in range [0, 2^k_sim - 1] - - Raises: - AssertionError: If vector shape doesn't match expected dimensionality - """ - assert vector.shape == ( - self.d, - ), f"Expected vector of shape ({self.d},), got {vector.shape}" - - # Project vector onto each hyperplane normal vector - dot_product = np.dot(vector, self.simhash_vectors) - - # Apply sign function to get binary values (1 if positive, 0 if negative) - binary_values = (dot_product > 0).astype(int) - - # Convert binary representation to decimal cluster ID - # Each bit position i contributes bit_value * 2^i to the final ID - cluster_id = 0 - for i, bit in enumerate(binary_values): - cluster_id += bit * (2**i) - return cluster_id - - -class MuveraAlgorithm: - """ - MUVERA (Multi-Vector Retrieval Architecture) algorithm implementation. - - This class creates Fixed Dimensional Encodings (FDEs) from variable-length - sequences of vectors by using SimHash clustering and random projections. - The process involves: - 1. Clustering vectors using multiple SimHash projections - 2. Computing cluster centers (with different strategies for docs vs queries) - 3. Applying random projections for dimensionality reduction - 4. Concatenating results from all projections - - Attributes: - k_sim (int): Number of SimHash functions per projection - d (int): Input vector dimensionality - d_proj (int): Output dimensionality after random projection - R_reps (int): Number of random projection repetitions - simhash_projections (List[SimHashProjection]): SimHash instances for clustering - S_projections (np.ndarray): Random projection matrices of shape (R_reps, d, d_proj) - """ - - def __init__( - self, - k_sim: int, - d: int, - d_proj: int, - R_reps: int, - random_generator: np.random.Generator, - ): - """ - Initialize MUVERA algorithm with specified parameters. - - Args: - k_sim (int): Number of SimHash functions (creates 2^k_sim clusters) - d (int): Dimensionality of input vectors - d_proj (int): Dimensionality after random projection (must be <= d) - R_reps (int): Number of random projection repetitions for robustness - random_generator (np.random.Generator): Random number random_generator for reproducibility - - Raises: - ValueError: If d_proj > d (cannot project to higher dimensionality) - """ - if d_proj > d: - raise ValueError( - f"Cannot project to a higher dimensionality (d_proj={d_proj} > d={d})" - ) - - self.k_sim = k_sim - self.d = d - self.d_proj = d_proj - self.R_reps = R_reps - # Create R_reps independent SimHash projections for robustness - self.simhash_projections = [ - SimHashProjection(k_sim=self.k_sim, d=self.d, random_generator=random_generator) - for _ in range(R_reps) - ] - # Random projection matrices with entries from {-1, +1} for each repetition - self.S_projections = random_generator.choice([-1, 1], size=(R_reps, d, d_proj)) - - def get_output_dimension(self) -> int: - """ - Get the output dimension of the MUVERA algorithm. - - Returns: - int: Output dimension (R_reps * B * d_proj) where B = 2^k_sim - """ - B = 2**self.k_sim - return self.R_reps * B * self.d_proj - - def encode_document(self, vectors: np.ndarray) -> np.ndarray: - """ - Encode a document's vectors into a Fixed Dimensional Encoding (FDE). - - Uses document-specific settings: normalizes cluster centers by vector count - and fills empty clusters using Hamming distance-based selection. - - Args: - vectors (np.ndarray): Document vectors of shape (n_tokens, d) - - Returns: - np.ndarray: Fixed dimensional encoding of shape (R_reps * B * d_proj,) - """ - return self.encode(vectors, fill_empty_clusters=True, normalize_by_count=True) - - def encode_query(self, vectors: np.ndarray) -> np.ndarray: - """ - Encode a query's vectors into a Fixed Dimensional Encoding (FDE). - - Uses query-specific settings: no normalization by count and no empty - cluster filling to preserve query vector magnitudes. - - Args: - vectors (np.ndarray): Query vectors of shape (n_tokens, d) - - Returns: - np.ndarray: Fixed dimensional encoding of shape (R_reps * B * d_proj,) - """ - return self.encode(vectors, fill_empty_clusters=False, normalize_by_count=False) - - def encode( - self, - vectors: np.ndarray, - fill_empty_clusters: bool = True, - normalize_by_count: bool = True, - ) -> np.ndarray: - """ - Core encoding method that transforms variable-length vector sequences into FDEs. - - The encoding process: - 1. For each of R_reps random projections: - a. Assign vectors to clusters using SimHash - b. Compute cluster centers (sum of vectors in each cluster) - c. Optionally normalize by cluster size - d. Fill empty clusters using Hamming distance if requested - e. Apply random projection for dimensionality reduction - f. Flatten cluster centers into a vector - 2. Concatenate all projection results - - Args: - vectors (np.ndarray): Input vectors of shape (n_vectors, d) - fill_empty_clusters (bool): Whether to fill empty clusters using nearest - vectors based on Hamming distance of cluster IDs - normalize_by_count (bool): Whether to normalize cluster centers by the - number of vectors assigned to each cluster - - Returns: - np.ndarray: Fixed dimensional encoding of shape (R_reps * B * d_proj,) - where B = 2^k_sim is the number of clusters - - Raises: - AssertionError: If input vectors don't have expected dimensionality - """ - assert ( - vectors.shape[1] == self.d - ), f"Expected vectors of shape (n, {self.d}), got {vectors.shape}" - - # Store results from each random projection - output_vectors = [] - - # B is the number of clusters (2^k_sim) - B = 2**self.k_sim - for projection_index, simhash in enumerate(self.simhash_projections): - # Initialize cluster centers and count vectors assigned to each cluster - cluster_centers = np.zeros((B, self.d)) - cluster_vector_counts = np.zeros(B) - - # Assign each vector to its cluster and accumulate cluster centers - for vector in vectors: - cluster_id = simhash.get_cluster_id(vector) - cluster_centers[cluster_id] += vector - cluster_vector_counts[cluster_id] += 1 - - # Normalize cluster centers by the number of vectors (for documents) - if normalize_by_count: - for i in range(B): - if cluster_vector_counts[i] == 0: - continue # Skip empty clusters - cluster_centers[i] /= cluster_vector_counts[i] - - # Fill empty clusters using vectors with minimum Hamming distance - if fill_empty_clusters: - for i in range(B): - if cluster_vector_counts[i] == 0: # Empty cluster found - min_hamming = float("inf") - best_vector = None - # Find vector whose cluster ID has minimum Hamming distance to i - for vector in vectors: - vector_cluster_id = simhash.get_cluster_id(vector) - # Hamming distance = number of differing bits in binary representation - hamming_dist = bin(i ^ vector_cluster_id).count("1") - if hamming_dist < min_hamming: - min_hamming = hamming_dist - best_vector = vector - # Assign the best matching vector to the empty cluster - if best_vector is not None: - cluster_centers[i] = best_vector - - # Apply random projection for dimensionality reduction if needed - if self.d_proj < self.d: - S = self.S_projections[ - projection_index - ] # Get projection matrix for this repetition - projected_centers = (1 / np.sqrt(self.d_proj)) * np.dot(cluster_centers, S) - - # Flatten cluster centers into a single vector and add to output - output_vectors.append(projected_centers.flatten()) - continue - - # If no projection needed (d_proj == d), use original cluster centers - output_vectors.append(cluster_centers.flatten()) - - # Concatenate results from all R_reps projections into final FDE - return np.concatenate(output_vectors) - - -class MuveraEmbedding(TextEmbeddingBase): - """ - MUVERA (Multi-Vector Retrieval Architecture) text embedding class. - - This class combines late interaction models (like ColBERT) with the MUVERA algorithm - to create fixed-dimensional embeddings from variable-length token sequences. - It's compatible with the fastembed TextEmbedding interface while supporting - multivector capabilities through late interaction models. - - The MUVERA algorithm transforms variable-length token embeddings into fixed-dimensional - embeddings using SimHash clustering and random projections, making it suitable for - traditional dense retrieval systems while preserving the benefits of late interaction. - - Attributes: - late_interaction_model (LateInteractionTextEmbedding): The underlying late interaction model - muvera_algorithm (MuveraAlgorithm): The MUVERA algorithm instance - k_sim (int): Number of SimHash functions (controls clustering) - d_proj (int): Dimensionality after random projection - R_reps (int): Number of random projection repetitions - """ - - def __init__( - self, - model_name: str = "colbert-ir/colbertv2.0", - cache_dir: Optional[str] = None, - threads: Optional[int] = None, - k_sim: int = 4, - d_proj: int = 32, - R_reps: int = 10, - random_seed: Optional[int] = None, - **kwargs: Any, - ): - """ - Initialize MuveraEmbedding with a late interaction model and MUVERA parameters. - - Args: - model_name (str): Name of the late interaction model to use. - Must be a supported late interaction model (ColBERT, JinaColBERT, etc.) - cache_dir (Optional[str]): Cache directory for model files - threads (Optional[int]): Number of threads for model inference - k_sim (int): Number of SimHash functions (creates 2^k_sim clusters). Default: 4 - d_proj (int): Dimensionality after random projection. Default: 32 - R_reps (int): Number of random projection repetitions for robustness. Default: 10 - random_seed (Optional[int]): Random seed for reproducibility. Default: None - **kwargs: Additional arguments passed to the late interaction model - - Raises: - ValueError: If the model_name is not a supported late interaction model - """ - super().__init__(model_name, cache_dir, threads, **kwargs) - - # Initialize the late interaction model (import locally to avoid circular imports) - try: - from fastembed.late_interaction.late_interaction_text_embedding import ( - LateInteractionTextEmbedding, - ) - - self.late_interaction_model = LateInteractionTextEmbedding( - model_name=model_name, cache_dir=cache_dir, threads=threads, **kwargs - ) - except ValueError as e: - raise ValueError( - f"Model {model_name} is not supported as a late interaction model. " - f"Please use a supported late interaction model like 'colbert-ir/colbertv2.0' or 'jinaai/jina-colbert-v2'. " - f"Original error: {e}" - ) - - # Store MUVERA parameters - self.k_sim = k_sim - self.d_proj = d_proj - self.R_reps = R_reps - - # Get the token embedding dimension from the late interaction model - self.token_dim = self.late_interaction_model.embedding_size - - # Initialize MUVERA algorithm - generator = np.random.default_rng(random_seed) - self.muvera_algorithm = MuveraAlgorithm( - k_sim=k_sim, - d=self.token_dim, - d_proj=d_proj, - R_reps=R_reps, - random_generator=generator, - ) - - # Cache the output embedding size - self._embedding_size: int = self.muvera_algorithm.get_output_dimension() - - @property - def embedding_size(self) -> int: - """Get the embedding size of the MUVERA output.""" - return self._embedding_size - - @classmethod - def get_embedding_size( - cls, model_name: str, k_sim: int = 4, d_proj: int = 32, R_reps: int = 10 - ) -> int: - """ - Get the embedding size for a given model and MUVERA parameters. - - Args: - model_name (str): Name of the late interaction model - k_sim (int): Number of SimHash functions. Default: 4 - d_proj (int): Dimensionality after random projection. Default: 32 - R_reps (int): Number of random projection repetitions. Default: 10 - - Returns: - int: The size of the MUVERA embedding (R_reps * 2^k_sim * d_proj) - - Raises: - ValueError: If the model name is not found in supported models - """ - # Calculate MUVERA output dimension - B = 2**k_sim - return R_reps * B * d_proj - - @classmethod - def _list_supported_models(cls) -> list[DenseModelDescription]: - """ - List supported models (same as late interaction models). - - Returns: - list[DenseModelDescription]: List of supported late interaction models - """ - from fastembed.late_interaction.late_interaction_text_embedding import ( - LateInteractionTextEmbedding, - ) - - return LateInteractionTextEmbedding._list_supported_models() - - @classmethod - def list_supported_models(cls) -> list[dict[str, Any]]: - """ - Lists the supported models. - - Returns: - list[dict[str, Any]]: A list of dictionaries containing the model information. - """ - return [asdict(model) for model in cls._list_supported_models()] - - def embed( - self, - documents: Union[str, Iterable[str]], - batch_size: int = 256, - parallel: Optional[int] = None, - **kwargs: Any, - ) -> Iterable[NumpyArray]: - """ - Encode a list of documents into MUVERA embeddings. - - This method uses the late interaction model to get token-level embeddings, - then applies the MUVERA algorithm to create fixed-dimensional embeddings. - Documents are encoded with normalization and empty cluster filling. - - Args: - documents: Iterator of documents or single document to embed - batch_size: Batch size for encoding -- higher values will use more memory, but be faster - parallel: If > 1, data-parallel encoding will be used, recommended for offline encoding of large datasets. - If 0, use all available cores. - If None, don't use data-parallel processing, use default onnxruntime threading instead. - - Returns: - Iterable[NumpyArray]: List of MUVERA embeddings, one per document - """ - # Handle single string input - if isinstance(documents, str): - documents = [documents] - - # Get token-level embeddings from the late interaction model - token_embeddings = self.late_interaction_model.embed( - documents=documents, batch_size=batch_size, parallel=parallel, **kwargs - ) - - # Apply MUVERA algorithm to each document's token embeddings - for token_embedding in token_embeddings: - # token_embedding shape: (n_tokens, token_dim) - muvera_embedding = self.muvera_algorithm.encode_document(token_embedding) - yield muvera_embedding.astype(np.float32) - - def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]: - """ - Embeds queries using MUVERA algorithm. - - This method uses the late interaction model to get token-level embeddings, - then applies the MUVERA algorithm with query-specific settings (no normalization - and no empty cluster filling) to preserve query vector magnitudes. - - Args: - query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries. - - Returns: - Iterable[NumpyArray]: The MUVERA query embeddings. - """ - # Handle single string input - if isinstance(query, str): - query = [query] - - # Get token-level embeddings from the late interaction model - token_embeddings = self.late_interaction_model.query_embed(query, **kwargs) - - # Apply MUVERA algorithm to each query's token embeddings - for token_embedding in token_embeddings: - # token_embedding shape: (n_tokens, token_dim) - muvera_embedding = self.muvera_algorithm.encode_query(token_embedding) - yield muvera_embedding.astype(np.float32) - - def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]: - """ - Embeds a list of text passages into MUVERA embeddings. - - This is an alias for the embed method, following the fastembed interface. - - Args: - texts (Iterable[str]): The list of texts to embed. - **kwargs: Additional keyword arguments to pass to the embed method. - - Yields: - Iterable[NumpyArray]: The MUVERA embeddings. - """ - yield from self.embed(texts, **kwargs) diff --git a/fastembed/text/text_embedding.py b/fastembed/text/text_embedding.py index de60d7e6b..117f5af79 100644 --- a/fastembed/text/text_embedding.py +++ b/fastembed/text/text_embedding.py @@ -8,7 +8,6 @@ from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding from fastembed.text.pooled_embedding import PooledEmbedding from fastembed.text.multitask_embedding import JinaEmbeddingV3 -from fastembed.text.muvera_embedding import MuveraEmbedding from fastembed.text.onnx_embedding import OnnxTextEmbedding from fastembed.text.text_embedding_base import TextEmbeddingBase from fastembed.common.model_description import DenseModelDescription, ModelSource, PoolingType @@ -21,7 +20,6 @@ class TextEmbedding(TextEmbeddingBase): PooledNormalizedEmbedding, PooledEmbedding, JinaEmbeddingV3, - MuveraEmbedding, CustomTextEmbedding, ] diff --git a/tests/test_postprocess.py b/tests/test_postprocess.py new file mode 100644 index 000000000..e69de29bb From 73fed1e966e8b7cff9a6ab792f0e5c2878a81285 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kacper=20=C5=81ukawski?= Date: Fri, 11 Jul 2025 14:28:12 +0200 Subject: [PATCH 07/13] feat: enhance MuveraEmbedding with multi-vector model support and improve parameter defaults --- fastembed/postprocess/muvera.py | 89 +++++++++++++++++++++++++++++---- 1 file changed, 79 insertions(+), 10 deletions(-) diff --git a/fastembed/postprocess/muvera.py b/fastembed/postprocess/muvera.py index 05ee52fe8..7d6a0f59c 100644 --- a/fastembed/postprocess/muvera.py +++ b/fastembed/postprocess/muvera.py @@ -1,6 +1,14 @@ import numpy as np from fastembed.common.types import NumpyArray +from fastembed.late_interaction.late_interaction_embedding_base import ( + LateInteractionTextEmbeddingBase, +) +from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding_base import ( + LateInteractionMultimodalEmbeddingBase, +) + +MultiVectorModel = LateInteractionTextEmbeddingBase | LateInteractionMultimodalEmbeddingBase class SimHashProjection: @@ -89,21 +97,25 @@ class MuveraPostprocessor: def __init__( self, - k_sim: int, d: int, - d_proj: int, - R_reps: int, - random_seed: np.random.Generator = 42, + k_sim: int = 5, + d_proj: int = 16, + R_reps: int = 20, # noqa[naming] + random_seed: int = 42, ): """ Initialize MUVERA algorithm with specified parameters. Args: - k_sim (int): Number of SimHash functions (creates 2^k_sim clusters) - d (int): Dimensionality of input vectors - d_proj (int): Dimensionality after random projection (must be <= d) - R_reps (int): Number of random projection repetitions for robustness - random_seed (np.random.Generator): Seed for random generator + d (int): Dimensionality of individual input vectors + k_sim (int, optional): Number of SimHash functions (creates 2^k_sim clusters). + Defaults to 5. + d_proj (int, optional): Dimensionality after random projection (must be <= d). + Defaults to 16. + R_reps (int, optional): Number of random projection repetitions for robustness. + Defaults to 20. + random_seed (int, optional): Seed for random number generator to ensure + reproducible results. Defaults to 42. Raises: ValueError: If d_proj > d (cannot project to higher dimensionality) @@ -124,7 +136,64 @@ def __init__( for _ in range(R_reps) ] # Random projection matrices with entries from {-1, +1} for each repetition - self.S_projections = random_seed.choice([-1, 1], size=(R_reps, d, d_proj)) + self.S_projections = generator.choice([-1, 1], size=(R_reps, d, d_proj)) + + @classmethod + def from_multivector_model( + cls, + model: MultiVectorModel, + k_sim: int = 5, + d_proj: int = 16, + R_reps: int = 20, # noqa[naming] + random_seed: int = 42, + ) -> "MuveraPostprocessor": + """ + Create a MuveraPostprocessor instance from a multi-vector embedding model. + + This class method provides a convenient way to initialize a MUVERA postprocessor + that is compatible with a given multi-vector model by automatically extracting + the embedding dimensionality from the model. + + Args: + model (MultiVectorModel): A late interaction text or multimodal embedding model + that provides multi-vector embeddings. Must have an + `embedding_size` attribute specifying the dimensionality + of individual vectors. + k_sim (int, optional): Number of SimHash functions (creates 2^k_sim clusters). + Defaults to 5. + d_proj (int, optional): Dimensionality after random projection (must be <= model's + embedding_size). Defaults to 16. + R_reps (int, optional): Number of random projection repetitions for robustness. + Defaults to 20. + random_seed (int, optional): Seed for random number generator to ensure + reproducible results. Defaults to 42. + + Returns: + MuveraPostprocessor: A configured MUVERA postprocessor instance ready to + process embeddings from the given model. + + Raises: + ValueError: If d_proj > model.embedding_size (cannot project to higher dimensionality) + + Example: + >>> from fastembed.late_interaction.colbert import Colbert + >>> model = Colbert(model_name="colbert-ir/colbertv2.0") + >>> muvera_postprocessor = MuveraPostprocessor.from_multivector_model( + ... model=model, + ... k_sim=6, + ... d_proj=32 + ... ) + >>> # Now use postprocessor with embeddings from the model + >>> embeddings = model.embed(["sample text"]) + >>> fde = muvera_postprocessor.process_document(embeddings[0]) + """ + return cls( + d=model.embedding_size, + k_sim=k_sim, + d_proj=d_proj, + R_reps=R_reps, + random_seed=random_seed, + ) def get_output_dimension(self) -> int: """ From b346f62cd4aeb1e9c1b7513c0aeef9fa421f8cc1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kacper=20=C5=81ukawski?= Date: Fri, 11 Jul 2025 14:28:54 +0200 Subject: [PATCH 08/13] feat: add embedding_size property to MuveraEmbedding --- fastembed/postprocess/muvera.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/fastembed/postprocess/muvera.py b/fastembed/postprocess/muvera.py index 7d6a0f59c..a1e49b70a 100644 --- a/fastembed/postprocess/muvera.py +++ b/fastembed/postprocess/muvera.py @@ -205,6 +205,10 @@ def get_output_dimension(self) -> int: B = 2**self.k_sim return self.R_reps * B * self.d_proj + @property + def embedding_size(self) -> int: + return self.get_output_dimension() + def process_document(self, vectors: np.ndarray) -> np.ndarray: """ Encode a document's vectors into a Fixed Dimensional Encoding (FDE). From b63558d98d486288c4457931a0bb45216aec6561 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Kacper=20=C5=81ukawski?= Date: Fri, 11 Jul 2025 17:09:42 +0200 Subject: [PATCH 09/13] feat: update MuveraPostprocessor to use model description for embedding size and add Jupyter notebook for MUVERA usage --- ...gs_from_Multi_Vector_Representations.ipynb | 203 ++++++++++++++++++ fastembed/postprocess/muvera.py | 3 +- 2 files changed, 205 insertions(+), 1 deletion(-) create mode 100644 docs/experimental/MUVERA_Embeddings_from_Multi_Vector_Representations.ipynb diff --git a/docs/experimental/MUVERA_Embeddings_from_Multi_Vector_Representations.ipynb b/docs/experimental/MUVERA_Embeddings_from_Multi_Vector_Representations.ipynb new file mode 100644 index 000000000..fa62d3119 --- /dev/null +++ b/docs/experimental/MUVERA_Embeddings_from_Multi_Vector_Representations.ipynb @@ -0,0 +1,203 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "ef06fb5d0f1a68c8", + "metadata": {}, + "source": "MUVERA (**Mu**lti-**Ve**ctor **R**etrieval **A**lgorithm) is an algorithm that transforms variable-length sequences of vectors into Fixed Dimensional Encodings (FDEs) described in the [paper](https://arxiv.org/abs/2405.19504). This is particularly useful for converting multi-vector representations from late interaction models (like ColBERT) into fixed-size embeddings that can be efficiently stored and searched." + }, + { + "cell_type": "markdown", + "id": "f2f81913084b1a22", + "metadata": {}, + "source": [ + "## 🤝 MUVERA with FastEmbed\n", + "\n", + "The original paper suggests using the created FDEs for initial retrieval and original multi-vector representations for reranking to achieve the best quality of the results. FastEmbed implements the MUVERA algorithm as a postprocessor, not a separate model, so you can pass the sequence of vectors from the late interaction model to the postprocessor and get the FDE as a result. By implementing it that way, we ensure you don't need to encode your data with a multi-vector model twice if you decide to keep both representations.\n", + "\n", + "If you used multi-vector model before, then you have probably created an instance of it like this:" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "23709588e9bfc403", + "metadata": { + "ExecuteTime": { + "end_time": "2025-07-11T15:08:53.742140Z", + "start_time": "2025-07-11T15:08:49.351855Z" + } + }, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Fetching 5 files: 100%|██████████| 5/5 [00:03<00:00, 1.52it/s]\n" + ] + } + ], + "source": [ + "from fastembed.late_interaction.colbert import Colbert\n", + "\n", + "model = Colbert(model_name=\"answerdotai/answerai-colbert-small-v1\")" + ] + }, + { + "cell_type": "markdown", + "id": "c73c7538a026c991", + "metadata": {}, + "source": "Adding MUVERA embeddings requires postprocessing the embeddings generated by your multi-vector encoder. Let's import the FastEmbed's MUVERA implementation and set it up specifically for the model we use. MUVERA needs to know the dimensionality of the individual vectors that your model creates, so we can eiter set it up manually or just pass a model instance to a helper class method `.from_multivector_model`:" + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "478914a84c8516a5", + "metadata": { + "ExecuteTime": { + "end_time": "2025-07-11T15:08:53.830455Z", + "start_time": "2025-07-11T15:08:53.823539Z" + } + }, + "outputs": [], + "source": [ + "from fastembed.postprocess.muvera import MuveraPostprocessor\n", + "\n", + "muvera_postprocessor = MuveraPostprocessor.from_multivector_model(model)" + ] + }, + { + "cell_type": "markdown", + "id": "cf56bfb49efb9db7", + "metadata": {}, + "source": [ + "## 🗂️ Ingesting the documents\n", + "\n", + "The original paper separates processing the document embeddings (the ones we store in the Qdrant collection to search over), from the query embeddings and calculate them in a slightly different way. Thus, there are different processing methods, depending on whether you encode document or queries." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "6d0518c6ef696a41", + "metadata": { + "ExecuteTime": { + "end_time": "2025-07-11T15:08:53.959852Z", + "start_time": "2025-07-11T15:08:53.872938Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "(10240,)\n" + ] + } + ], + "source": [ + "multivectors = model.passage_embed(\n", + " [\n", + " \"Paris is a capital of France\",\n", + " \"Berlin is a capital of Germany\",\n", + " \"The best chestnuts are in Place Pigalle\",\n", + " ]\n", + ")\n", + "fde_vectors = [muvera_postprocessor.process_document(v) for v in multivectors]\n", + "print(fde_vectors[0].shape)" + ] + }, + { + "cell_type": "markdown", + "id": "411028b2bdb33ac1", + "metadata": {}, + "source": [ + "## 🔎 Querying\n", + "\n", + "When querying, we use a different method of the `MuveraPostprocessor` to convert the multi-vector representations into FDEs:" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "ae95c8ee5bb59a29", + "metadata": { + "ExecuteTime": { + "end_time": "2025-07-11T15:08:53.981466Z", + "start_time": "2025-07-11T15:08:53.965051Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "(10240,)\n" + ] + } + ], + "source": [ + "query_multivector = next(model.query_embed(\"French cuisine\"))\n", + "query_fde = muvera_postprocessor.process_query(query_multivector)\n", + "print(query_fde.shape)" + ] + }, + { + "cell_type": "markdown", + "id": "199c5365fbebcfc4", + "metadata": {}, + "source": [ + "## 🦀 Usage with Qdrant\n", + "\n", + "If you want to reproduce the whole process as described in the MUVERA paper, you have to create a single Qdrant collection with two [named vectors](https://qdrant.tech/documentation/concepts/vectors/#named-vectors): one for the multi-vector representation and one for the MUVERA embedding. The latter one has a dimensionality that depends on how you configure the parameters of the MUVERA projection and might be checked by inspecting the `.embedding_size` property of the postprocessor. This will be useful for properly configuring the collection." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "a6e363c3db0cd7bc", + "metadata": { + "ExecuteTime": { + "end_time": "2025-07-11T15:08:54.025254Z", + "start_time": "2025-07-11T15:08:54.021436Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "10240" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "muvera_postprocessor.embedding_size" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 2 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython2", + "version": "2.7.6" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/fastembed/postprocess/muvera.py b/fastembed/postprocess/muvera.py index a1e49b70a..966b63f82 100644 --- a/fastembed/postprocess/muvera.py +++ b/fastembed/postprocess/muvera.py @@ -187,8 +187,9 @@ def from_multivector_model( >>> embeddings = model.embed(["sample text"]) >>> fde = muvera_postprocessor.process_document(embeddings[0]) """ + model_desc = model._get_model_description(model.model_name) return cls( - d=model.embedding_size, + d=model_desc.dim, k_sim=k_sim, d_proj=d_proj, R_reps=R_reps, From eae8d48f470d5637f2e80b7febb6acb22ce3fbe1 Mon Sep 17 00:00:00 2001 From: George Date: Thu, 24 Jul 2025 19:05:55 +0300 Subject: [PATCH 10/13] fix: fix types, doctest, rename variables, refactor (#545) * fix: fix types, doctest, rename variables, refactor * fix: fix python3.9 compatibility --- ...gs_from_Multi_Vector_Representations.ipynb | 203 ------------------ fastembed/postprocess/__init__.py | 3 + fastembed/postprocess/muvera.py | 158 +++++++------- 3 files changed, 83 insertions(+), 281 deletions(-) delete mode 100644 docs/experimental/MUVERA_Embeddings_from_Multi_Vector_Representations.ipynb diff --git a/docs/experimental/MUVERA_Embeddings_from_Multi_Vector_Representations.ipynb b/docs/experimental/MUVERA_Embeddings_from_Multi_Vector_Representations.ipynb deleted file mode 100644 index fa62d3119..000000000 --- a/docs/experimental/MUVERA_Embeddings_from_Multi_Vector_Representations.ipynb +++ /dev/null @@ -1,203 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "ef06fb5d0f1a68c8", - "metadata": {}, - "source": "MUVERA (**Mu**lti-**Ve**ctor **R**etrieval **A**lgorithm) is an algorithm that transforms variable-length sequences of vectors into Fixed Dimensional Encodings (FDEs) described in the [paper](https://arxiv.org/abs/2405.19504). This is particularly useful for converting multi-vector representations from late interaction models (like ColBERT) into fixed-size embeddings that can be efficiently stored and searched." - }, - { - "cell_type": "markdown", - "id": "f2f81913084b1a22", - "metadata": {}, - "source": [ - "## 🤝 MUVERA with FastEmbed\n", - "\n", - "The original paper suggests using the created FDEs for initial retrieval and original multi-vector representations for reranking to achieve the best quality of the results. FastEmbed implements the MUVERA algorithm as a postprocessor, not a separate model, so you can pass the sequence of vectors from the late interaction model to the postprocessor and get the FDE as a result. By implementing it that way, we ensure you don't need to encode your data with a multi-vector model twice if you decide to keep both representations.\n", - "\n", - "If you used multi-vector model before, then you have probably created an instance of it like this:" - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "id": "23709588e9bfc403", - "metadata": { - "ExecuteTime": { - "end_time": "2025-07-11T15:08:53.742140Z", - "start_time": "2025-07-11T15:08:49.351855Z" - } - }, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "Fetching 5 files: 100%|██████████| 5/5 [00:03<00:00, 1.52it/s]\n" - ] - } - ], - "source": [ - "from fastembed.late_interaction.colbert import Colbert\n", - "\n", - "model = Colbert(model_name=\"answerdotai/answerai-colbert-small-v1\")" - ] - }, - { - "cell_type": "markdown", - "id": "c73c7538a026c991", - "metadata": {}, - "source": "Adding MUVERA embeddings requires postprocessing the embeddings generated by your multi-vector encoder. Let's import the FastEmbed's MUVERA implementation and set it up specifically for the model we use. MUVERA needs to know the dimensionality of the individual vectors that your model creates, so we can eiter set it up manually or just pass a model instance to a helper class method `.from_multivector_model`:" - }, - { - "cell_type": "code", - "execution_count": 2, - "id": "478914a84c8516a5", - "metadata": { - "ExecuteTime": { - "end_time": "2025-07-11T15:08:53.830455Z", - "start_time": "2025-07-11T15:08:53.823539Z" - } - }, - "outputs": [], - "source": [ - "from fastembed.postprocess.muvera import MuveraPostprocessor\n", - "\n", - "muvera_postprocessor = MuveraPostprocessor.from_multivector_model(model)" - ] - }, - { - "cell_type": "markdown", - "id": "cf56bfb49efb9db7", - "metadata": {}, - "source": [ - "## 🗂️ Ingesting the documents\n", - "\n", - "The original paper separates processing the document embeddings (the ones we store in the Qdrant collection to search over), from the query embeddings and calculate them in a slightly different way. Thus, there are different processing methods, depending on whether you encode document or queries." - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "id": "6d0518c6ef696a41", - "metadata": { - "ExecuteTime": { - "end_time": "2025-07-11T15:08:53.959852Z", - "start_time": "2025-07-11T15:08:53.872938Z" - } - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "(10240,)\n" - ] - } - ], - "source": [ - "multivectors = model.passage_embed(\n", - " [\n", - " \"Paris is a capital of France\",\n", - " \"Berlin is a capital of Germany\",\n", - " \"The best chestnuts are in Place Pigalle\",\n", - " ]\n", - ")\n", - "fde_vectors = [muvera_postprocessor.process_document(v) for v in multivectors]\n", - "print(fde_vectors[0].shape)" - ] - }, - { - "cell_type": "markdown", - "id": "411028b2bdb33ac1", - "metadata": {}, - "source": [ - "## 🔎 Querying\n", - "\n", - "When querying, we use a different method of the `MuveraPostprocessor` to convert the multi-vector representations into FDEs:" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "id": "ae95c8ee5bb59a29", - "metadata": { - "ExecuteTime": { - "end_time": "2025-07-11T15:08:53.981466Z", - "start_time": "2025-07-11T15:08:53.965051Z" - } - }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "(10240,)\n" - ] - } - ], - "source": [ - "query_multivector = next(model.query_embed(\"French cuisine\"))\n", - "query_fde = muvera_postprocessor.process_query(query_multivector)\n", - "print(query_fde.shape)" - ] - }, - { - "cell_type": "markdown", - "id": "199c5365fbebcfc4", - "metadata": {}, - "source": [ - "## 🦀 Usage with Qdrant\n", - "\n", - "If you want to reproduce the whole process as described in the MUVERA paper, you have to create a single Qdrant collection with two [named vectors](https://qdrant.tech/documentation/concepts/vectors/#named-vectors): one for the multi-vector representation and one for the MUVERA embedding. The latter one has a dimensionality that depends on how you configure the parameters of the MUVERA projection and might be checked by inspecting the `.embedding_size` property of the postprocessor. This will be useful for properly configuring the collection." - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "id": "a6e363c3db0cd7bc", - "metadata": { - "ExecuteTime": { - "end_time": "2025-07-11T15:08:54.025254Z", - "start_time": "2025-07-11T15:08:54.021436Z" - } - }, - "outputs": [ - { - "data": { - "text/plain": [ - "10240" - ] - }, - "execution_count": 5, - "metadata": {}, - "output_type": "execute_result" - } - ], - "source": [ - "muvera_postprocessor.embedding_size" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 2 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython2", - "version": "2.7.6" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/fastembed/postprocess/__init__.py b/fastembed/postprocess/__init__.py index e69de29bb..5d03c50e8 100644 --- a/fastembed/postprocess/__init__.py +++ b/fastembed/postprocess/__init__.py @@ -0,0 +1,3 @@ +from fastembed.postprocess.muvera import Muvera + +__all__ = ["Muvera"] diff --git a/fastembed/postprocess/muvera.py b/fastembed/postprocess/muvera.py index 966b63f82..3da15864d 100644 --- a/fastembed/postprocess/muvera.py +++ b/fastembed/postprocess/muvera.py @@ -1,3 +1,5 @@ +from typing import Union + import numpy as np from fastembed.common.types import NumpyArray @@ -8,7 +10,7 @@ LateInteractionMultimodalEmbeddingBase, ) -MultiVectorModel = LateInteractionTextEmbeddingBase | LateInteractionMultimodalEmbeddingBase +MultiVectorModel = Union[LateInteractionTextEmbeddingBase, LateInteractionMultimodalEmbeddingBase] class SimHashProjection: @@ -21,23 +23,23 @@ class SimHashProjection: Attributes: k_sim (int): Number of SimHash functions (hyperplanes) - d (int): Dimensionality of input vectors - simhash_vectors (np.ndarray): Random hyperplane normal vectors of shape (d, k_sim) + dim (int): Dimensionality of input vectors + simhash_vectors (np.ndarray): Random hyperplane normal vectors of shape (dim, k_sim) """ - def __init__(self, k_sim: int, d: int, random_generator: np.random.Generator): + def __init__(self, k_sim: int, dim: int, random_generator: np.random.Generator): """ Initialize SimHash projection with random hyperplanes. Args: k_sim (int): Number of SimHash functions, determines 2^k_sim clusters - d (int): Dimensionality of input vectors + dim (int): Dimensionality of input vectors random_generator (np.random.Generator): Random number generator for reproducibility """ self.k_sim = k_sim - self.d = d + self.dim = dim # Generate k_sim random hyperplanes (normal vectors) from standard normal distribution - self.simhash_vectors = random_generator.normal(size=(d, k_sim)) + self.simhash_vectors = random_generator.normal(size=(dim, k_sim)) def get_cluster_id(self, vector: np.ndarray) -> int: """ @@ -48,7 +50,7 @@ def get_cluster_id(self, vector: np.ndarray) -> int: the resulting binary string as an integer. Args: - vector (np.ndarray): Input vector of shape (d,) + vector (np.ndarray): Input vector of shape (dim,) Returns: int: Cluster ID in range [0, 2^k_sim - 1] @@ -57,24 +59,24 @@ def get_cluster_id(self, vector: np.ndarray) -> int: AssertionError: If a vector shape doesn't match expected dimensionality """ assert vector.shape == ( - self.d, - ), f"Expected vector of shape ({self.d},), got {vector.shape}" + self.dim, + ), f"Expected vector of shape ({self.dim},), got {vector.shape}" # Project vector onto each hyperplane normal vector dot_product = np.dot(vector, self.simhash_vectors) # Apply sign function to get binary values (1 if positive, 0 if negative) binary_values = (dot_product > 0).astype(int) - # Convert binary representation to decimal cluster ID # Each bit position i contributes bit_value * 2^i to the final ID cluster_id = 0 for i, bit in enumerate(binary_values): cluster_id += bit * (2**i) + return cluster_id -class MuveraPostprocessor: +class Muvera: """ MUVERA (Multi-Vector Retrieval Architecture) algorithm implementation. @@ -88,69 +90,69 @@ class MuveraPostprocessor: Attributes: k_sim (int): Number of SimHash functions per projection - d (int): Input vector dimensionality - d_proj (int): Output dimensionality after random projection - R_reps (int): Number of random projection repetitions + dim (int): Input vector dimensionality + dim_proj (int): Output dimensionality after random projection + r_reps (int): Number of random projection repetitions simhash_projections (List[SimHashProjection]): SimHash instances for clustering - S_projections (np.ndarray): Random projection matrices of shape (R_reps, d, d_proj) + dim_reduction_projections (np.ndarray): Random projection matrices of shape (R_reps, d, d_proj) """ def __init__( self, - d: int, + dim: int, k_sim: int = 5, - d_proj: int = 16, - R_reps: int = 20, # noqa[naming] + dim_proj: int = 16, + r_reps: int = 20, random_seed: int = 42, ): """ Initialize MUVERA algorithm with specified parameters. Args: - d (int): Dimensionality of individual input vectors + dim (int): Dimensionality of individual input vectors k_sim (int, optional): Number of SimHash functions (creates 2^k_sim clusters). Defaults to 5. - d_proj (int, optional): Dimensionality after random projection (must be <= d). + dim_proj (int, optional): Dimensionality after random projection (must be <= dim). Defaults to 16. - R_reps (int, optional): Number of random projection repetitions for robustness. + r_reps (int, optional): Number of random projection repetitions for robustness. Defaults to 20. random_seed (int, optional): Seed for random number generator to ensure reproducible results. Defaults to 42. Raises: - ValueError: If d_proj > d (cannot project to higher dimensionality) + ValueError: If dim_proj > dim (cannot project to higher dimensionality) """ - if d_proj > d: + if dim_proj > dim: raise ValueError( - f"Cannot project to a higher dimensionality (d_proj={d_proj} > d={d})" + f"Cannot project to a higher dimensionality (dim_proj={dim_proj} > dim={dim})" ) self.k_sim = k_sim - self.d = d - self.d_proj = d_proj - self.R_reps = R_reps - # Create R_reps independent SimHash projections for robustness + self.dim = dim + self.dim_proj = dim_proj + self.r_reps = r_reps + # Create r_reps independent SimHash projections for robustness generator = np.random.default_rng(random_seed) self.simhash_projections = [ - SimHashProjection(k_sim=self.k_sim, d=self.d, random_generator=generator) - for _ in range(R_reps) + SimHashProjection(k_sim=self.k_sim, dim=self.dim, random_generator=generator) + for _ in range(r_reps) ] # Random projection matrices with entries from {-1, +1} for each repetition - self.S_projections = generator.choice([-1, 1], size=(R_reps, d, d_proj)) + self.dim_reduction_projections = generator.choice([-1, 1], size=(r_reps, dim, dim_proj)) @classmethod def from_multivector_model( cls, model: MultiVectorModel, k_sim: int = 5, - d_proj: int = 16, - R_reps: int = 20, # noqa[naming] + dim_proj: int = 16, + r_reps: int = 20, # noqa[naming] random_seed: int = 42, - ) -> "MuveraPostprocessor": + ) -> "Muvera": """ - Create a MuveraPostprocessor instance from a multi-vector embedding model. + Create a Muvera instance from a multi-vector embedding model. - This class method provides a convenient way to initialize a MUVERA postprocessor + This class method provides a convenient way to initialize a MUVERA that is compatible with a given multi-vector model by automatically extracting the embedding dimensionality from the model. @@ -161,38 +163,36 @@ def from_multivector_model( of individual vectors. k_sim (int, optional): Number of SimHash functions (creates 2^k_sim clusters). Defaults to 5. - d_proj (int, optional): Dimensionality after random projection (must be <= model's + dim_proj (int, optional): Dimensionality after random projection (must be <= model's embedding_size). Defaults to 16. - R_reps (int, optional): Number of random projection repetitions for robustness. + r_reps (int, optional): Number of random projection repetitions for robustness. Defaults to 20. random_seed (int, optional): Seed for random number generator to ensure reproducible results. Defaults to 42. Returns: - MuveraPostprocessor: A configured MUVERA postprocessor instance ready to - process embeddings from the given model. + Muvera: A configured MUVERA instance ready to process embeddings from the given model. Raises: - ValueError: If d_proj > model.embedding_size (cannot project to higher dimensionality) + ValueError: If dim_proj > model.embedding_size (cannot project to higher dimensionality) Example: - >>> from fastembed.late_interaction.colbert import Colbert - >>> model = Colbert(model_name="colbert-ir/colbertv2.0") - >>> muvera_postprocessor = MuveraPostprocessor.from_multivector_model( + >>> from fastembed import LateInteractionTextEmbedding + >>> model = LateInteractionTextEmbedding(model_name="colbert-ir/colbertv2.0") + >>> muvera = Muvera.from_multivector_model( ... model=model, ... k_sim=6, - ... d_proj=32 + ... dim_proj=32 ... ) >>> # Now use postprocessor with embeddings from the model - >>> embeddings = model.embed(["sample text"]) - >>> fde = muvera_postprocessor.process_document(embeddings[0]) + >>> embeddings = np.array(list(model.embed(["sample text"]))) + >>> fde = muvera.process_document(embeddings[0]) """ - model_desc = model._get_model_description(model.model_name) return cls( - d=model_desc.dim, + dim=model.embedding_size, k_sim=k_sim, - d_proj=d_proj, - R_reps=R_reps, + dim_proj=dim_proj, + r_reps=r_reps, random_seed=random_seed, ) @@ -201,16 +201,16 @@ def get_output_dimension(self) -> int: Get the output dimension of the MUVERA algorithm. Returns: - int: Output dimension (R_reps * B * d_proj) where B = 2^k_sim + int: Output dimension (r_reps * num_partitions * dim_proj) where b = 2^k_sim """ - B = 2**self.k_sim - return self.R_reps * B * self.d_proj + num_partitions = 2**self.k_sim + return self.r_reps * num_partitions * self.dim_proj @property def embedding_size(self) -> int: return self.get_output_dimension() - def process_document(self, vectors: np.ndarray) -> np.ndarray: + def process_document(self, vectors: NumpyArray) -> NumpyArray: """ Encode a document's vectors into a Fixed Dimensional Encoding (FDE). @@ -218,14 +218,14 @@ def process_document(self, vectors: np.ndarray) -> np.ndarray: and fills empty clusters using Hamming distance-based selection. Args: - vectors (np.ndarray): Document vectors of shape (n_tokens, d) + vectors (NumpyArray): Document vectors of shape (n_tokens, dim) Returns: - np.ndarray: Fixed dimensional encoding of shape (R_reps * B * d_proj,) + NumpyArray: Fixed dimensional encodings of shape (r_reps * b * dim_proj,) """ return self.process(vectors, fill_empty_clusters=True, normalize_by_count=True) - def process_query(self, vectors: np.ndarray) -> np.ndarray: + def process_query(self, vectors: NumpyArray) -> NumpyArray: """ Encode a query's vectors into a Fixed Dimensional Encoding (FDE). @@ -233,10 +233,10 @@ def process_query(self, vectors: np.ndarray) -> np.ndarray: cluster filling to preserve query vector magnitudes. Args: - vectors (np.ndarray): Query vectors of shape (n_tokens, d) + vectors (NumpyArray]): Query vectors of shape (n_tokens, dim) Returns: - np.ndarray: Fixed dimensional encoding of shape (R_reps * B * d_proj,) + NumpyArray: Fixed dimensional encoding of shape (r_reps * b * dim_proj,) """ return self.process(vectors, fill_empty_clusters=False, normalize_by_count=False) @@ -250,7 +250,7 @@ def process( Core encoding method that transforms variable-length vector sequences into FDEs. The encoding process: - 1. For each of R_reps random projections: + 1. For each of r_reps random projections: a. Assign vectors to clusters using SimHash b. Compute cluster centers (sum of vectors in each cluster) c. Optionally normalize by cluster size @@ -260,55 +260,55 @@ def process( 2. Concatenate all projection results Args: - vectors (np.ndarray): Input vectors of shape (n_vectors, d) + vectors (np.ndarray): Input vectors of shape (n_vectors, dim) fill_empty_clusters (bool): Whether to fill empty clusters using nearest vectors based on Hamming distance of cluster IDs normalize_by_count (bool): Whether to normalize cluster centers by the number of vectors assigned to each cluster Returns: - np.ndarray: Fixed dimensional encoding of shape (R_reps * B * d_proj) + np.ndarray: Fixed dimensional encoding of shape (r_reps * b * dim_proj) where B = 2^k_sim is the number of clusters Raises: AssertionError: If input vectors don't have expected dimensionality """ assert ( - vectors.shape[1] == self.d - ), f"Expected vectors of shape (n, {self.d}), got {vectors.shape}" + vectors.shape[1] == self.dim + ), f"Expected vectors of shape (n, {self.dim}), got {vectors.shape}" # Store results from each random projection output_vectors = [] - # B is the number of clusters (2^k_sim) - B = 2**self.k_sim + # num of space partitions in SimHash + num_partitions = 2**self.k_sim for projection_index, simhash in enumerate(self.simhash_projections): # Initialize cluster centers and count vectors assigned to each cluster - cluster_centers = np.zeros((B, self.d)) - cluster_vector_counts = np.zeros(B) + cluster_centers = np.zeros((num_partitions, self.dim)) + cluster_vector_counts = np.zeros(num_partitions) # Assign each vector to its cluster and accumulate cluster centers for vector in vectors: - cluster_id = simhash.get_cluster_id(vector) + cluster_id = simhash.get_cluster_id(vector) # type: ignore cluster_centers[cluster_id] += vector cluster_vector_counts[cluster_id] += 1 # Normalize cluster centers by the number of vectors (for documents) if normalize_by_count: - for i in range(B): + for i in range(num_partitions): if cluster_vector_counts[i] == 0: continue # Skip empty clusters cluster_centers[i] /= cluster_vector_counts[i] # Fill empty clusters using vectors with minimum Hamming distance if fill_empty_clusters: - for i in range(B): + for i in range(num_partitions): if cluster_vector_counts[i] == 0: # Empty cluster found min_hamming = float("inf") best_vector = None # Find vector whose cluster ID has minimum Hamming distance to i for vector in vectors: - vector_cluster_id = simhash.get_cluster_id(vector) + vector_cluster_id = simhash.get_cluster_id(vector) # type: ignore # Hamming distance = number of differing bits in binary representation hamming_dist = bin(i ^ vector_cluster_id).count("1") if hamming_dist < min_hamming: @@ -319,17 +319,19 @@ def process( cluster_centers[i] = best_vector # Apply random projection for dimensionality reduction if needed - if self.d_proj < self.d: - S = self.S_projections[ + if self.dim_proj < self.dim: + dim_reduction_projection = self.dim_reduction_projections[ projection_index ] # Get projection matrix for this repetition - projected_centers = (1 / np.sqrt(self.d_proj)) * np.dot(cluster_centers, S) + projected_centers = (1 / np.sqrt(self.dim_proj)) * np.dot( + cluster_centers, dim_reduction_projection + ) # Flatten cluster centers into a single vector and add to output output_vectors.append(projected_centers.flatten()) continue - # If no projection needed (d_proj == d), use original cluster centers + # If no projection needed (dim_proj == dim), use original cluster centers output_vectors.append(cluster_centers.flatten()) # Concatenate results from all R_reps projections into final FDE From f0fb50e476acead3fa54a30ec87f0599ea5bfae7 Mon Sep 17 00:00:00 2001 From: George Panchuk Date: Thu, 31 Jul 2025 16:05:05 +0300 Subject: [PATCH 11/13] fix: make get_output_dimension protected --- fastembed/postprocess/muvera.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/fastembed/postprocess/muvera.py b/fastembed/postprocess/muvera.py index 3da15864d..78cf5109f 100644 --- a/fastembed/postprocess/muvera.py +++ b/fastembed/postprocess/muvera.py @@ -196,7 +196,7 @@ def from_multivector_model( random_seed=random_seed, ) - def get_output_dimension(self) -> int: + def _get_output_dimension(self) -> int: """ Get the output dimension of the MUVERA algorithm. @@ -208,7 +208,7 @@ def get_output_dimension(self) -> int: @property def embedding_size(self) -> int: - return self.get_output_dimension() + return self._get_output_dimension() def process_document(self, vectors: NumpyArray) -> NumpyArray: """ From 05c4e04b17b5094b11c47216a3558ac00a81e00c Mon Sep 17 00:00:00 2001 From: George Date: Wed, 20 Aug 2025 11:42:57 +0300 Subject: [PATCH 12/13] Optimize muvera (#551) * vectorize operations * fix: fill empty clusters with dataset vectors * rollback get_output_dimension * fix: fix type hints * fix: review comments --- fastembed/postprocess/muvera.py | 120 +++++++++++++++++++------------- 1 file changed, 73 insertions(+), 47 deletions(-) diff --git a/fastembed/postprocess/muvera.py b/fastembed/postprocess/muvera.py index 78cf5109f..35dfaabe9 100644 --- a/fastembed/postprocess/muvera.py +++ b/fastembed/postprocess/muvera.py @@ -10,7 +10,25 @@ LateInteractionMultimodalEmbeddingBase, ) + MultiVectorModel = Union[LateInteractionTextEmbeddingBase, LateInteractionMultimodalEmbeddingBase] +MAX_HAMMING_DISTANCE = 65 # 64 bits + 1 +POPCOUNT_LUT = np.array([bin(x).count("1") for x in range(256)], dtype=np.uint8) + + +def hamming_distance_matrix(ids: np.ndarray) -> np.ndarray: + """Compute full Hamming distance matrix + + Args: + ids: shape (n,) - array of ids, only size of the array matters + + Return: + np.ndarray (n, n) - hamming distance matrix + """ + n = len(ids) + xor_vals = np.bitwise_xor(ids[:, None], ids[None, :]) # (n, n) uint64 + bytes_view = xor_vals.view(np.uint8).reshape(n, n, 8) # (n, n, 8) + return POPCOUNT_LUT[bytes_view].sum(axis=2) class SimHashProjection: @@ -41,39 +59,28 @@ def __init__(self, k_sim: int, dim: int, random_generator: np.random.Generator): # Generate k_sim random hyperplanes (normal vectors) from standard normal distribution self.simhash_vectors = random_generator.normal(size=(dim, k_sim)) - def get_cluster_id(self, vector: np.ndarray) -> int: + def get_cluster_ids(self, vectors: np.ndarray) -> np.ndarray: """ - Compute the cluster ID for a given vector using SimHash. + Compute the cluster IDs for a given vector using SimHash. The cluster ID is determined by computing the dot product of the vector with each hyperplane normal vector, taking the sign, and interpreting the resulting binary string as an integer. Args: - vector (np.ndarray): Input vector of shape (dim,) + vectors (np.ndarray): Input vectors of shape (n, dim,) Returns: - int: Cluster ID in range [0, 2^k_sim - 1] + np.ndarray: Cluster IDs in range [0, 2^k_sim - 1] Raises: AssertionError: If a vector shape doesn't match expected dimensionality """ - assert vector.shape == ( - self.dim, - ), f"Expected vector of shape ({self.dim},), got {vector.shape}" - - # Project vector onto each hyperplane normal vector - dot_product = np.dot(vector, self.simhash_vectors) - - # Apply sign function to get binary values (1 if positive, 0 if negative) - binary_values = (dot_product > 0).astype(int) - # Convert binary representation to decimal cluster ID - # Each bit position i contributes bit_value * 2^i to the final ID - cluster_id = 0 - for i, bit in enumerate(binary_values): - cluster_id += bit * (2**i) - - return cluster_id + dot_product = ( + vectors @ self.simhash_vectors + ) # (token_num, dim) x (dim, k_sim) -> (token_num, k_sim) + cluster_ids = (dot_product > 0) @ (1 << np.arange(self.k_sim)) + return cluster_ids class Muvera: @@ -93,6 +100,7 @@ class Muvera: dim (int): Input vector dimensionality dim_proj (int): Output dimensionality after random projection r_reps (int): Number of random projection repetitions + random_seed (int): Random seed for consistent random matrix generation simhash_projections (List[SimHashProjection]): SimHash instances for clustering dim_reduction_projections (np.ndarray): Random projection matrices of shape (R_reps, d, d_proj) """ @@ -282,49 +290,59 @@ def process( # num of space partitions in SimHash num_partitions = 2**self.k_sim + cluster_center_ids = np.arange(num_partitions) + precomputed_hamming_matrix = ( + hamming_distance_matrix(cluster_center_ids) if fill_empty_clusters else None + ) + for projection_index, simhash in enumerate(self.simhash_projections): # Initialize cluster centers and count vectors assigned to each cluster cluster_centers = np.zeros((num_partitions, self.dim)) - cluster_vector_counts = np.zeros(num_partitions) + cluster_center_id_to_vectors: dict[int, list[int]] = { + cluster_center_id: [] for cluster_center_id in cluster_center_ids + } + cluster_vector_counts = None + empty_mask = None # Assign each vector to its cluster and accumulate cluster centers - for vector in vectors: - cluster_id = simhash.get_cluster_id(vector) # type: ignore - cluster_centers[cluster_id] += vector - cluster_vector_counts[cluster_id] += 1 + vector_cluster_ids = simhash.get_cluster_ids(vectors) + for cluster_id, (vec_idx, vec) in zip(vector_cluster_ids, enumerate(vectors)): + cluster_centers[cluster_id] += vec + cluster_center_id_to_vectors[cluster_id].append(vec_idx) + + if normalize_by_count or fill_empty_clusters: + cluster_vector_counts = np.bincount(vector_cluster_ids, minlength=num_partitions) + empty_mask = cluster_vector_counts == 0 - # Normalize cluster centers by the number of vectors (for documents) if normalize_by_count: - for i in range(num_partitions): - if cluster_vector_counts[i] == 0: - continue # Skip empty clusters - cluster_centers[i] /= cluster_vector_counts[i] + assert empty_mask is not None + assert cluster_vector_counts is not None + non_empty_mask = ~empty_mask + cluster_centers[non_empty_mask] /= cluster_vector_counts[non_empty_mask][:, None] # Fill empty clusters using vectors with minimum Hamming distance if fill_empty_clusters: - for i in range(num_partitions): - if cluster_vector_counts[i] == 0: # Empty cluster found - min_hamming = float("inf") - best_vector = None - # Find vector whose cluster ID has minimum Hamming distance to i - for vector in vectors: - vector_cluster_id = simhash.get_cluster_id(vector) # type: ignore - # Hamming distance = number of differing bits in binary representation - hamming_dist = bin(i ^ vector_cluster_id).count("1") - if hamming_dist < min_hamming: - min_hamming = hamming_dist - best_vector = vector - # Assign the best matching vector to the empty cluster - if best_vector is not None: - cluster_centers[i] = best_vector + assert empty_mask is not None + assert precomputed_hamming_matrix is not None + masked_hamming = np.where( + empty_mask[None, :], MAX_HAMMING_DISTANCE, precomputed_hamming_matrix + ) + nearest_non_empty = np.argmin(masked_hamming, axis=1) + fill_vectors = np.array( + [ + vectors[cluster_center_id_to_vectors[cluster_id][0]] + for cluster_id in nearest_non_empty[empty_mask] + ] + ).reshape(-1, self.dim) + cluster_centers[empty_mask] = fill_vectors # Apply random projection for dimensionality reduction if needed if self.dim_proj < self.dim: dim_reduction_projection = self.dim_reduction_projections[ projection_index ] # Get projection matrix for this repetition - projected_centers = (1 / np.sqrt(self.dim_proj)) * np.dot( - cluster_centers, dim_reduction_projection + projected_centers = (1 / np.sqrt(self.dim_proj)) * ( + cluster_centers @ dim_reduction_projection ) # Flatten cluster centers into a single vector and add to output @@ -336,3 +354,11 @@ def process( # Concatenate results from all R_reps projections into final FDE return np.concatenate(output_vectors) + + +if __name__ == "__main__": + v_arrs = np.random.randn(10, 100, 128) + muvera = Muvera(128, 4, 8, 20, 42) + + for v_arr in v_arrs: + muvera.process(v_arr) # type: ignore From e9245247bca59719dec47be60a0a55f5cb0cd2a3 Mon Sep 17 00:00:00 2001 From: George Panchuk Date: Thu, 21 Aug 2025 02:53:19 +0300 Subject: [PATCH 13/13] tests: add tests --- tests/test_postprocess.py | 38 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/tests/test_postprocess.py b/tests/test_postprocess.py index e69de29bb..82576f115 100644 --- a/tests/test_postprocess.py +++ b/tests/test_postprocess.py @@ -0,0 +1,38 @@ +import numpy as np + +from fastembed import LateInteractionTextEmbedding +from fastembed.postprocess import Muvera + +CANONICAL_VALUES = [-2.61810007e-04, 1.89005750e00, -2.32070747e00] +CANONICAL_QUERY_VALUES = [ + -0.85783903, + 1.1077204, + -0.09522747, +] # part of the values are zeros, should be compared with the result of nonzero mask + +DIM = 128 +K_SIM = 5 +DIM_PROJ = 16 +R_REPS = 20 + + +def test_single_input(): + model = LateInteractionTextEmbedding("colbert-ir/colbertv2.0", lazy_load=True) + random_generator = np.random.default_rng(42) + multivector = random_generator.random((10, 128)) + + for muvera in ( + Muvera(dim=DIM, k_sim=K_SIM, dim_proj=DIM_PROJ, r_reps=R_REPS, random_seed=42), + Muvera.from_multivector_model(model, k_sim=K_SIM, dim_proj=DIM_PROJ, r_reps=R_REPS), + ): + fde = muvera.process(multivector) + assert fde.shape[0] == muvera.embedding_size + assert np.allclose(fde[:3], CANONICAL_VALUES) + + fde_doc = muvera.process_document(multivector) + assert fde_doc.shape[0] == muvera.embedding_size + assert np.allclose(fde, fde_doc) + + fde_query = muvera.process_query(multivector) + assert fde_query.shape[0] == muvera.embedding_size + assert np.allclose(fde_query[np.nonzero(fde_query)][:3], CANONICAL_QUERY_VALUES)