def add(self, other, positions=None, last=None):
other = ensure_gpu(other, self.gpu_id)
T.cuda.synchronize()
if positions is not None:
positions = ensure_gpu(positions, self.gpu_id)
assert positions.size(0) == other.size(0), "Mismatch in number of positions and vectors"
self.index.add_with_ids_c(other.size(0), cast_float(ptr(other)), cast_long(ptr(positions + 1)))
else:
other = other[:last, :] if last is not None else other
self.index.add_c(other.size(0), cast_float(ptr(other)))
T.cuda.synchronize()
def search(self, query, k=None):
query = ensure_gpu(query, self.gpu_id)
k = k if k else self.K
(b, n) = query.size()
distances = T.FloatTensor(b, k)
labels = T.LongTensor(b, k)
if self.gpu_id != -1: distances = distances.cuda(self.gpu_id)
if self.gpu_id != -1: labels = labels.cuda(self.gpu_id)
T.cuda.synchronize()
self.index.search_c(b, cast_float(ptr(query)), k,
cast_float(ptr(distances)), cast_long(ptr(labels)))
T.cuda.synchronize()
return (distances, (labels - 1))