Add SVS NearestNeighbors emulator

Author: Alexsandruss

configs/experiments/nearest_neighbors.json

@@ -35,6 +35,17 @@
                         "m_subvectors": 0.2
                     }
                 }
+            },
+            {
+                "algorithm": {
+                    "library": "sklbench.emulators.svs",
+                    "device": "cpu",
+                    "estimator_params": {
+                        "algorithm": "vamana",
+                        "graph_max_degree": 128,
+                        "window_size": 256
+                    }
+                }
             }
         ],
         "nearest neighbors common parameters": {

sklbench/emulators/__init__.py

@@ -14,4 +14,4 @@
 # limitations under the License.
 # ===============================================================================
 
-__all__ = ["common", "faiss", "raft"]
+__all__ = ["common", "faiss", "raft", "svs"]

sklbench/emulators/common/neighbors.py

@@ -27,12 +27,15 @@ def get_params(self):
             "metric": self.metric,
             "metric_params": None,
             "p": 2 if "euclidean" in self.metric else None,
-            "n_lists": self.n_lists,
-            "n_probes": self.n_probes,
-            "m_subvectors": self.m_subvectors,
-            "n_bits": self.n_bits,
         }
-        optional_keys = ["intermediate_graph_degree", "graph_degree"]
+        optional_keys = [
+            "n_lists",
+            "n_probes",
+            "m_subvectors",
+            "n_bits",
+            "intermediate_graph_degree",
+            "graph_degree",
+        ]
         for optional_key in optional_keys:
             if hasattr(self, optional_key):
                 result[optional_key] = getattr(self, optional_key)

sklbench/emulators/svs/__init__.py

@@ -0,0 +1,20 @@
+# ===============================================================================
+# Copyright 2024 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ===============================================================================
+
+from .neighbors import NearestNeighbors
+
+
+__all__ = ["NearestNeighbors"]

sklbench/emulators/svs/neighbors.py

@@ -0,0 +1,63 @@
+# ===============================================================================
+# Copyright 2024 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ===============================================================================
+
+import pysvs
+from psutil import cpu_count
+from ..common.neighbors import NearestNeighborsBase
+
+
+class NearestNeighbors(NearestNeighborsBase):
+    """
+    Minimal class emulating `sklearn.neighbors.NearestNeighbors` estimator
+    """
+
+    def __init__(
+        self,
+        n_neighbors=5,
+        algorithm="vamana",
+        metric="euclidean",
+        graph_max_degree=64,
+        window_size=128,
+        n_jobs=cpu_count(logical=False),
+    ):
+        self.n_neighbors = n_neighbors
+        self.algorithm = algorithm
+        self.metric = metric
+        self.graph_max_degree = graph_max_degree
+        self.window_size = window_size
+        self.n_jobs = n_jobs
+
+    def fit(self, X, y=None):
+        build_params = pysvs.VamanaBuildParameters(
+            graph_max_degree=self.graph_max_degree,
+            window_size=self.window_size,
+            num_threads=self.n_jobs,
+        )
+        self._index = pysvs.Vamana.build(
+            build_params,
+            X,
+            pysvs.DistanceType.L2,
+            num_threads=self.n_jobs,
+        )
+        return self
+
+    def kneighbors(self, X, n_neighbors=None, return_distance=True):
+        k = self.n_neighbors if n_neighbors is None else n_neighbors
+        indices, distances = self._index.search(X, k)
+        if return_distance:
+            return distances, indices
+        else:
+            return indices