Skip to content

Commit

Permalink
FIX correct call to brute force in generate groundtruth
Browse files Browse the repository at this point in the history
  • Loading branch information
dantegd committed Oct 25, 2024
1 parent e7f1085 commit 791d8e2
Showing 1 changed file with 4 additions and 3 deletions.
7 changes: 4 additions & 3 deletions python/cuvs_bench/cuvs_bench/generate_groundtruth/__main__.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@
from pylibraft.common import DeviceResources
from rmm.allocators.cupy import rmm_cupy_allocator

from cuvs.neighbors.brute_force import knn
from cuvs.neighbors.brute_force import build, search

from .utils import memmap_bin_file, suffix_from_dtype, write_bin

Expand All @@ -49,7 +49,7 @@ def choose_random_queries(dataset, n_queries):


def calc_truth(dataset, queries, k, metric="sqeuclidean"):
handle = DeviceResources()
resources = DeviceResources()
n_samples = dataset.shape[0]
n = 500000 # batch size for processing neighbors
i = 0
Expand All @@ -63,7 +63,8 @@ def calc_truth(dataset, queries, k, metric="sqeuclidean"):

X = cp.asarray(dataset[i : i + n_batch, :], cp.float32)

D, Ind = knn(X, queries, k, metric=metric, handle=handle)
index = build(X, metric=metric, resources=resources)
D, Ind = search(index, queries, k, resources=resources)
handle.sync()

D, Ind = cp.asarray(D), cp.asarray(Ind)
Expand Down

0 comments on commit 791d8e2

Please sign in to comment.