Error using assign method with IndexBinary in FAISS

[landEpita]

Description:

I am encountering an issue when attempting to use the assign method on an IndexBinary object in FAISS. The objective is to assign each input vector to the label of its nearest neighbor within the index. However, when I call the assign method, I receive a TypeError indicating a mismatch in the number or type of arguments.

Steps to Reproduce:

1. Import FAISS and other necessary libraries.
2. Generate random floating-point vectors and quantize them using sentence_transformers.quantization.quantize_embeddings.
3. Create an IndexBinary object and set its parameters.
4. Attempt to train the index with quantized vectors.
5. Initialize an empty numpy array for nearest neighbor labels.
6. Call the assign method with the quantized vectors and the labels array.

Expected Behavior:

The assign method should correctly assign each input vector to the label of its nearest neighbor within the index without any errors.

Actual Behavior:

When calling the assign method, the following TypeError is thrown:

TypeError: Wrong number or type of arguments for overloaded function 'IndexBinary_assign'.
Possible C/C++ prototypes are:
    faiss::IndexBinary::assign(faiss::idx_t,uint8_t const *,faiss::idx_t *,faiss::idx_t) const
    faiss::IndexBinary::assign(faiss::idx_t,uint8_t const *,faiss::idx_t *) const

Environment:

• FAISS version: 1.8.0
• Python version: 3.11.6

Additional Context:

I am using FAISS for vector similarity search in a binary space. The vectors are first quantized to a binary format before being used with the IndexBinary object. Below is a snippet of the code that leads to the error:

import faiss
from sentence_transformers.quantization import quantize_embeddings
import numpy as np

d = 1024                          
nb = 100000  
k=1

xb = np.random.random((nb, d)).astype('float32')
vectors = quantize_embeddings(xb, "ubinary")
count = len(vectors)

index = faiss.index_binary_factory(d, "BIVF64")
index.nprobe = 4
index.train(vectors)

list_nos = np.empty((count, k), dtype=np.int64)

index.assign(n=count, x=vectors, labels=list_nos, k=k)

I have ensured that the quantized vectors are of type uint8 as required. Any insights into what might be causing this error or how to resolve it would be greatly appreciated.

[mdouze]

Right, assign is not wrapped properly here
https://github.com/facebookresearch/faiss/blob/main/faiss/python/class_wrappers.py#L835
marking this as an enhancement.
In the meantime you can just use index.search(x, 1) to get the assignment.