def testAddWithMetaData(index, x, s, out, algo, distmethod):
if index != None:
i = SPTAG.AnnIndex.Load(index)
else:
i = SPTAG.AnnIndex(algo, 'Float', x.shape[1])
i = SPTAG.AnnIndex(algo, 'Float', x.shape[1])
i.SetBuildParam("NumberOfThreads", '4')
i.SetBuildParam("DistCalcMethod", distmethod)
if i.AddWithMetaData(x, s, x.shape[0]):
i.Save(out)
def build_advanced_index(self, vecs: 'np.ndarray') -> 'SPTAG.AnnIndex':
"""Build an advanced index structure from a numpy array.
:param vecs: numpy array containing the vectors to index
:return: advanced index
"""
import SPTAG
_index = SPTAG.AnnIndex(self.method, 'Float', vecs.shape[1])
# Set the parameters
_index.SetBuildParam("Samples", str(self.samples))
_index.SetBuildParam("TPTNumber", str(self.tpt_number))
_index.SetBuildParam("TPTLeafSize", str(self.tpt_leaf_size))
_index.SetBuildParam("NeighborhoodSize", str(self.neighborhood_size))
_index.SetBuildParam("GraphNeighborhoodScale",
str(self.graph_neighborhood_size))
_index.SetBuildParam("CEF", str(self.cef))
_index.SetBuildParam("MaxCheckForRefineGraph",
str(self.max_check_for_refined_graph))
_index.SetBuildParam("NumberOfThreads", str(self.num_threads))
_index.SetBuildParam("DistCalcMethod", self.space)
_index.SetSearchParam("MaxCheck", str(self.max_check))
_index.SetBuildParam("BKTNumber", str(self.bkt_number))
_index.SetBuildParam("BKTMeansK", str(self.bkt_meansk))
_index.SetBuildParam("KDTNumber", str(self.kdt_number))
if _index.Build(vecs.astype('float32'), vecs.shape[0]):
return _index
def build_advanced_index(self, vecs: 'np.ndarray'):
import SPTAG
_index = SPTAG.AnnIndex(self.method, 'Float', vecs.shape[1])
# Set the thread number to speed up the build procedure in parallel
_index.SetBuildParam("NumberOfThreads", str(self.num_threads))
_index.SetBuildParam("DistCalcMethod", self.space)
if _index.Build(vecs.astype('float32'), vecs.shape[0]):
return _index
def get_query_handler(self):
vecs = super().get_query_handler()
if vecs is not None:
import SPTAG
_index = SPTAG.AnnIndex(self.method, 'Float', vecs.shape[1])
# Set the thread number to speed up the build procedure in parallel
_index.SetBuildParam("NumberOfThreads", str(self.num_threads))
_index.SetBuildParam("DistCalcMethod", self.method)
if _index.Build(vecs, vecs.shape[0]):
return _index
else:
return None
def exposed_add_data(self, index_name, p_data, p_meta, algo="BKT", dist="L2"):
print("*"*100)
print("add_data")
p_data = np.array(p_data, dtype=np.float32)
index_name = os.path.join(DATA_DIR, index_name)
if os.path.exists(index_name):
i = SPTAG.AnnIndex.Load(index_name)
else:
i = SPTAG.AnnIndex(algo, 'Float', p_data.shape[1])
i.SetBuildParam("NumberOfThreads", '4')
i.SetBuildParam("DistCalcMethod", dist)
p_num = p_data.shape[0]
success = i.AddWithMetaData(p_data, p_meta, p_num)
if success:
i.Save(index_name)
return success
def __init__(self, config=None):
if config is not None:
self.config = config
else:
self.config = configparser.ConfigParser()
self.config.read("config.ini")
self.index_data = self.config["SPTAG"]['Index']
self.algo = self.config["SPTAG"]['Algo']
self.dist_method = self.config["SPTAG"]['DistMethod']
self.data_type = self.config["SPTAG"]['DataType']
self.dimensions = int(self.config["SPTAG"]['Dimensions'])
self.threads = self.config["SPTAG"]['Threads']
self.threshold = float(self.config["SPTAG"]['Threshold'])
# if self.index_data is not None:
if os.path.exists(self.index_data):
self.index = SPTAG.AnnIndex.Load(self.index_data)
else:
# print(self.algo, self.data_type, self.dimensions)
self.index = SPTAG.AnnIndex(self.algo, self.data_type, self.dimensions)
self.index.SetBuildParam("NumberOfThreads", self.threads)
self.index.SetBuildParam("DistCalcMethod", self.dist_method)
self.build(np.ones((1, self.dimensions), dtype=np.float32), "0\n".encode())
def testBuildWithMetaData(algo, distmethod, x, s, out):
i = SPTAG.AnnIndex(algo, 'Float', x.shape[1])
i.SetBuildParam("NumberOfThreads", '4')
i.SetBuildParam("DistCalcMethod", distmethod)
if i.BuildWithMetaData(x, s, x.shape[0]):
i.Save(out)
def testBuild(algo, distmethod, x, out):
i = SPTAG.AnnIndex(algo, 'Float', x.shape[1])
i.SetBuildParam("NumberOfThreads", '4')
i.SetBuildParam("DistCalcMethod", distmethod)
ret = i.Build(x, x.shape[0])
i.Save(out)
def fit(self, X):
self._sptag = SPTAG.AnnIndex(self._algo, 'Float', X.shape[1])
self._sptag.SetBuildParam("NumberOfThreads", '32')
self._sptag.SetBuildParam("DistCalcMethod", self._metric)
self._sptag.Build(X, X.shape[0])