muvera-python 0.1.3__tar.gz → 0.2.0__tar.gz

muvera_python-0.2.0/.github/workflows/publish.yml

@@ -0,0 +1,144 @@
+name: Publish to PyPI
+
+on:
+  push:
+    branches: [main]
+
+jobs:
+  check-version:
+    runs-on: ubuntu-latest
+    outputs:
+      should_release: ${{ steps.check.outputs.should_release }}
+      version: ${{ steps.check.outputs.version }}
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: Check if version tag already exists
+        id: check
+        run: |
+          VERSION=$(grep '^version = ' pyproject.toml | cut -d'"' -f2)
+          echo "version=$VERSION" >> "$GITHUB_OUTPUT"
+          if git rev-parse "v$VERSION" >/dev/null 2>&1; then
+            echo "Tag v$VERSION already exists, skipping release"
+            echo "should_release=false" >> "$GITHUB_OUTPUT"
+          else
+            echo "Tag v$VERSION does not exist, proceeding with release"
+            echo "should_release=true" >> "$GITHUB_OUTPUT"
+          fi
+
+  test:
+    needs: check-version
+    if: needs.check-version.outputs.should_release == 'true'
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Install Rust toolchain
+        uses: dtolnay/rust-toolchain@stable
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install ".[dev]"
+
+      - name: Run tests
+        run: pytest -v
+
+  build-wheels:
+    needs: [check-version, test]
+    if: needs.check-version.outputs.should_release == 'true'
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - os: ubuntu-latest
+            target: x86_64
+          - os: ubuntu-latest
+            target: aarch64
+          - os: macos-latest
+            target: x86_64
+          - os: macos-latest
+            target: aarch64
+    runs-on: ${{ matrix.os }}
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Build wheels
+        uses: PyO3/maturin-action@v1
+        with:
+          target: ${{ matrix.target }}
+          args: --release --out dist --interpreter 3.9 3.10 3.11 3.12 3.13
+          manylinux: auto
+
+      - name: Upload wheels
+        uses: actions/upload-artifact@v4
+        with:
+          name: wheels-${{ matrix.os }}-${{ matrix.target }}
+          path: dist
+
+  build-sdist:
+    needs: [check-version, test]
+    if: needs.check-version.outputs.should_release == 'true'
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Build sdist
+        uses: PyO3/maturin-action@v1
+        with:
+          command: sdist
+          args: --out dist
+
+      - name: Upload sdist
+        uses: actions/upload-artifact@v4
+        with:
+          name: wheels-sdist
+          path: dist
+
+  publish:
+    needs: [check-version, build-wheels, build-sdist]
+    if: always() && needs.check-version.outputs.should_release == 'true' && (needs.build-wheels.result == 'success' || needs.build-sdist.result == 'success')
+    runs-on: ubuntu-latest
+    permissions:
+      id-token: write
+      contents: write
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Create version tag
+        run: |
+          VERSION=${{ needs.check-version.outputs.version }}
+          git tag "v$VERSION"
+          git push origin "v$VERSION"
+
+      - name: Download all artifacts
+        uses: actions/download-artifact@v4
+        with:
+          pattern: wheels-*
+          merge-multiple: true
+          path: dist
+
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1
+
+      - name: Create GitHub Release
+        uses: softprops/action-gh-release@v2
+        with:
+          tag_name: v${{ needs.check-version.outputs.version }}
+          files: dist/*
+          generate_release_notes: true
+          body: |
+            ## Installation
+            ```bash
+            pip install muvera-python==${{ needs.check-version.outputs.version }}
+            ```
+
+            See [PyPI](https://pypi.org/project/muvera-python/${{ needs.check-version.outputs.version }}/) for full package details.

muvera_python-0.2.0/.github/workflows/test.yml

@@ -0,0 +1,42 @@
+name: Tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.9", "3.10", "3.11", "3.12", "3.13"]
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install Rust toolchain
+        uses: dtolnay/rust-toolchain@stable
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install ".[dev]"
+
+      - name: Run ruff (lint)
+        run: ruff check .
+
+      - name: Run ruff (format check)
+        run: ruff format --check .
+
+      - name: Run mypy
+        run: mypy muvera
+
+      - name: Run pytest
+        run: pytest -v --tb=short

{muvera_python-0.1.3 → muvera_python-0.2.0}/.gitignore

@@ -21,4 +21,11 @@ examples/.cache/
 .vscode/
 
 # Benchmark results (generated, not committed)
-benchmarks/results/*.json
+benchmarks/results/*.json
+
+target/
+
+# Rust/maturin build artifacts
+*.so
+*.dylib
+*.pyd

muvera_python-0.2.0/.pre-commit-config.yaml

@@ -0,0 +1,24 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.9.7
+    hooks:
+      - id: ruff
+        args: [--fix, --exit-non-zero-on-fix]
+      - id: ruff-format
+
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.13.0
+    hooks:
+      - id: mypy
+        additional_dependencies: [numpy>=1.22.0]
+        files: ^muvera/
+
+  - repo: local
+    hooks:
+      - id: pytest
+        name: pytest
+        entry: pytest
+        language: system
+        pass_filenames: false
+        always_run: true
+        args: [tests/, -q]

muvera_python-0.2.0/CLAUDE.md

@@ -0,0 +1,253 @@
+# CLAUDE.md
+
+This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
+
+## Project Overview
+
+MuVERA (Multi-Vector Retrieval via Fixed Dimensional Encoding Algorithm) is a Python library that converts multi-vector embeddings (point clouds) into fixed-dimensional single vectors. This enables using existing single-vector search infrastructure (MIPS, ANN) without modification.
+
+### Project Goals
+
+**Simplicity First**: This project pursues a simple, intuitive interface to make MuVERA easy to use. Unlike the reference implementation which exposes low-level config objects and separate functions for queries vs. documents, this library wraps everything behind a single `Muvera` class:
+- `encode_documents()` - encodes document embeddings using AVERAGE method
+- `encode_queries()` - encodes query embeddings using SUM method
+
+No config dataclasses, no encoding-type enums, no manual seed juggling. Just NumPy arrays in, NumPy arrays out.
+
+**Distribution Plan**: The library will be published to PyPI for easy installation via `pip install muvera-python`.
+
+**Key use case**: Efficiently encode ColBERT-style multi-vector embeddings for retrieval without specialized infrastructure.
+
+## Git & PR Conventions
+- **Do NOT** add `Co-Authored-By` lines to commit messages.
+- **Do NOT** add "Generated with Claude Code" or similar attribution to PR descriptions.
+
+## Development Commands
+
+### Environment Activation
+**IMPORTANT**: Always activate the virtual environment before running any commands:
+```bash
+source .venv/bin/activate
+```
+
+### Post-Code-Writing Checklist
+**IMPORTANT**: After writing or modifying any code, always run:
+```bash
+ruff check .                  # Lint check (must pass)
+pytest                        # Tests (must pass)
+```
+
+### Setup
+```bash
+pip install maturin            # Required for building Rust extension
+maturin develop --release      # Build Rust extension (needs Rust toolchain)
+pip install -e ".[dev]"        # Install with dev dependencies
+```
+
+### Testing
+```bash
+pytest                        # Run all tests
+pytest tests/test_muvera.py   # Run specific test file
+pytest -v --tb=short          # Verbose output with short traceback
+pytest -k test_name           # Run specific test by name
+```
+
+### Code Quality
+```bash
+ruff check .                  # Lint all files
+ruff check --fix .            # Lint and auto-fix issues
+ruff format .                 # Format code
+mypy muvera                   # Type checking
+pre-commit run --all-files    # Run all pre-commit hooks
+```
+
+### Running Examples
+```bash
+python examples/basic_usage.py
+python examples/colbert_nanobeir.py
+```
+
+### Deployment
+
+**Version bump → merge → auto-release:**
+
+```bash
+# 1. Update version in pyproject.toml
+# version = "0.2.0"
+
+# 2. Create PR and merge to main
+# GitHub Actions will automatically:
+# - Detect the new version (tag doesn't exist yet)
+# - Run full test suite
+# - Create git tag v0.2.0
+# - Build wheel and sdist
+# - Publish to PyPI via OIDC
+# - Create GitHub Release
+#
+# If version is unchanged, all release steps are skipped.
+```
+
+**Local build (for testing):**
+```bash
+maturin build --release        # Build wheel with Rust extension
+```
+
+## Architecture
+
+### Core Components
+
+**`muvera/muvera.py`** - Main `Muvera` class implementing Fixed Dimensional Encoding (FDE)
+- Two encoding paths: single document, variable-length batch
+- Document encoding uses AVERAGE aggregation within partitions
+- Query encoding uses SUM aggregation within partitions
+- Optional final dimensionality reduction via Count Sketch
+- Hot-path methods (`_aggregate_single`, `_scatter_add`, `_fill_empty_batch`) delegate to Rust kernels when available
+
+**`muvera/helper.py`** - Low-level utilities (not public API)
+- Gray code manipulation for partition indexing
+- Random projection matrices (SimHash, AMS Sketch, Count Sketch)
+- Vectorized batch partition indexing
+- `partition_index_gray` and `partition_indices_gray_batch` delegate to Rust when available
+
+**`src/`** - Rust extension module (`muvera._rust_kernels`) via PyO3/maturin
+- `gray_code.rs` — Gray code append and binary conversion
+- `partition.rs` — Single and batch Gray-code partition indexing
+- `scatter.rs` — Scatter-add kernel for batch aggregation
+- `fill_empty.rs` — Single-point-cloud aggregation + batch empty partition filling
+- `lib.rs` — PyO3 module definition exposing 5 functions
+
+**`muvera/_rust_kernels.pyi`** - Type stubs for the Rust extension module
+
+### Algorithm Flow
+
+1. **SimHash Projection**: Maps each vector to a partition using random Gaussian projections
+2. **Partition Assignment**: Uses Gray code to assign vectors to one of `2^num_simhash_projections` partitions
+3. **Inner Projection**: Optionally reduces dimension via AMS Sketch (or uses identity)
+4. **Aggregation**:
+   - Documents: compute centroid (average) of vectors in each partition
+   - Queries: compute sum of vectors in each partition
+5. **Empty Partition Filling** (documents only): Fill empty partitions with nearest vector by Hamming distance
+6. **Repetitions**: Repeat steps 1-5 with different random seeds, concatenating results
+7. **Final Projection** (optional): Apply Count Sketch to reduce final dimension
+
+### Rust Acceleration
+
+Performance-critical inner loops are implemented in Rust via PyO3, with automatic fallback to pure Python:
+
+```python
+# muvera/__init__.py
+try:
+    import muvera._rust_kernels
+    _RUST_AVAILABLE = True
+except ImportError:
+    _RUST_AVAILABLE = False
+```
+
+**Accelerated functions:**
+| Rust function | Python fallback | Speedup |
+|---|---|---|
+| `aggregate_single` | `Muvera._aggregate_single_python` | 8-17x (single doc) |
+| `scatter_add_partitions` | `Muvera._scatter_add` (np.add.at loop) | 1-2.5x (batch) |
+| `fill_empty_partitions_batch` | `Muvera._fill_empty_batch` (Python loop) | 1-2.5x (batch) |
+| `partition_index_gray` | `helper._partition_index_gray_python` | part of aggregate |
+| `partition_indices_gray_batch` | `helper._partition_indices_gray_batch_python` | part of batch |
+
+**What is NOT in Rust** (intentionally kept in NumPy for seed compatibility):
+- `simhash_matrix_from_seed`, `ams_projection_matrix_from_seed` — depend on `np.random.default_rng`
+- `count_sketch_vector_from_seed` — same reason
+- `Muvera.__init__`, public API signatures — 100% unchanged
+
+### Batch Processing
+
+The library supports two input formats:
+- **Single**: `(num_vectors, dimension)` - processes one point cloud
+- **Variable-length batch**: `list[np.ndarray]` - each point cloud has different length (recommended for real-world data)
+
+Variable-length batch processing flattens all point clouds, processes them together, then aggregates per-document using Rust `scatter_add_partitions` (or `np.add.at()` fallback).
+
+## Code Conventions
+
+### Python
+- NumPy-style docstrings (configured in pyproject.toml)
+- Type hints required (Python 3.9+ syntax with `|` for unions)
+- Line length: 100 characters
+- Use `np.float32` for all embeddings (memory efficiency)
+- Use `np.uint32` for partition indices
+- Random number generation via `np.random.default_rng(seed)` for reproducibility
+
+### Rust
+- Edition 2021
+- Dependencies: `pyo3` 0.23, `numpy` 0.23 (Rust crate, not Python package), `ndarray` 0.16
+- All three crates are version-locked together (upgrade all at once)
+- Use `f32` for all floating-point data, `u32` for partition indices, `i32` for counts, `i64` for boundaries
+- PyO3 functions accept `PyReadonlyArray*` for input arrays and `&Bound<PyArray*>` for in-place mutation
+
+## Testing
+
+### Test Organization
+
+- **`test_helper.py`**: Low-level helper function tests (Gray code, projections, etc.)
+- **`test_muvera.py`**: Core Muvera class tests (shapes, validation, reproducibility)
+- **`test_reference.py`**: Validation against reference implementation (sionic-ai/muvera-py)
+- **`test_real_colbert.py`**: Real-world ColBERT embedding tests using NanoBEIR fixtures
+- **`test_rust_equivalence.py`**: Numerical equivalence tests between Rust kernels and Python fallbacks (skipped if Rust extension is unavailable)
+
+### Real Data Testing
+
+`test_real_colbert.py` uses cached ColBERT embeddings from NanoBEIR to validate performance on real data:
+- **Fixtures**: 35 documents, 5 queries, 35 relevance judgments (~2.2MB cached in git)
+- **Tests**: FDE encoding, correlation with native MaxSim, Recall@K metrics
+- **Generation**: `python scripts/generate_test_fixtures.py` (requires pylate, datasets, torch)
+
+Fixtures are cached in `tests/fixtures/colbert_nanobeir/` to avoid slow model inference during CI/testing.
+
+## Key Parameters
+
+- `num_repetitions`: Controls accuracy/dimension trade-off (default: 20)
+- `num_simhash_projections`: Determines partition count as `2^n` (default: 5, range: [0, 31))
+- `fill_empty_partitions`: Whether to fill empty partitions with nearest vector (default: True, documents only)
+- `projection_type`: "identity" or "ams_sketch" for dimensionality reduction
+- `final_projection_dimension`: Optional Count Sketch final projection
+
+Output dimension: `num_repetitions * 2^num_simhash_projections * projection_dimension` (or `final_projection_dimension` if set)
+
+## CI/CD and Deployment
+
+### GitHub Actions Workflows
+
+**`.github/workflows/test.yml`** - Continuous Integration
+- Triggers: Push to main, all pull requests
+- Tests across Python 3.9-3.13
+- Installs Rust toolchain via `dtolnay/rust-toolchain@stable`
+- Builds Rust extension via `pip install ".[dev]"` (maturin build backend)
+- Runs ruff (lint + format check), mypy (type checking), pytest
+
+**`.github/workflows/publish.yml`** - PyPI Publishing
+- Triggers: Push to main
+- Checks if `v{version}` tag already exists; skips release if it does
+- Runs full test suite
+- Builds cross-platform wheels via `PyO3/maturin-action@v1` (Linux x86_64/aarch64, macOS x86_64/aarch64, Windows x86_64)
+- Builds sdist separately
+- Creates git tag, publishes to PyPI via OIDC, creates GitHub Release
+
+### Deployment Policy
+
+**Auto-release on version bump:**
+- Merging a PR that changes the version in `pyproject.toml` triggers a release
+- If the version is unchanged, all release steps are skipped
+- Uses OpenID Connect (OIDC) for secure, token-free authentication to PyPI
+
+**Pre-deployment checklist:**
+1. Update `version` in `pyproject.toml`
+2. Run tests locally: `pytest`
+3. Check code quality: `ruff check . && mypy muvera`
+4. Create PR and merge to main
+
+**OIDC Setup (one-time):**
+1. Go to [PyPI](https://pypi.org) → Account settings → Publishing
+2. Add a new pending publisher:
+   - PyPI Project Name: `muvera-python`
+   - Owner: `craftsangjae`
+   - Repository: `muvera-python`
+   - Workflow: `publish.yml`
+   - Environment: (leave empty)