Use cuvs for knn when available (#103)

Author: praateekmahajan

crossfit/backend/__init__.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 # flake8: noqa
+import logging
 
 from crossfit.backend.dask.dataframe import *
 from crossfit.backend.numpy.sparse import *
@@ -23,17 +24,20 @@
     from crossfit.backend.cudf.array import *
     from crossfit.backend.cudf.dataframe import *
 except ImportError:
+    logging.warning("Import Error for cudf backend in Crossfit. Skipping it.")
     pass
 
 try:
     from crossfit.backend.cupy.array import *
     from crossfit.backend.cupy.sparse import *
 except ImportError:
+    logging.warning("Import Error for cupy backend in Crossfit. Skipping it.")
     pass
 
 try:
     from crossfit.backend.torch.array import *
 except ImportError:
+    logging.warning("Import Error for Torch backend in Crossfit. Skipping it.")
     pass
 
 # from crossfit.backend.tf.array import *

crossfit/op/vector_search.py

@@ -16,17 +16,26 @@
 
 import cudf
 import cupy as cp
+import cuvs
 import dask.dataframe as dd
+import pylibraft
 from cuml.dask.neighbors import NearestNeighbors
 from dask import delayed
 from dask_cudf import from_delayed
-from pylibraft.neighbors.brute_force import knn
+from packaging.version import parse as parse_version
 
 from crossfit.backend.cudf.series import create_list_series_from_1d_or_2d_ar
 from crossfit.backend.dask.cluster import global_dask_client
 from crossfit.dataset.base import EmbeddingDatataset
 from crossfit.op.base import Op
 
+if (parse_version(pylibraft.__version__).base_version >= "24.12") or (
+    parse_version(cuvs.__version__).base_version >= "24.12"
+):
+    from cuvs.neighbors.brute_force import search
+else:
+    from pylibraft.neighbors.brute_force import knn as search
+
 
 class VectorSearchOp(Op):
     @overload
@@ -171,7 +180,7 @@ def __init__(
         self.normalize = normalize
 
     def search_tensors(self, queries, corpus):
-        results, indices = knn(dataset=corpus, queries=queries, k=self.k, metric=self.metric)
+        results, indices = search(dataset=corpus, queries=queries, k=self.k, metric=self.metric)
 
         return cp.asarray(results), cp.asarray(indices)