faiss_l2_index_knn_search_quickstart.py

1import numpy as np
2import faiss
3
4# Step 1: Prepare data
5d = 64                           # dimension
6nb = 100000                      # database size
7nq = 10000                       # nb of queries
8np.random.seed(1234)             # make reproducible
9xb = np.random.random((nb, d)).astype('float32')
10xb[:, 0] += np.arange(nb) / 1000.
11xq = np.random.random((nq, d)).astype('float32')
12xq[:, 0] += np.arange(nq) / 1000.
13
14# Step 2: Build the index
15index = faiss.IndexFlatL2(d)   # build a flat (CPU) index
16print(f"Is the index trained? {index.is_trained}")
17
18# Step 3: Add vectors to the index
19index.add(xb)                  # add vectors to the index
20print(f"Total vectors in index: {index.ntotal}")
21
22# Step 4: Search for the k-nearest neighbors
23k = 4                          # we want to see 4 nearest neighbors
24D, I = index.search(xq[:5], k) # search for the first 5 query vectors
25
26# Print results
27print("Indices of nearest neighbors for the first 5 queries:")
28print(I)
29print("L2 distances to these neighbors:")
30print(D)