mojolearn.IVFRaBitQIndex#

class mojolearn.IVFRaBitQIndex(n_lists, n_probes, n_neighbors=8, kmeans_n_iters=20, random_state=0)[source]#

Bases: _AnnResident, _AnnSaved, NumericModeMixin

Approximate nearest-neighbor index with inverted lists and 1-bit RaBitQ codes (IVF-RaBitQ).

Follows cuVS’s ivf_rabitq (RaBitQ, Gao and Long). fit trains a k-means coarse quantizer, rotates each row’s residual by a randomized Hadamard transform, and keeps its sign bits together with its norm and an inner-product correction factor. search scans the n_probes lists nearest each query and returns the RaBitQ distance estimates, the n_neighbors smallest by (estimate, id). The distances are estimates; pass the candidates to refine for exact distances. Results are reproducible. On a GPU the fitted index stays on the device between searches; each search uploads only its queries.

Parameters:

  • n_lists (int) – Number of inverted lists (k-means clusters). Must be in [1, n_samples].

  • n_probes (int) – Number of lists scanned per query. Required, with no default.

  • n_neighbors (int, default 8) – Number of neighbors search returns per query.

  • kmeans_n_iters (int, default 20) – k-means iterations for the coarse quantizer.

  • random_state (int, default 0) – Seed for the k-means initialization and the rotation’s random signs.

  • numeric_mode (str or None, default None) – Numeric mode for this estimator; None uses the package default set with mojolearn.set_numeric_mode.

Variables:

  • centers_ (ndarray of shape (n_lists, n_features), float32) – Coarse quantizer centers.

  • list_offsets_ (ndarray of shape (n_lists + 1,), int32) – Start of each list in list_indices_.

  • list_indices_ (ndarray of shape (n_samples,), int32) – Original row ids, grouped by list.

  • codes_ (ndarray of shape (n_samples, n_words), int32) – Sign bits of each rotated residual, packed 32 per word; the residual is zero-padded to the next power of two.

  • norms_ (ndarray of shape (n_samples,), float32) – Norm of each residual.

  • ip_factors_ (ndarray of shape (n_samples,), float32) – Inner-product correction factor of each row.

  • n_features_in_ (int) – Number of features seen during fit.

  • n_rows_ (int) – Number of indexed rows.

  • n_lists_ (int) – The fitted value of n_lists.

  • seed_ (int) – The seed used for the rotation.

  • n_candidates_ (ndarray of shape (n_queries,), int32) – Set by search: how many indexed rows each query scanned.

Raises:

  • TypeError – If an integer parameter is not an int.

  • ValueError – If n_lists is out of range.

See also

refine

Re-rank candidates with exact distances.

Examples

>>> import numpy as np, mojolearn
>>> X = np.random.default_rng(0).random((1000, 8), dtype=np.float32)
>>> index = mojolearn.IVFRaBitQIndex(n_lists=16, n_probes=4, n_neighbors=20).fit(X)
>>> _, candidates = index.search(X[:10])
>>> distances, indices = mojolearn.refine(X, X[:10], candidates, k=5)
__init__(n_lists, n_probes, n_neighbors=8, kmeans_n_iters=20, random_state=0)[source]#

Methods

fit(X[, y])

Build the index from X.

get_params([deep])

load(path)

Load an index written by save.

numeric_mode_used()

The numeric mode this estimator runs in.

save(path)

Write the fitted index to path as an .npz file.

search(queries[, filter])

Find the approximate nearest indexed rows of each query.

set_params(**params)

Set constructor parameters.

vendor_used()

The GPU backend this estimator runs on.
fit(X, y=None)[source]#

Build the index from X.

Parameters:

  • X (array-like of shape (n_samples, n_features)) – Rows to index, converted to float32.

  • y (None) – Ignored.

Returns:

self (IVFRaBitQIndex) – The fitted index.

search(queries, filter=None)[source]#

Find the approximate nearest indexed rows of each query.

Parameters:

  • queries (array-like of shape (n_queries, n_features)) – Query rows, converted to float32.

  • filter (array-like of bool, shape (n_rows_,), optional) – One flag per indexed row; a row flagged False is never returned (cuVS’s sample filter, applied before scoring).

Returns:

  • distances (ndarray of shape (n_queries, n_neighbors), float32) – Estimated squared Euclidean distances, ascending.

  • indices (ndarray of shape (n_queries, n_neighbors), int32) – Original row ids of the neighbors. A query whose probed lists hold fewer than n_neighbors rows is padded with distance inf and id -1.

Raises:

ValueError – If the index is not fitted, the queries have a different number of features, or filter has the wrong shape or dtype.