ann.py

# https://github.com/matsui528/faiss_tips
import faiss
import numpy as np

D = 256
N = 10000
X = np.random.random((N, D)).astype(np.float32)

# Setup
index = faiss.index_factory(D, "IVF25,Flat", faiss.METRIC_INNER_PRODUCT)
index.train(X)
index.add(X)

# Search
topk = 4
# Use the top three vectors for querying
dists, ids = index.search(x=X[:3], k=topk)
# <class 'numpy.ndarray'> float32 (3, 4)
print(type(dists), dists.dtype, dists.shape)
print(type(ids), ids.dtype, ids.shape)

# Show params
print("N:", index.ntotal)
print("D:", index.d)

faiss.write_index(index, "/tmp/test.ivf")