vicinity 0.11.0

Approximate nearest-neighbor search
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76

# /// script
# requires-python = ">=3.10"
# dependencies = ["pyvicinity", "numpy"]
# ///
"""Drop pyvicinity into the ann-benchmarks harness contract.

Anyone benchmarking ANN libraries against ann-benchmarks /
big-ann-benchmarks / VIBE expects a class with ``fit``, ``query``, and
``batch_query`` methods. ``pyvicinity.ann_benchmarks`` provides HNSW and
IVF-PQ wrappers for that contract. This example exercises the full harness
flow against a 5k synthetic corpus.

Run with:

    uv run examples/python/03_ann_benchmarks_harness.py
"""

from __future__ import annotations

import numpy as np

from pyvicinity.ann_benchmarks import VicinityHNSW, VicinityIVFPQ


def run_harness_flow(name: str, algo, train: np.ndarray, test: np.ndarray) -> None:
    # 1. fit() -- the harness builds the index once.
    algo.fit(train)
    print(f"{name} fit ok:  {algo}")

    # 2. set_query_arguments() -- one search parameter value per recall point.
    if isinstance(algo, VicinityIVFPQ):
        algo.set_query_arguments(8, rerank_pool=100)
    else:
        algo.set_query_arguments(50)

    # 3. query() -- single-query path used by erikbern/ann-benchmarks.
    ids = algo.query(test[0], 10)
    print(f"{name} single-query ids[:10]: {ids.tolist()}")

    # 4. batch_query() -- preferred path for big-ann-benchmarks and VIBE.
    algo.batch_query(test, 10)
    batch = algo.get_batch_results()
    print(f"{name} batch_query shape:     {batch.shape}")

    # 5. done() -- some harnesses release the index between configs.
    algo.done()


def main() -> None:
    rng = np.random.default_rng(0)
    train = rng.standard_normal((5_000, 32), dtype=np.float32)
    test = rng.standard_normal((50, 32), dtype=np.float32)

    run_harness_flow(
        "hnsw", VicinityHNSW("cosine", {"M": 16, "efConstruction": 100}), train, test
    )
    run_harness_flow(
        "ivfpq",
        VicinityIVFPQ(
            "cosine",
            {
                "num_clusters": 32,
                "num_codebooks": 8,
                "codebook_size": 32,
                "training_sample_size": 1_000,
                "kmeans_max_iter": 5,
                "nprobe": 8,
            },
        ),
        train,
        test,
    )


if __name__ == "__main__":
    main()