muvera-python 0.1.1__tar.gz

muvera_python-0.1.1/.gitignore

@@ -0,0 +1,24 @@
+# Python-generated files
+__pycache__/
+*.py[oc]
+build/
+dist/
+wheels/
+*.egg-info
+
+# Virtual environments
+.venv
+
+# Lock file (library repo — dependencies should stay flexible)
+uv.lock
+
+# Local Python version (managed by pyenv/uv, not part of the library)
+.python-version
+
+# Embedding cache (generated by examples)
+examples/.cache/
+
+.vscode/
+
+# Benchmark results (generated, not committed)
+benchmarks/results/*.json

muvera_python-0.1.1/LICENSE

@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2026 craftsangjae
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

muvera_python-0.1.1/PKG-INFO

@@ -0,0 +1,126 @@
+Metadata-Version: 2.4
+Name: muvera-python
+Version: 0.1.1
+Summary: MuVERA: Multi-Vector Retrieval via Fixed Dimensional Encodings
+Project-URL: Homepage, https://github.com/craftsangjae/muvera-python
+Project-URL: Repository, https://github.com/craftsangjae/muvera-python
+Project-URL: Issues, https://github.com/craftsangjae/muvera-python/issues
+Author: craftsangjae
+License-Expression: MIT
+License-File: LICENSE
+Keywords: embedding,fixed-dimensional-encoding,multi-vector,muvera,muvera-python,retrieval
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Information Analysis
+Requires-Python: >=3.9
+Requires-Dist: numpy>=1.22.0
+Provides-Extra: dev
+Requires-Dist: mypy>=1.0; extra == 'dev'
+Requires-Dist: pre-commit>=3.0; extra == 'dev'
+Requires-Dist: pytest>=7.0; extra == 'dev'
+Requires-Dist: ruff>=0.4; extra == 'dev'
+Requires-Dist: scipy>=1.13.1; extra == 'dev'
+Description-Content-Type: text/markdown
+
+# MuVERA
+
+A Python implementation of **Mu**lti-**Ve**ctor **R**etrieval via Fixed Dimensional Encoding **A**lgorithm.
+
+Converts multi-vector embeddings (point clouds) into fixed-dimensional single vectors, enabling the use of existing single-vector search infrastructure (MIPS, ANN, etc.) as-is.
+
+## Why this library?
+
+The original MuVERA algorithm is described in a [research paper](https://arxiv.org/abs/2405.19504) and implemented in C++ within Google's [graph-mining](https://github.com/google/graph-mining) repository. While a Python reference exists, it exposes low-level config objects and separate functions for queries vs. documents, making it cumbersome to integrate into real workflows.
+
+This library wraps the full algorithm behind a **single `Muvera` class** with a minimal, intuitive interface — initialize once, then call `encode_documents()` and `encode_queries()`. No config dataclasses, no encoding-type enums, no manual seed juggling. Just NumPy arrays in, NumPy arrays out.
+
+## Installation
+
+```bash
+pip install muvera
+```
+
+Development install:
+
+```bash
+git clone https://github.com/craftsangjae/muvera-python.git
+cd muvera-python
+pip install -e .
+```
+
+## Quick Start
+
+```python
+import numpy as np
+from muvera import Muvera
+
+# Initialize encoder
+encoder = Muvera(
+    num_repetitions=10,
+    num_simhash_projections=4,
+    dimension=128,
+    seed=42,
+)
+
+# Encode documents (batch)
+# shape: (num_documents, num_vectors_per_doc, embedding_dim)
+documents = np.random.randn(100, 80, 128).astype(np.float32)
+doc_fdes = encoder.encode_documents(documents)  # (100, output_dimension)
+
+# Encode queries (batch)
+queries = np.random.randn(10, 32, 128).astype(np.float32)
+query_fdes = encoder.encode_queries(queries)  # (10, output_dimension)
+
+# Compute similarity (dot product)
+scores = query_fdes @ doc_fdes.T  # (10, 100)
+```
+
+## Parameters
+
+| Parameter | Default | Description |
+|---|---|---|
+| `num_repetitions` | 20 | Number of FDE repetitions. Higher values improve accuracy but increase output dimension |
+| `num_simhash_projections` | 5 | Number of SimHash projections. Number of partitions = 2^n |
+| `dimension` | 16 | Input embedding dimension |
+| `projection_type` | `"identity"` | `"identity"` or `"ams_sketch"` |
+| `projection_dimension` | None | Projected dimension when using AMS Sketch |
+| `fill_empty_partitions` | True | Whether to fill empty partitions with the nearest vector |
+| `final_projection_dimension` | None | Final dimension reduction via Count Sketch |
+| `seed` | 42 | Random seed for reproducibility |
+
+## Benchmark
+
+End-to-end retrieval on [NanoFiQA2018](https://huggingface.co/datasets/zeta-alpha-ai/NanoFiQA2018) (4598 documents, 50 queries) using `raphaelsty/neural-cherche-colbert` (dim=128):
+
+```
+=====================================================================================
+                                       RESULTS
+                            (zeta-alpha-ai/NanoFiQA2018)
+=====================================================================================
+Retriever                      | Index (s)    | Query (ms)   | Recall@25
+-------------------------------------------------------------------------------------
+ColBERT (Native MaxSim)        | 240.04       | 836.94       | 0.8400
+ColBERT + Muvera FDE           | 77.29        | 69.48        | 0.7600
+=====================================================================================
+```
+
+FDE achieves **90% of native MaxSim recall** while being **12x faster** at query time. See `examples/colbert_nanobeir.py` to reproduce.
+
+## Acknowledgments
+
+This library was inspired by [sionic-ai/muvera-py](https://github.com/sionic-ai/muvera-py), the first Python implementation of the MuVERA algorithm. Their faithful port of the C++ reference made it possible to validate correctness and understand the algorithm deeply. This project builds on that foundation with a simplified API designed for easier integration.
+
+## References
+
+- [MuVERA: Multi-Vector Retrieval via Fixed Dimensional Encodings](https://arxiv.org/abs/2405.19504)
+- [Google graph-mining C++ implementation](https://github.com/google/graph-mining/blob/main/sketching/point_cloud/fixed_dimensional_encoding.cc)
+- [sionic-ai/muvera-py](https://github.com/sionic-ai/muvera-py)

muvera_python-0.1.1/README.md

@@ -0,0 +1,93 @@
+# MuVERA
+
+A Python implementation of **Mu**lti-**Ve**ctor **R**etrieval via Fixed Dimensional Encoding **A**lgorithm.
+
+Converts multi-vector embeddings (point clouds) into fixed-dimensional single vectors, enabling the use of existing single-vector search infrastructure (MIPS, ANN, etc.) as-is.
+
+## Why this library?
+
+The original MuVERA algorithm is described in a [research paper](https://arxiv.org/abs/2405.19504) and implemented in C++ within Google's [graph-mining](https://github.com/google/graph-mining) repository. While a Python reference exists, it exposes low-level config objects and separate functions for queries vs. documents, making it cumbersome to integrate into real workflows.
+
+This library wraps the full algorithm behind a **single `Muvera` class** with a minimal, intuitive interface — initialize once, then call `encode_documents()` and `encode_queries()`. No config dataclasses, no encoding-type enums, no manual seed juggling. Just NumPy arrays in, NumPy arrays out.
+
+## Installation
+
+```bash
+pip install muvera
+```
+
+Development install:
+
+```bash
+git clone https://github.com/craftsangjae/muvera-python.git
+cd muvera-python
+pip install -e .
+```
+
+## Quick Start
+
+```python
+import numpy as np
+from muvera import Muvera
+
+# Initialize encoder
+encoder = Muvera(
+    num_repetitions=10,
+    num_simhash_projections=4,
+    dimension=128,
+    seed=42,
+)
+
+# Encode documents (batch)
+# shape: (num_documents, num_vectors_per_doc, embedding_dim)
+documents = np.random.randn(100, 80, 128).astype(np.float32)
+doc_fdes = encoder.encode_documents(documents)  # (100, output_dimension)
+
+# Encode queries (batch)
+queries = np.random.randn(10, 32, 128).astype(np.float32)
+query_fdes = encoder.encode_queries(queries)  # (10, output_dimension)
+
+# Compute similarity (dot product)
+scores = query_fdes @ doc_fdes.T  # (10, 100)
+```
+
+## Parameters
+
+| Parameter | Default | Description |
+|---|---|---|
+| `num_repetitions` | 20 | Number of FDE repetitions. Higher values improve accuracy but increase output dimension |
+| `num_simhash_projections` | 5 | Number of SimHash projections. Number of partitions = 2^n |
+| `dimension` | 16 | Input embedding dimension |
+| `projection_type` | `"identity"` | `"identity"` or `"ams_sketch"` |
+| `projection_dimension` | None | Projected dimension when using AMS Sketch |
+| `fill_empty_partitions` | True | Whether to fill empty partitions with the nearest vector |
+| `final_projection_dimension` | None | Final dimension reduction via Count Sketch |
+| `seed` | 42 | Random seed for reproducibility |
+
+## Benchmark
+
+End-to-end retrieval on [NanoFiQA2018](https://huggingface.co/datasets/zeta-alpha-ai/NanoFiQA2018) (4598 documents, 50 queries) using `raphaelsty/neural-cherche-colbert` (dim=128):
+
+```
+=====================================================================================
+                                       RESULTS
+                            (zeta-alpha-ai/NanoFiQA2018)
+=====================================================================================
+Retriever                      | Index (s)    | Query (ms)   | Recall@25
+-------------------------------------------------------------------------------------
+ColBERT (Native MaxSim)        | 240.04       | 836.94       | 0.8400
+ColBERT + Muvera FDE           | 77.29        | 69.48        | 0.7600
+=====================================================================================
+```
+
+FDE achieves **90% of native MaxSim recall** while being **12x faster** at query time. See `examples/colbert_nanobeir.py` to reproduce.
+
+## Acknowledgments
+
+This library was inspired by [sionic-ai/muvera-py](https://github.com/sionic-ai/muvera-py), the first Python implementation of the MuVERA algorithm. Their faithful port of the C++ reference made it possible to validate correctness and understand the algorithm deeply. This project builds on that foundation with a simplified API designed for easier integration.
+
+## References
+
+- [MuVERA: Multi-Vector Retrieval via Fixed Dimensional Encodings](https://arxiv.org/abs/2405.19504)
+- [Google graph-mining C++ implementation](https://github.com/google/graph-mining/blob/main/sketching/point_cloud/fixed_dimensional_encoding.cc)
+- [sionic-ai/muvera-py](https://github.com/sionic-ai/muvera-py)

muvera_python-0.1.1/muvera/__init__.py

@@ -0,0 +1,6 @@
+"""MuVERA: Multi-Vector Retrieval via Fixed Dimensional Encodings."""
+
+from muvera.muvera import Muvera
+
+__all__ = ["Muvera"]
+__version__ = "0.1.0"

muvera_python-0.1.1/muvera/helper.py

@@ -0,0 +1,205 @@
+"""Internal helper functions for MuVERA Fixed Dimensional Encoding.
+
+This module contains low-level utilities for Gray code manipulation,
+random projection matrix generation, Count Sketch, and SimHash-based
+partition indexing. These are not part of the public API.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+# ---------------------------------------------------------------------------
+# Gray code utilities
+# ---------------------------------------------------------------------------
+
+
+def append_to_gray_code(gray_code: int, bit: bool) -> int:
+    """Append a single bit to a Gray code value.
+
+    Parameters
+    ----------
+    gray_code : int
+        Current Gray code value.
+    bit : bool
+        Bit to append (True=1, False=0).
+
+    Returns
+    -------
+    int
+        Updated Gray code.
+    """
+    return (gray_code << 1) + (int(bit) ^ (gray_code & 1))
+
+
+def gray_code_to_binary(num: int) -> int:
+    """Convert a Gray code value to its binary representation.
+
+    Parameters
+    ----------
+    num : int
+        Gray code value.
+
+    Returns
+    -------
+    int
+        Corresponding binary representation.
+    """
+    mask = num >> 1
+    while mask != 0:
+        num = num ^ mask
+        mask >>= 1
+    return num
+
+
+# ---------------------------------------------------------------------------
+# Random projection matrices
+# ---------------------------------------------------------------------------
+
+
+def simhash_matrix_from_seed(dimension: int, num_projections: int, seed: int) -> np.ndarray:
+    """Generate a Gaussian random projection matrix for SimHash.
+
+    Parameters
+    ----------
+    dimension : int
+        Input vector dimension.
+    num_projections : int
+        Number of SimHash projections.
+    seed : int
+        Random seed.
+
+    Returns
+    -------
+    numpy.ndarray
+        Float32 matrix of shape ``(dimension, num_projections)``.
+    """
+    rng = np.random.default_rng(seed)
+    return rng.normal(loc=0.0, scale=1.0, size=(dimension, num_projections)).astype(np.float32)
+
+
+def ams_projection_matrix_from_seed(dimension: int, projection_dim: int, seed: int) -> np.ndarray:
+    """Generate an AMS Sketch projection matrix.
+
+    Each row has exactly one non-zero entry (+1 or -1), forming a sparse
+    random projection.
+
+    Parameters
+    ----------
+    dimension : int
+        Input vector dimension.
+    projection_dim : int
+        Output (projected) dimension.
+    seed : int
+        Random seed.
+
+    Returns
+    -------
+    numpy.ndarray
+        Float32 matrix of shape ``(dimension, projection_dim)``.
+    """
+    rng = np.random.default_rng(seed)
+    out = np.zeros((dimension, projection_dim), dtype=np.float32)
+    indices = rng.integers(0, projection_dim, size=dimension)
+    signs = rng.choice(np.array([-1.0, 1.0], dtype=np.float32), size=dimension)
+    out[np.arange(dimension), indices] = signs
+    return out
+
+
+def count_sketch_vector_from_seed(
+    input_vector: np.ndarray, final_dimension: int, seed: int
+) -> np.ndarray:
+    """Project a vector to a lower dimension using Count Sketch.
+
+    Parameters
+    ----------
+    input_vector : numpy.ndarray
+        Input vector (1-D).
+    final_dimension : int
+        Output dimension.
+    seed : int
+        Random seed.
+
+    Returns
+    -------
+    numpy.ndarray
+        Float32 vector of shape ``(final_dimension,)``.
+    """
+    rng = np.random.default_rng(seed)
+    out = np.zeros(final_dimension, dtype=np.float32)
+    indices = rng.integers(0, final_dimension, size=input_vector.shape[0])
+    signs = rng.choice(np.array([-1.0, 1.0], dtype=np.float32), size=input_vector.shape[0])
+    np.add.at(out, indices, signs * input_vector)
+    return out
+
+
+# ---------------------------------------------------------------------------
+# Partition indexing
+# ---------------------------------------------------------------------------
+
+
+def partition_index_gray(sketch_vector: np.ndarray) -> int:
+    """Compute a Gray-code-based partition index from a SimHash sketch vector.
+
+    Parameters
+    ----------
+    sketch_vector : numpy.ndarray
+        SimHash projection result vector (1-D).
+
+    Returns
+    -------
+    int
+        Partition index.
+    """
+    partition_index = 0
+    for val in sketch_vector:
+        partition_index = append_to_gray_code(partition_index, val > 0)
+    return partition_index
+
+
+def distance_to_partition(sketch_vector: np.ndarray, partition_index: int) -> int:
+    """Compute the Hamming distance between a sketch vector and a partition.
+
+    Parameters
+    ----------
+    sketch_vector : numpy.ndarray
+        SimHash projection result vector (1-D).
+    partition_index : int
+        Target partition index.
+
+    Returns
+    -------
+    int
+        Hamming distance.
+    """
+    num_projections = sketch_vector.size
+    binary_representation = gray_code_to_binary(partition_index)
+    sketch_bits = (sketch_vector > 0).astype(int)
+    binary_array = (binary_representation >> np.arange(num_projections - 1, -1, -1)) & 1
+    return int(np.sum(sketch_bits != binary_array))
+
+
+# ---------------------------------------------------------------------------
+# Vectorised batch helpers
+# ---------------------------------------------------------------------------
+
+
+def partition_indices_gray_batch(sketches: np.ndarray) -> np.ndarray:
+    """Compute Gray-code partition indices for a batch of sketch vectors.
+
+    Parameters
+    ----------
+    sketches : numpy.ndarray
+        SimHash sketch matrix of shape ``(N, num_projections)``.
+
+    Returns
+    -------
+    numpy.ndarray
+        Uint32 partition index array of shape ``(N,)``.
+    """
+    num_projections = sketches.shape[1]
+    bits = (sketches > 0).astype(np.uint32)
+    partition_indices = np.zeros(sketches.shape[0], dtype=np.uint32)
+    for bit_idx in range(num_projections):
+        partition_indices = (partition_indices << 1) + (bits[:, bit_idx] ^ (partition_indices & 1))
+    return partition_indices

muvera_python-0.1.1/muvera/muvera.py

@@ -0,0 +1,411 @@
+"""MuVERA (Multi-Vector Retrieval via Fixed Dimensional Encodings).
+
+This module provides the Fixed Dimensional Encoding (FDE) algorithm that
+converts multi-vector embeddings (point clouds) into single fixed-dimensional
+vectors.
+
+References
+----------
+.. [1] Google graph-mining: fixed_dimensional_encoding.cc
+   https://github.com/google/graph-mining/blob/main/sketching/point_cloud/fixed_dimensional_encoding.cc
+.. [2] sionic-ai/muvera-py: fde_generator.py
+   https://github.com/sionic-ai/muvera-py/blob/master/fde_generator.py
+"""
+
+from __future__ import annotations
+
+from typing import Literal
+
+import numpy as np
+
+from muvera.helper import (
+    ams_projection_matrix_from_seed,
+    count_sketch_vector_from_seed,
+    distance_to_partition,
+    gray_code_to_binary,
+    partition_index_gray,
+    partition_indices_gray_batch,
+    simhash_matrix_from_seed,
+)
+
+
+class Muvera:
+    """Encoder that converts multi-vector embeddings into Fixed Dimensional Encodings.
+
+    Uses the MuVERA algorithm to encode variable-length multi-vector
+    representations (point clouds) into fixed-dimensional single vectors.
+    The dot product between encoded vectors approximates the Chamfer
+    similarity (MaxSim) between the original multi-vector sets.
+
+    Parameters
+    ----------
+    num_repetitions : int, default=20
+        Number of repetitions for FDE generation. Higher values improve
+        accuracy but proportionally increase the output dimension.
+    num_simhash_projections : int, default=5
+        Number of SimHash projections. The number of partitions is
+        ``2 ** num_simhash_projections``. Must be in ``[0, 31)``.
+    dimension : int, default=16
+        Dimension of the input embedding vectors.
+    projection_type : {'identity', 'ams_sketch'}, default='identity'
+        Inner projection method.
+
+        - ``'identity'``: Uses the original dimension as-is.
+        - ``'ams_sketch'``: Uses AMS Sketch for dimensionality reduction.
+    projection_dimension : int or None, default=None
+        Projected dimension when ``projection_type='ams_sketch'``.
+        Ignored when ``'identity'``.
+    fill_empty_partitions : bool, default=True
+        Whether to fill empty partitions with the nearest vector during
+        document encoding.
+    final_projection_dimension : int or None, default=None
+        Final Count Sketch projection dimension. ``None`` disables
+        final projection.
+    seed : int, default=42
+        Random seed for reproducibility.
+
+    Attributes
+    ----------
+    output_dimension : int
+        Dimension of the encoded output vector.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> from muvera import Muvera
+    >>> encoder = Muvera(num_repetitions=10, num_simhash_projections=4,
+    ...                  dimension=128, seed=42)
+    >>> # Single document/query
+    >>> doc = np.random.randn(80, 128).astype(np.float32)
+    >>> query = np.random.randn(32, 128).astype(np.float32)
+    >>> doc_fde = encoder.encode_documents(doc)
+    >>> query_fde = encoder.encode_queries(query)
+    >>> score = query_fde @ doc_fde  # similarity score
+    >>>
+    >>> # Batch of documents with variable lengths
+    >>> docs = [np.random.randn(80, 128).astype(np.float32) for _ in range(5)]
+    >>> doc_fdes = encoder.encode_documents(docs)  # (5, output_dimension)
+    """
+
+    def __init__(
+        self,
+        num_repetitions: int = 20,
+        num_simhash_projections: int = 5,
+        dimension: int = 16,
+        projection_type: Literal["identity", "ams_sketch"] = "identity",
+        projection_dimension: int | None = None,
+        fill_empty_partitions: bool = True,
+        final_projection_dimension: int | None = None,
+        seed: int = 42,
+    ):
+        if num_repetitions <= 0:
+            raise ValueError(f"num_repetitions must be greater than 0, got {num_repetitions}")
+        if not (0 <= num_simhash_projections < 31):
+            raise ValueError(
+                f"num_simhash_projections must be in [0, 31), got {num_simhash_projections}"
+            )
+        if projection_type not in ("identity", "ams_sketch"):
+            raise ValueError(
+                f"projection_type must be 'identity' or 'ams_sketch', got '{projection_type}'"
+            )
+        if projection_type == "ams_sketch" and (
+            projection_dimension is None or projection_dimension <= 0
+        ):
+            raise ValueError(
+                "A positive projection_dimension is required when projection_type='ams_sketch'"
+            )
+
+        self.num_repetitions = num_repetitions
+        self.num_simhash_projections = num_simhash_projections
+        self.dimension = dimension
+        self.projection_type = projection_type
+        self.projection_dimension = projection_dimension
+        self.fill_empty_partitions = fill_empty_partitions
+        self.final_projection_dimension = final_projection_dimension
+        self.seed = seed
+
+        # Derived constants
+        self._num_partitions: int = 2**num_simhash_projections
+        self._use_identity: bool = projection_type == "identity"
+        self._proj_dim: int = dimension if self._use_identity else projection_dimension  # type: ignore[assignment]
+        self._fde_dim: int = num_repetitions * self._num_partitions * self._proj_dim
+
+    @property
+    def output_dimension(self) -> int:
+        """Dimension of the encoded output vector."""
+        if self.final_projection_dimension is not None and self.final_projection_dimension > 0:
+            return self.final_projection_dimension
+        return self._fde_dim
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def encode_documents(self, documents: np.ndarray | list[np.ndarray]) -> np.ndarray:
+        """Encode document point clouds into Fixed Dimensional Encodings.
+
+        Uses AVERAGE aggregation within partitions. Empty partitions are filled
+        with the nearest vector when ``fill_empty_partitions=True``.
+
+        Parameters
+        ----------
+        documents : np.ndarray or list[np.ndarray]
+            - Single document: shape ``(num_vectors, dimension)``
+            - Batch of documents: ``list[np.ndarray]`` where each element has
+              shape ``(num_vectors_i, dimension)``
+
+        Returns
+        -------
+        np.ndarray
+            Encoded FDE vector(s). Single document returns shape ``(output_dimension,)``,
+            batch returns shape ``(num_documents, output_dimension)``.
+        """
+        if isinstance(documents, list):
+            return self._encode_batch(documents, is_query=False)
+        return self._encode_single(documents, is_query=False)
+
+    def encode_queries(self, queries: np.ndarray | list[np.ndarray]) -> np.ndarray:
+        """Encode query point clouds into Fixed Dimensional Encodings.
+
+        Uses SUM aggregation within partitions. Empty partitions are not filled.
+
+        Parameters
+        ----------
+        queries : np.ndarray or list[np.ndarray]
+            - Single query: shape ``(num_vectors, dimension)``
+            - Batch of queries: ``list[np.ndarray]`` where each element has
+              shape ``(num_vectors_i, dimension)``
+
+        Returns
+        -------
+        np.ndarray
+            Encoded FDE vector(s). Single query returns shape ``(output_dimension,)``,
+            batch returns shape ``(num_queries, output_dimension)``.
+        """
+        if isinstance(queries, list):
+            return self._encode_batch(queries, is_query=True)
+        return self._encode_single(queries, is_query=True)
+
+    # ------------------------------------------------------------------
+    # Core Encoding
+    # ------------------------------------------------------------------
+
+    def _encode_single(self, point_cloud: np.ndarray, *, is_query: bool) -> np.ndarray:
+        """Encode a single point cloud."""
+        if point_cloud.ndim != 2 or point_cloud.shape[1] != self.dimension:
+            raise ValueError(f"Expected shape (N, {self.dimension}), got {point_cloud.shape}")
+
+        point_cloud = np.asarray(point_cloud, dtype=np.float32)
+        out = np.zeros(self._fde_dim, dtype=np.float32)
+
+        for rep in range(self.num_repetitions):
+            current_seed = self.seed + rep
+            sketches = self._compute_sketches(point_cloud, current_seed)
+            projected = self._compute_projection(point_cloud, current_seed)
+            rep_fde = self._aggregate_single(sketches, projected, is_query)
+
+            rep_start = rep * self._num_partitions * self._proj_dim
+            out[rep_start : rep_start + rep_fde.size] = rep_fde
+
+        return self._apply_final_projection(out)
+
+    def _encode_batch(self, point_clouds: list[np.ndarray], *, is_query: bool) -> np.ndarray:
+        """Encode a batch of variable-length point clouds."""
+        batch_size = len(point_clouds)
+        if batch_size == 0:
+            return np.zeros((0, self.output_dimension), dtype=np.float32)
+
+        # Validate and flatten
+        for i, pc in enumerate(point_clouds):
+            if pc.ndim != 2 or pc.shape[1] != self.dimension:
+                raise ValueError(f"Element {i}: expected (N, {self.dimension}), got {pc.shape}")
+
+        flat_points, doc_indices, doc_boundaries = self._flatten_batch(point_clouds)
+        out = np.zeros((batch_size, self._fde_dim), dtype=np.float32)
+
+        for rep in range(self.num_repetitions):
+            current_seed = self.seed + rep
+            sketches = self._compute_sketches(flat_points, current_seed)
+            projected = self._compute_projection(flat_points, current_seed)
+            rep_fde = self._aggregate_batch(
+                sketches, projected, doc_indices, doc_boundaries, is_query
+            )
+
+            rep_start = rep * self._num_partitions * self._proj_dim
+            out[:, rep_start : rep_start + rep_fde.size // batch_size] = rep_fde
+
+        return self._apply_final_projection(out)
+
+    # ------------------------------------------------------------------
+    # Helper Methods
+    # ------------------------------------------------------------------
+
+    def _compute_sketches(self, vectors: np.ndarray, seed: int) -> np.ndarray:
+        """Compute SimHash sketches for partition assignment."""
+        sim_matrix = simhash_matrix_from_seed(self.dimension, self.num_simhash_projections, seed)
+        return vectors @ sim_matrix
+
+    def _compute_projection(self, vectors: np.ndarray, seed: int) -> np.ndarray:
+        """Project vectors using identity or AMS sketch."""
+        if self._use_identity:
+            return vectors
+
+        ams_matrix = ams_projection_matrix_from_seed(self.dimension, self._proj_dim, seed)
+        return vectors @ ams_matrix
+
+    def _aggregate_single(
+        self, sketches: np.ndarray, projected: np.ndarray, is_query: bool
+    ) -> np.ndarray:
+        """Aggregate vectors into partitions for a single point cloud."""
+        num_points = sketches.shape[0]
+        partition_counts = np.zeros(self._num_partitions, dtype=np.int32)
+        rep_fde = np.zeros(self._num_partitions * self._proj_dim, dtype=np.float32)
+
+        # Assign vectors to partitions
+        for i in range(num_points):
+            pidx = partition_index_gray(sketches[i])
+            start = pidx * self._proj_dim
+            rep_fde[start : start + self._proj_dim] += projected[i]
+            partition_counts[pidx] += 1
+
+        # Apply AVERAGE for documents, SUM for queries
+        if not is_query:
+            self._apply_average_and_fill(rep_fde, partition_counts, sketches, projected)
+
+        return rep_fde
+
+    def _aggregate_batch(
+        self,
+        all_sketches: np.ndarray,
+        all_projected: np.ndarray,
+        doc_indices: np.ndarray,
+        doc_boundaries: np.ndarray,
+        is_query: bool,
+    ) -> np.ndarray:
+        """Aggregate vectors into partitions for a batch of point clouds."""
+        batch_size = len(doc_boundaries) - 1
+        part_indices = partition_indices_gray_batch(all_sketches)
+        partition_counts = np.zeros((batch_size, self._num_partitions), dtype=np.int32)
+        rep_fde = np.zeros((batch_size, self._num_partitions, self._proj_dim), dtype=np.float32)
+
+        # Count partitions
+        np.add.at(partition_counts, (doc_indices, part_indices), 1)
+
+        # Scatter-add projected vectors
+        self._scatter_add(rep_fde, doc_indices, part_indices, all_projected)
+
+        # Apply AVERAGE for documents
+        if not is_query:
+            self._apply_average_batch(rep_fde, partition_counts)
+            self._fill_empty_batch(
+                rep_fde, partition_counts, all_sketches, all_projected, doc_boundaries
+            )
+
+        return rep_fde.reshape(batch_size, -1)
+
+    def _apply_average_and_fill(
+        self,
+        rep_fde: np.ndarray,
+        partition_counts: np.ndarray,
+        sketches: np.ndarray,
+        projected: np.ndarray,
+    ) -> None:
+        """Apply AVERAGE aggregation and fill empty partitions for single cloud."""
+        num_points = sketches.shape[0]
+
+        for pidx in range(self._num_partitions):
+            start = pidx * self._proj_dim
+            if partition_counts[pidx] > 0:
+                rep_fde[start : start + self._proj_dim] /= partition_counts[pidx]
+            elif self.fill_empty_partitions and num_points > 0 and self.num_simhash_projections > 0:
+                distances = np.array(
+                    [distance_to_partition(sketches[j], pidx) for j in range(num_points)]
+                )
+                nearest = np.argmin(distances)
+                rep_fde[start : start + self._proj_dim] = projected[nearest]
+
+    def _apply_average_batch(self, rep_fde: np.ndarray, partition_counts: np.ndarray) -> None:
+        """Apply AVERAGE aggregation for batch."""
+        counts_3d = partition_counts[:, :, np.newaxis]
+        np.divide(rep_fde, counts_3d, out=rep_fde, where=counts_3d > 0)
+
+    def _fill_empty_batch(
+        self,
+        rep_fde: np.ndarray,
+        partition_counts: np.ndarray,
+        all_sketches: np.ndarray,
+        all_projected: np.ndarray,
+        doc_boundaries: np.ndarray,
+    ) -> None:
+        """Fill empty partitions for batch."""
+        if not self.fill_empty_partitions or self.num_simhash_projections == 0:
+            return
+
+        empty_docs, empty_parts = np.where(partition_counts == 0)
+        for doc_idx, pidx in zip(empty_docs, empty_parts):
+            doc_start, doc_end = doc_boundaries[doc_idx], doc_boundaries[doc_idx + 1]
+            if doc_start == doc_end:
+                continue
+
+            doc_sketches = all_sketches[doc_start:doc_end]
+            binary_rep = gray_code_to_binary(int(pidx))
+            target_bits = (binary_rep >> np.arange(self.num_simhash_projections - 1, -1, -1)) & 1
+            distances = np.sum((doc_sketches > 0).astype(int) != target_bits, axis=1)
+            nearest_local = np.argmin(distances)
+            rep_fde[doc_idx, pidx, :] = all_projected[doc_start + nearest_local]
+
+    def _scatter_add(
+        self,
+        rep_fde: np.ndarray,
+        doc_indices: np.ndarray,
+        part_indices: np.ndarray,
+        all_projected: np.ndarray,
+    ) -> None:
+        """Scatter-add projected vectors into partitions."""
+        doc_part = doc_indices * self._num_partitions + part_indices
+        base = doc_part * self._proj_dim
+        flat_rep_fde = rep_fde.reshape(-1)
+
+        for d in range(self._proj_dim):
+            np.add.at(flat_rep_fde, base + d, all_projected[:, d])
+
+    def _flatten_batch(
+        self, point_clouds: list[np.ndarray]
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """Flatten batch of point clouds into single array with metadata."""
+        doc_lengths = np.array([pc.shape[0] for pc in point_clouds], dtype=np.int32)
+        doc_boundaries = np.insert(np.cumsum(doc_lengths), 0, 0)
+        doc_indices = np.repeat(np.arange(len(point_clouds)), doc_lengths)
+        flat_points = np.vstack(point_clouds).astype(np.float32)
+
+        return flat_points, doc_indices, doc_boundaries
+
+    def _apply_final_projection(self, fdes: np.ndarray) -> np.ndarray:
+        """Apply optional Count Sketch final projection."""
+        if self.final_projection_dimension is None or self.final_projection_dimension <= 0:
+            return fdes
+
+        if fdes.ndim == 1:
+            return count_sketch_vector_from_seed(fdes, self.final_projection_dimension, self.seed)
+
+        # Batch
+        result = np.zeros((fdes.shape[0], self.final_projection_dimension), dtype=np.float32)
+        for i in range(fdes.shape[0]):
+            result[i] = count_sketch_vector_from_seed(
+                fdes[i], self.final_projection_dimension, self.seed
+            )
+        return result
+
+    def __repr__(self) -> str:  # noqa: D105
+        return (
+            f"Muvera("
+            f"num_repetitions={self.num_repetitions}, "
+            f"num_simhash_projections={self.num_simhash_projections}, "
+            f"dimension={self.dimension}, "
+            f"projection_type='{self.projection_type}', "
+            f"projection_dimension={self.projection_dimension}, "
+            f"fill_empty_partitions={self.fill_empty_partitions}, "
+            f"final_projection_dimension={self.final_projection_dimension}, "
+            f"seed={self.seed}"
+            f")"
+        )

muvera_python-0.1.1/pyproject.toml

@@ -0,0 +1,97 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "muvera-python"
+version = "0.1.1"
+description = "MuVERA: Multi-Vector Retrieval via Fixed Dimensional Encodings"
+readme = "README.md"
+license = "MIT"
+requires-python = ">=3.9"
+authors = [
+    { name = "craftsangjae" },
+]
+keywords = ["muvera", "muvera-python", "multi-vector", "retrieval", "embedding", "fixed-dimensional-encoding"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Scientific/Engineering :: Information Analysis",
+]
+dependencies = [
+    "numpy>=1.22.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=7.0",
+    "mypy>=1.0",
+    "ruff>=0.4",
+    "pre-commit>=3.0",
+    "scipy>=1.13.1",
+]
+
+[project.urls]
+Homepage = "https://github.com/craftsangjae/muvera-python"
+Repository = "https://github.com/craftsangjae/muvera-python"
+Issues = "https://github.com/craftsangjae/muvera-python/issues"
+
+[tool.hatch.build.targets.wheel]
+packages = ["muvera"]
+
+[tool.hatch.build.targets.sdist]
+include = ["muvera/", "README.md", "LICENSE"]
+
+# ---------------------------------------------------------------------------
+# Ruff
+# ---------------------------------------------------------------------------
+
+[tool.ruff]
+target-version = "py39"
+line-length = 100
+
+[tool.ruff.lint]
+select = [
+    "E",      # pycodestyle errors
+    "W",      # pycodestyle warnings
+    "F",      # pyflakes
+    "I",      # isort
+    "N",      # pep8-naming
+    "UP",     # pyupgrade
+    "B",      # flake8-bugbear
+    "SIM",    # flake8-simplify
+    "D",      # pydocstyle
+]
+ignore = [
+    "D100",   # Missing docstring in public module (module docstring already present)
+    "D104",   # Missing docstring in public package
+    "D203",   # 1 blank line required before class docstring (conflicts with D211)
+    "D213",   # Multi-line docstring summary should start at the second line (conflicts with D212)
+]
+
+[tool.ruff.lint.pydocstyle]
+convention = "numpy"
+
+[tool.ruff.lint.isort]
+known-first-party = ["muvera"]
+
+[tool.ruff.lint.per-file-ignores]
+"tests/**" = ["D"]
+"examples/**" = ["D"]
+
+# ---------------------------------------------------------------------------
+# Pytest
+# ---------------------------------------------------------------------------
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "-v --tb=short"