faiss: как получить вектор по id из python

Question

У меня есть индекс faiss, и я хочу использовать некоторые вложения в своем скрипте Python. Выбор Embeddings должен быть сделан по id. Поскольку faiss написан на C++, в качестве API используется swig.

Я предполагаю, что функция, которая мне нужна, - это реконструировать :

      /** Reconstruct a stored vector (or an approximation if lossy coding)
     *
     * this function may not be defined for some indexes
     * @param key         id of the vector to reconstruct
     * @param recons      reconstucted vector (size d)
     */
    virtual void reconstruct(idx_t key, float* recons) const;

Поэтому я вызываю этот метод в питоне, например:

      vector = index.reconstruct(0)

Но это приводит к следующей ошибке:

vector = index.reconstruct(0)File"lib/python3.8/site-packages/faiss/init.py " , строка 406, в replace_reconstructself.reconstruct_c(key, swig_ptr(x)) File "lib/python3.8/site-packages/faiss/swigfaiss.py", строка 1897, при реконструкции возвращается _swigfaiss.IndexFlat_reconstruct(self, key, recons)

TypeError: в методе «IndexFlat_reconstruct», аргумент 2 типа «faiss::Index::idx_t» python-BaseException

Кто-нибудь знает, что не так с моим подходом?

Answer 1

Это единственный способ, который я нашел вручную.

      import faiss
import numpy as np

a = np.random.uniform(size=30)
a = a.reshape(-1,10).astype(np.float32)
d = 10
index = faiss.index_factory(d,'Flat', faiss.METRIC_L2)
index.add(a)

xb = index.xb
print(xb.at(0) == a[0][0])

Выход:

Вы можете получить любой вектор с циклом

      required_vector_id = 1
vector = np.array([xb.at(required_vector_id*index.d + i) for i in range(index.d)])
    
print(np.all(vector== a[1]))

Выход:

      True

Answer 2

Вы можете получить все вложения, которые вы добавили в индекс, используя это:

      # Number of docs added to your index
num_docs = index.ntotal
# Get the dimension of your embeddings
embedding_dimension = index.d

embeddings = faiss.rev_swig_ptr(index.get_xb(), num_docs*embedding_dimension).reshape(num_docs, embedding_dimension)

Ссылка

• Тестовый пример FAISS в репозитории GitHub