Upgrade to use newer NMSLIB api
This package was using an old version of the nmslib api. While I kept that version around
for backwards compatibility, the newer version has a couple of advantages: you don't
need to call freeIndex, the newer version accepts numpy arrays as inputs etc etc.
@@ -4,10 +4,12 @@ import nmslib
from ann_benchmarks.constants import INDEX_DIR
from ann_benchmarks.algorithms.base import BaseANN
+
class NmslibReuseIndex(BaseANN):
@staticmethod
def encode(d):
return ["%s=%s" % (a, b) for (a, b) in d.iteritems()]
+
def __init__(self, metric, method_name, index_param, save_index, query_param):
self._nmslib_metric = {'angular': 'cosinesimil', 'euclidean': 'l2'}[metric]
self._method_name = method_name
@@ -29,27 +31,23 @@ class NmslibReuseIndex(BaseANN):
# Aborted (core dumped)
self._index_param.append('bucketSize=%d' % min(int(X.shape[0] * 0.0005), 1000))
- self._index = nmslib.init(space=self._nmslib_metric,
- method=self._method_name)
-
- for i, x in enumerate(X):
- nmslib.addDataPoint(self._index, i, x.tolist())
+ self._index = nmslib.init(space=self._nmslib_metric, method=self._method_name)
+ self._index.addDataPointBatch(X)
if os.path.exists(self._index_name):
print('Loading index from file')
- nmslib.loadIndex(self._index, self._index_name)
+ self._index.loadIndex(self._index_name)
else:
- nmslib.createIndex(self._index, self._index_param)
+ self._index.createIndex(self._index_param)
if self._save_index:
- nmslib.saveIndex(self._index, self._index_name)
+ self._index.saveIndex(self._index_name)
- nmslib.setQueryTimeParams(self._index, self._query_param)
+ self._index.setQueryTimeParams(self._query_param)
def query(self, v, n):
- return nmslib.knnQuery(self._index, n, v.tolist())
+ ids, distances = self._index.knnQuery(v, n)
+ return ids
- def freeIndex(self):
- nmslib.freeIndex(self._index)
class NmslibNewIndex(BaseANN):
def __init__(self, metric, method_name, method_param):
@@ -66,15 +64,11 @@ class NmslibNewIndex(BaseANN):
# Aborted (core dumped)
self._method_param.append('bucketSize=%d' % min(int(X.shape[0] * 0.0005), 1000))
- self._index = nmslib.init(self._nmslib_metric, [], self._method_name, nmslib.DataType.DENSE_VECTOR, nmslib.DistType.FLOAT)
-
- for i, x in enumerate(X):
- nmslib.addDataPoint(self._index, i, x.tolist())
+ self._index = nmslib.init(space=self._nmslib_metric, method=self._method_name)
+ self._index.addDataPointBatch(X)
nmslib.createIndex(self._index, self._method_param)
def query(self, v, n):
- return nmslib.knnQuery(self._index, n, v.tolist())
-
- def freeIndex(self):
- nmslib.freeIndex(self._index)
+ ids, distances = self._index.knnQuery(v, n)
+ return ids
