Add jaccard to PUFFINN
@@ -5,6 +5,9 @@ import numpy
class Puffinn(BaseANN):
def __init__(self, metric, space=10**6, hash_function="fht_crosspolytope", hash_source='pool', hash_args=None):
+ if metric not in ['jaccard', 'angular']:
+ raise NotImplementedError(
+ "Puffinn doesn't support metric %s" % metric)
self.metric = metric
self.space = space
self.hash_function = hash_function
@@ -13,21 +16,25 @@ class Puffinn(BaseANN):
def fit(self, X):
if self.hash_args:
- self.index = puffinn.Index(self.metric, X.shape[1], self.space,\
+ self.index = puffinn.Index(self.metric, len(X[0]), self.space,\
hash_function=self.hash_function, hash_source=self.hash_source,\
hash_args=self.hash_args)
else:
- self.index = puffinn.Index(self.metric, X.shape[1], self.space,\
+ self.index = puffinn.Index(self.metric, len(X[0]), self.space,\
hash_function=self.hash_function, hash_source=self.hash_source)
for i, x in enumerate(X):
- self.index.insert(x.tolist())
+ if self.metric == 'angular':
+ x = x.tolist()
+ self.index.insert(x)
self.index.rebuild(10)
def set_query_arguments(self, recall):
self.recall = recall
def query(self, v, n):
- return self.index.search(v.tolist(), n, self.recall)
+ if self.metric == 'angular':
+ v = v.tolist()
+ return self.index.search(v, n, self.recall)
def __str__(self):
return 'PUFFINN(space=%d, recall=%f, hf=%s, hashsource=%s)' % (self.space, self.recall, self.hash_function, self.hash_source)
