IVF PQ

1
cdef class Index

1
def trained(self)

1
def n_lists(self)

1
def dim(self)

1
def pq_dim(self)

1
def pq_len(self)

1
def pq_bits(self)

1
def centers(self)

1
def centers_padded(self)

1
def pq_centers(self)

1
def centers_rot(self)

1
def rotation_matrix(self)

1
def list_sizes(self)

1
def lists(self, resources=None)

1
def list_data(self, label, n_rows=0, offset=0, out_codes=None, resources=None)

1
def list_indices(self, label, n_rows=0)

1
cdef class IndexParams

1
def __init__(self, *, n_lists=1024, metric="sqeuclidean", metric_arg=2.0, kmeans_n_iters=20, kmeans_trainset_fraction=0.5, pq_bits=8, pq_dim=0, codebook_kind="subspace", force_random_rotation=False, add_data_on_build=True, conservative_memory_allocation=False, max_train_points_per_pq_code=256, codes_layout="interleaved")

1
def get_handle(self)

1
def metric(self)

1
def metric_arg(self)

1
def add_data_on_build(self)

1
def n_lists(self)

1
def kmeans_n_iters(self)

1
def kmeans_trainset_fraction(self)

1
def pq_bits(self)

1
def pq_dim(self)

1
def codebook_kind(self)

1
def force_random_rotation(self)

1
def add_data_on_build(self)

1
def conservative_memory_allocation(self)

1
def max_train_points_per_pq_code(self)

1
def codes_layout(self)

1
def get_handle(self)

1
cdef class SearchParams

1
def __init__(self, *, n_probes=20, lut_dtype=np.float32, internal_distance_dtype=np.float32, coarse_search_dtype=np.float32, max_internal_batch_size=4096)

1
def get_handle(self)

1
def n_probes(self)

1
def lut_dtype(self)

1
def internal_distance_dtype(self)

1
def coarse_search_dtype(self)

1
def max_internal_batch_size(self)

1
def get_handle(self)

1
def build(IndexParams index_params, dataset, resources=None)

1
>>> import cupy as cp
2
>>> from cuvs.neighbors import ivf_pq
3
>>> n_samples = 50000
4
>>> n_features = 50
5
>>> n_queries = 1000
6
>>> k = 10
7
>>> dataset = cp.random.random_sample((n_samples, n_features),
8
...                                   dtype=cp.float32)
9
>>> build_params = ivf_pq.IndexParams(metric="sqeuclidean")
10
>>> index = ivf_pq.build(build_params, dataset)
11
>>> distances, neighbors = ivf_pq.search(ivf_pq.SearchParams(),
12
...                                        index, dataset,
13
...                                        k)
14
>>> distances = cp.asarray(distances)
15
>>> neighbors = cp.asarray(neighbors)

1
def build_precomputed(IndexParams index_params, uint32_t dim, pq_centers, centers, centers_rot, rotation_matrix, resources=None)

1
>>> import cupy as cp
2
>>> from cuvs.neighbors import ivf_pq
3
>>> n_lists = 100
4
>>> dim = 128
5
>>> pq_dim = 16
6
>>> pq_bits = 8
7
>>> pq_len = (dim + pq_dim - 1) // pq_dim  # ceil division
8
>>> pq_book_size = 1 << pq_bits
9
>>> rot_dim = pq_len * pq_dim
10
>>> dim_ext = ((dim + 1 + 7) // 8) * 8  # round_up(dim + 1, 8)
11
>>>
12
>>> # Prepare precomputed matrices (example with random data)
13
>>> pq_centers = cp.random.random((pq_dim, pq_len, pq_book_size),
14
...                               dtype=cp.float32)
15
>>> centers = cp.random.random((n_lists, dim_ext), dtype=cp.float32)
16
>>> centers_rot = cp.random.random((n_lists, rot_dim), dtype=cp.float32)
17
>>> rotation_matrix = cp.random.random((rot_dim, dim), dtype=cp.float32)
18
>>>
19
>>> # Build index from precomputed data
20
>>> build_params = ivf_pq.IndexParams(n_lists=n_lists, pq_dim=pq_dim,
21
...                                    pq_bits=pq_bits,
22
...                                    codebook_kind="subspace")
23
>>> index = ivf_pq.build_precomputed(build_params, dim, pq_centers,
24
...                                   centers, centers_rot, rotation_matrix)

1
def extend(Index index, new_vectors, new_indices, resources=None)

1
>>> import cupy as cp
2
>>> from cuvs.neighbors import ivf_pq
3
>>> n_samples = 50000
4
>>> n_features = 50
5
>>> n_queries = 1000
6
>>> dataset = cp.random.random_sample((n_samples, n_features),
7
...                                   dtype=cp.float32)
8
>>> index = ivf_pq.build(ivf_pq.IndexParams(), dataset)
9
>>> n_rows = 100
10
>>> more_data = cp.random.random_sample((n_rows, n_features),
11
...                                     dtype=cp.float32)
12
>>> indices = n_samples + cp.arange(n_rows, dtype=cp.int64)
13
>>> index = ivf_pq.extend(index, more_data, indices)
14
>>> # Search using the built index
15
>>> queries = cp.random.random_sample((n_queries, n_features),
16
...                                   dtype=cp.float32)
17
>>> distances, neighbors = ivf_pq.search(ivf_pq.SearchParams(),
18
...                                      index, queries,
19
...                                      k=10)

1
def load(filename, resources=None)

1
def save(filename, Index index, bool include_dataset=True, resources=None)

1
>>> import cupy as cp
2
>>> from cuvs.neighbors import ivf_pq
3
>>> n_samples = 50000
4
>>> n_features = 50
5
>>> dataset = cp.random.random_sample((n_samples, n_features),
6
...                                   dtype=cp.float32)
7
>>> # Build index
8
>>> index = ivf_pq.build(ivf_pq.IndexParams(), dataset)
9
>>> # Serialize and deserialize the ivf_pq index built
10
>>> ivf_pq.save("my_index.bin", index)
11
>>> index_loaded = ivf_pq.load("my_index.bin")

1
def search(SearchParams search_params, Index index, queries, k, neighbors=None, distances=None, resources=None)

1
>>> import cupy as cp
2
>>> from cuvs.neighbors import ivf_pq
3
>>> n_samples = 50000
4
>>> n_features = 50
5
>>> n_queries = 1000
6
>>> dataset = cp.random.random_sample((n_samples, n_features),
7
...                                   dtype=cp.float32)
8
>>> # Build the index
9
>>> index = ivf_pq.build(ivf_pq.IndexParams(), dataset)
10
>>>
11
>>> # Search using the built index
12
>>> queries = cp.random.random_sample((n_queries, n_features),
13
...                                   dtype=cp.float32)
14
>>> k = 10
15
>>> search_params = ivf_pq.SearchParams(n_probes=20)
16
>>>
17
>>> distances, neighbors = ivf_pq.search(search_params, index, queries,
18
...                                     k)

1
def transform(Index index, input_dataset, output_labels=None, output_dataset=None, resources=None)