add id

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):

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)