diff --git a/pyengine/engine/faiss_wrapper/faiss_wrapper.py b/pyengine/engine/faiss_wrapper/faiss_wrapper.py index 6a60e53d35ba091bed4c1795b53a7ab1a16f5760..576f22fa7a9d6b99f5c09f14ae9c3d0be1f5338e 100644 --- a/pyengine/engine/faiss_wrapper/faiss_wrapper.py +++ b/pyengine/engine/faiss_wrapper/faiss_wrapper.py @@ -2,27 +2,32 @@ import faiss import numpy as np class FaissSearch(): - def __init__(self, index, id_to_vector_map): - pass + def __init__(self, index, id_to_vector_map=None): + self.__index = index + if id_to_vector_map is None: + self.__id_to_vector_map = [] # def search_by_ids(self, id_list, k): # pass def search_by_vectors(self, vector_list, k): - # return both id and vector - pass - - def __search__(self, id_list, vector_list, k): - pass - - -class FaissIndex(): - def build_index(self, vector_list, dimension): - # return index - pass - - # def build_index_cpu(self): - # pass - - # def build_index_gpu(self): - # pass \ No newline at end of file + id_list = [None] * len(vector_list) + + result = self.__search(id_list, vector_list, k) + return result + + def __search(self, id_list, vector_list, k): + D, I = self.__index.search(vector_list, k) + return I + + +# class FaissIndex(): +# def build_index(self, vector_list, dimension): +# # return index +# pass +# +# def build_index_cpu(self): +# pass +# +# def build_index_gpu(self): +# pass \ No newline at end of file