diff --git a/pyengine/engine/ingestion/build_index.py b/pyengine/engine/ingestion/build_index.py index 63c431d517..6ea9f43e1c 100644 --- a/pyengine/engine/ingestion/build_index.py +++ b/pyengine/engine/ingestion/build_index.py @@ -36,9 +36,10 @@ class DefaultIndex(Index): pass def build(self, d, vectors, vector_ids, DEVICE=INDEXDEVICES.CPU): - index = faiss.IndexFlatL2(d) # trained - index.add(vectors) - return index + index = faiss.IndexFlatL2(d) + index2 = faiss.IndexIDMap(index) + index2.add_with_ids(vectors, vector_ids) + return index2 class LowMemoryIndex(Index): diff --git a/pyengine/engine/ingestion/tests/test_build.py b/pyengine/engine/ingestion/tests/test_build.py index f4b92bcbbd..a4bb070bdb 100644 --- a/pyengine/engine/ingestion/tests/test_build.py +++ b/pyengine/engine/ingestion/tests/test_build.py @@ -16,14 +16,16 @@ class TestBuildIndex(unittest.TestCase): nb = 10000 nq = 100 _, xb, xq = get_dataset(d, nb, 500, nq) + ids = np.arange(xb.shape[0]) # Expected result index = faiss.IndexFlatL2(d) - index.add(xb) + index2 = faiss.IndexIDMap(index) + index2.add_with_ids(xb, ids) Dref, Iref = index.search(xq, 5) builder = DefaultIndex() - index2 = builder.build(d, xb) + index2 = builder.build(d, xb, ids) Dnew, Inew = index2.search(xq, 5) assert np.all(Dnew == Dref) and np.all(Inew == Iref)