hnswlib_vector_index_knn_search_quickstart.py

1import hnswlib
2import numpy as np
3
4dim = 128
5num_elements = 10000
6
7# Generating sample data
8data = np.float32(np.random.random((num_elements, dim)))
9ids = np.arange(num_elements)
10
11# Declaring index
12p = hnswlib.Index(space='l2', dim=dim)  # possible spaces are l2, cosine or ip
13
14# Initing index - the maximum number of elements should be known beforehand
15p.init_index(max_elements=num_elements, ef_construction=200, M=16)
16
17# Element insertion (can be called in several batches)
18p.add_items(data, ids)
19
20# Controlling the recall by setting ef:
21p.set_ef(50) # ef should always be > k
22
23# Query dataset, asking for 10 closest elements
24labels, distances = p.knn_query(data, k=10)
25
26print("Recall for the query:", np.mean(labels[:, 0] == ids))