codekingpro/portable-devtools

1import numpy as np2 3from fastembed.common.types import NumpyArray4from fastembed.late_interaction.late_interaction_embedding_base import (5    LateInteractionTextEmbeddingBase,6)7from fastembed.late_interaction_multimodal.late_interaction_multimodal_embedding_base import (8    LateInteractionMultimodalEmbeddingBase,9)10 11 12MultiVectorModel = LateInteractionTextEmbeddingBase | LateInteractionMultimodalEmbeddingBase13MAX_HAMMING_DISTANCE = 65  # 64 bits + 114POPCOUNT_LUT = np.array([bin(x).count("1") for x in range(256)], dtype=np.uint8)15 16 17def hamming_distance_matrix(ids: np.ndarray) -> np.ndarray:18    """Compute full Hamming distance matrix19 20    Args:21    ids: shape (n,) - array of ids, only size of the array matters22 23    Return:24        np.ndarray (n, n) - hamming distance matrix25    """26    n = len(ids)27    xor_vals = np.bitwise_xor(ids[:, None], ids[None, :])  # (n, n) uint6428    bytes_view = xor_vals.view(np.uint8).reshape(n, n, 8)  # (n, n, 8)29    return POPCOUNT_LUT[bytes_view].sum(axis=2)30 31 32class SimHashProjection:33    """34    SimHash projection component for MUVERA clustering.35 36    This class implements locality-sensitive hashing using random hyperplanes37    to partition the vector space into 2^k_sim clusters. Each vector is assigned38    to a cluster based on which side of k_sim random hyperplanes it falls on.39 40    Attributes:41        k_sim (int): Number of SimHash functions (hyperplanes)42        dim (int): Dimensionality of input vectors43        simhash_vectors (np.ndarray): Random hyperplane normal vectors of shape (dim, k_sim)44    """45 46    def __init__(self, k_sim: int, dim: int, random_generator: np.random.Generator):47        """48        Initialize SimHash projection with random hyperplanes.49 50        Args:51            k_sim (int): Number of SimHash functions, determines 2^k_sim clusters52            dim (int): Dimensionality of input vectors53            random_generator (np.random.Generator): Random number generator for reproducibility54        """55        self.k_sim = k_sim56        self.dim = dim57        # Generate k_sim random hyperplanes (normal vectors) from standard normal distribution58        self.simhash_vectors = random_generator.normal(size=(dim, k_sim))59 60    def get_cluster_ids(self, vectors: np.ndarray) -> np.ndarray:61        """62        Compute the cluster IDs for a given vector using SimHash.63 64        The cluster ID is determined by computing the dot product of the vector65        with each hyperplane normal vector, taking the sign, and interpreting66        the resulting binary string as an integer.67 68        Args:69            vectors (np.ndarray): Input vectors of shape (n, dim,)70 71        Returns:72            np.ndarray: Cluster IDs in range [0, 2^k_sim - 1]73 74        Raises:75            AssertionError: If a vector shape doesn't match expected dimensionality76        """77        dot_product = (78            vectors @ self.simhash_vectors79        )  # (token_num, dim) x (dim, k_sim) -> (token_num, k_sim)80        cluster_ids = (dot_product > 0) @ (1 << np.arange(self.k_sim))81        return cluster_ids82 83 84class Muvera:85    """86    MUVERA (Multi-Vector Retrieval Architecture) algorithm implementation.87 88    This class creates Fixed Dimensional Encodings (FDEs) from variable-length89    sequences of vectors by using SimHash clustering and random projections.90    The process involves:91    1. Clustering vectors using multiple SimHash projections92    2. Computing cluster centers (with different strategies for docs vs queries)93    3. Applying random projections for dimensionality reduction94    4. Concatenating results from all projections95 96    Attributes:97        k_sim (int): Number of SimHash functions per projection98        dim (int): Input vector dimensionality99        dim_proj (int): Output dimensionality after random projection100        r_reps (int): Number of random projection repetitions101        random_seed (int): Random seed for consistent random matrix generation102        simhash_projections (List[SimHashProjection]): SimHash instances for clustering103        dim_reduction_projections (np.ndarray): Random projection matrices of shape (R_reps, d, d_proj)104    """105 106    def __init__(107        self,108        dim: int,109        k_sim: int = 5,110        dim_proj: int = 16,111        r_reps: int = 20,112        random_seed: int = 42,113    ):114        """115        Initialize MUVERA algorithm with specified parameters.116 117        Args:118            dim (int): Dimensionality of individual input vectors119            k_sim (int, optional): Number of SimHash functions (creates 2^k_sim clusters).120                                   Defaults to 5.121            dim_proj (int, optional): Dimensionality after random projection (must be <= dim).122                                    Defaults to 16.123            r_reps (int, optional): Number of random projection repetitions for robustness.124                                    Defaults to 20.125            random_seed (int, optional): Seed for random number generator to ensure126                                         reproducible results. Defaults to 42.127 128        Raises:129            ValueError: If dim_proj > dim (cannot project to higher dimensionality)130        """131        if dim_proj > dim:132            raise ValueError(133                f"Cannot project to a higher dimensionality (dim_proj={dim_proj} > dim={dim})"134            )135 136        self.k_sim = k_sim137        self.dim = dim138        self.dim_proj = dim_proj139        self.r_reps = r_reps140        # Create r_reps independent SimHash projections for robustness141        generator = np.random.default_rng(random_seed)142        self.simhash_projections = [143            SimHashProjection(k_sim=self.k_sim, dim=self.dim, random_generator=generator)144            for _ in range(r_reps)145        ]146        # Random projection matrices with entries from {-1, +1} for each repetition147        self.dim_reduction_projections = generator.choice([-1, 1], size=(r_reps, dim, dim_proj))148 149    @classmethod150    def from_multivector_model(151        cls,152        model: MultiVectorModel,153        k_sim: int = 5,154        dim_proj: int = 16,155        r_reps: int = 20,  # noqa[naming]156        random_seed: int = 42,157    ) -> "Muvera":158        """159        Create a Muvera instance from a multi-vector embedding model.160 161        This class method provides a convenient way to initialize a MUVERA162        that is compatible with a given multi-vector model by automatically extracting163        the embedding dimensionality from the model.164 165        Args:166            model (MultiVectorModel): A late interaction text or multimodal embedding model167                                    that provides multi-vector embeddings. Must have an168                                    `embedding_size` attribute specifying the dimensionality169                                    of individual vectors.170            k_sim (int, optional): Number of SimHash functions (creates 2^k_sim clusters).171                                   Defaults to 5.172            dim_proj (int, optional): Dimensionality after random projection (must be <= model's173                                    embedding_size). Defaults to 16.174            r_reps (int, optional): Number of random projection repetitions for robustness.175                                    Defaults to 20.176            random_seed (int, optional): Seed for random number generator to ensure177                                         reproducible results. Defaults to 42.178 179        Returns:180            Muvera: A configured MUVERA instance ready to process embeddings from the given model.181 182        Raises:183            ValueError: If dim_proj > model.embedding_size (cannot project to higher dimensionality)184 185        Example:186            >>> from fastembed import LateInteractionTextEmbedding187            >>> model = LateInteractionTextEmbedding(model_name="colbert-ir/colbertv2.0")188            >>> muvera = Muvera.from_multivector_model(189            ...     model=model,190            ...     k_sim=6,191            ...     dim_proj=32192            ... )193            >>> # Now use postprocessor with embeddings from the model194            >>> embeddings = np.array(list(model.embed(["sample text"])))195            >>> fde = muvera.process_document(embeddings[0])196        """197        return cls(198            dim=model.embedding_size,199            k_sim=k_sim,200            dim_proj=dim_proj,201            r_reps=r_reps,202            random_seed=random_seed,203        )204 205    def _get_output_dimension(self) -> int:206        """207        Get the output dimension of the MUVERA algorithm.208 209        Returns:210            int: Output dimension (r_reps * num_partitions * dim_proj) where b = 2^k_sim211        """212        num_partitions = 2**self.k_sim213        return self.r_reps * num_partitions * self.dim_proj214 215    @property216    def embedding_size(self) -> int:217        return self._get_output_dimension()218 219    def process_document(self, vectors: NumpyArray) -> NumpyArray:220        """221        Encode a document's vectors into a Fixed Dimensional Encoding (FDE).222 223        Uses document-specific settings: normalizes cluster centers by vector count224        and fills empty clusters using Hamming distance-based selection.225 226        Args:227            vectors (NumpyArray): Document vectors of shape (n_tokens, dim)228 229        Returns:230            NumpyArray: Fixed dimensional encodings of shape (r_reps * b * dim_proj,)231        """232        return self.process(vectors, fill_empty_clusters=True, normalize_by_count=True)233 234    def process_query(self, vectors: NumpyArray) -> NumpyArray:235        """236        Encode a query's vectors into a Fixed Dimensional Encoding (FDE).237 238        Uses query-specific settings: no normalization by count and no empty239        cluster filling to preserve query vector magnitudes.240 241        Args:242            vectors (NumpyArray]): Query vectors of shape (n_tokens, dim)243 244        Returns:245            NumpyArray: Fixed dimensional encoding of shape (r_reps * b * dim_proj,)246        """247        return self.process(vectors, fill_empty_clusters=False, normalize_by_count=False)248 249    def process(250        self,251        vectors: NumpyArray,252        fill_empty_clusters: bool = True,253        normalize_by_count: bool = True,254    ) -> NumpyArray:255        """256        Core encoding method that transforms variable-length vector sequences into FDEs.257 258        The encoding process:259        1. For each of r_reps random projections:260           a. Assign vectors to clusters using SimHash261           b. Compute cluster centers (sum of vectors in each cluster)262           c. Optionally normalize by cluster size263           d. Fill empty clusters using Hamming distance if requested264           e. Apply random projection for dimensionality reduction265           f. Flatten cluster centers into a vector266        2. Concatenate all projection results267 268        Args:269            vectors (np.ndarray): Input vectors of shape (n_vectors, dim)270            fill_empty_clusters (bool): Whether to fill empty clusters using nearest271                                      vectors based on Hamming distance of cluster IDs272            normalize_by_count (bool): Whether to normalize cluster centers by the273                                     number of vectors assigned to each cluster274 275        Returns:276            np.ndarray: Fixed dimensional encoding of shape (r_reps * b * dim_proj)277                        where B = 2^k_sim is the number of clusters278 279        Raises:280            AssertionError: If input vectors don't have expected dimensionality281        """282        assert (283            vectors.shape[1] == self.dim284        ), f"Expected vectors of shape (n, {self.dim}), got {vectors.shape}"285 286        # Store results from each random projection287        output_vectors = []288 289        # num of space partitions in SimHash290        num_partitions = 2**self.k_sim291        cluster_center_ids = np.arange(num_partitions)292        precomputed_hamming_matrix = (293            hamming_distance_matrix(cluster_center_ids) if fill_empty_clusters else None294        )295 296        for projection_index, simhash in enumerate(self.simhash_projections):297            # Initialize cluster centers and count vectors assigned to each cluster298            cluster_centers = np.zeros((num_partitions, self.dim))299            cluster_center_id_to_vectors: dict[int, list[int]] = {300                cluster_center_id: [] for cluster_center_id in cluster_center_ids301            }302            cluster_vector_counts = None303            empty_mask = None304 305            # Assign each vector to its cluster and accumulate cluster centers306            vector_cluster_ids = simhash.get_cluster_ids(vectors)307            for cluster_id, (vec_idx, vec) in zip(vector_cluster_ids, enumerate(vectors)):308                cluster_centers[cluster_id] += vec309                cluster_center_id_to_vectors[cluster_id].append(vec_idx)310 311            if normalize_by_count or fill_empty_clusters:312                cluster_vector_counts = np.bincount(vector_cluster_ids, minlength=num_partitions)313                empty_mask = cluster_vector_counts == 0314 315            if normalize_by_count:316                assert empty_mask is not None317                assert cluster_vector_counts is not None318                non_empty_mask = ~empty_mask319                cluster_centers[non_empty_mask] /= cluster_vector_counts[non_empty_mask][:, None]320 321            # Fill empty clusters using vectors with minimum Hamming distance322            if fill_empty_clusters:323                assert empty_mask is not None324                assert precomputed_hamming_matrix is not None325                masked_hamming = np.where(326                    empty_mask[None, :], MAX_HAMMING_DISTANCE, precomputed_hamming_matrix327                )328                nearest_non_empty = np.argmin(masked_hamming, axis=1)329                fill_vectors = np.array(330                    [331                        vectors[cluster_center_id_to_vectors[cluster_id][0]]332                        for cluster_id in nearest_non_empty[empty_mask]333                    ]334                ).reshape(-1, self.dim)335                cluster_centers[empty_mask] = fill_vectors336 337            # Apply random projection for dimensionality reduction if needed338            if self.dim_proj < self.dim:339                dim_reduction_projection = self.dim_reduction_projections[340                    projection_index341                ]  # Get projection matrix for this repetition342                projected_centers = (1 / np.sqrt(self.dim_proj)) * (343                    cluster_centers @ dim_reduction_projection344                )345 346                # Flatten cluster centers into a single vector and add to output347                output_vectors.append(projected_centers.flatten())348                continue349 350            # If no projection needed (dim_proj == dim), use original cluster centers351            output_vectors.append(cluster_centers.flatten())352 353        # Concatenate results from all R_reps projections into final FDE354        return np.concatenate(output_vectors)355 356 357if __name__ == "__main__":358    v_arrs = np.random.randn(10, 100, 128)359    muvera = Muvera(128, 4, 8, 20, 42)360 361    for v_arr in v_arrs:362        muvera.process(v_arr)  # type: ignore363