[Python] Using an iterator to get the next(k) in FAISS similarity search

In FAISS after constructing an index (IVF, HNSW, or LSH) you can obtain the k-nearest neighbour based on a query q_x.

After my initial vector search in FAISS of let's say k = 50, I want to get more NN (i.e. 51, 52, 53..) until I satisfy the criteria I have set for the spatial calculations. (I am doing a combination of spatial and semantic similarity search.)

Is there a way you can do that getting the next(k) from the list of NN, without having to start the search all over again with k=51, then k=52 and so on?

import faiss
import numpy as np


d = 128 # dimension of vectors
index = faiss.IndexFlatL2(d)

# Add some vectors to the index
np.random.seed(123)
xb = np.random.random((1000, d)).astype('float32')
index.add(xb)

# Query vector
xq = np.random.random((1, d)).astype('float32')

def incremental_knn(index, xq, k, step_size):
"""Simulate an iterator-like behavior to get the next k-NN."""
start = 0
while True:
# Incrementally increase the number of nearest neighbors
k_next = start + step_size
D, I = index.search(xq, k_next) # Perform the search
yield I[0][start:k_next], D[0][start:k_next] # Return new results
start = k_next # Update starting point for the next iteration
if k_next >= k:
break

# Simulate fetching k-NN incrementally
step_size = 5 # Fetch 5 neighbors at a time
k = 20 # Total number of neighbors to fetch
for neighbors, distances in incremental_knn(index, xq, k, step_size):
print(f"Next batch of neighbors: {neighbors}, distances: {distances}")


The only way I have found so far is to increase the k and do the calculations with the new vectors.

Continue reading...