Repository: secretflow/spu (language: python)

Issue:
Bump release tag to 0.1.0b1



Discussion / hints:

Fix patch:
diff --git a/sml/neighbors/BUILD.bazel b/sml/neighbors/BUILD.bazel
index 7832e732..ff98f7eb 100644
--- a/sml/neighbors/BUILD.bazel
+++ b/sml/neighbors/BUILD.bazel
@@ -11,3 +11,12 @@
 # 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.
+
+load("@rules_python//python:defs.bzl", "py_library")
+
+package(default_visibility = ["//visibility:public"])
+
+py_library(
+    name = "knn",
+    srcs = ["knn.py"],
+)
diff --git a/sml/neighbors/emulations/BUILD.bazel b/sml/neighbors/emulations/BUILD.bazel
index 7832e732..22d3b0c0 100644
--- a/sml/neighbors/emulations/BUILD.bazel
+++ b/sml/neighbors/emulations/BUILD.bazel
@@ -11,3 +11,16 @@
 # 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.
+
+load("@rules_python//python:defs.bzl", "py_binary")
+
+package(default_visibility = ["//visibility:public"])
+
+py_binary(
+    name = "knn_emul",
+    srcs = ["knn_emul.py"],
+    deps = [
+        "//sml/neighbors:knn",
+        "//sml/utils:emulation",
+    ],
+)
diff --git a/sml/neighbors/emulations/knn_emul.py b/sml/neighbors/emulations/knn_emul.py
new file mode 100644
index 00000000..c9a225b9
--- /dev/null
+++ b/sml/neighbors/emulations/knn_emul.py
@@ -0,0 +1,113 @@
+# Copyright 2023 Ant Group Co., Ltd.
+#
+# 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
+#
+#      https://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 os
+import sys
+from collections import defaultdict
+
+import jax.numpy as jnp
+import numpy as np
+from sklearn.neighbors import KNeighborsClassifier
+
+# from sklearn.metrics import roc_auc_score, explained_variance_score
+
+# Add the library directory to the path
+sys.path.append(os.path.join(os.path.dirname(__file__), '../../../'))
+
+import sml.utils.emulation as emulation
+from sml.neighbors.knn import KNNClassifer
+
+
+# TODO: design the enumation framework, just like py.unittest
+# all emulation action should begin with `emul_` (for reflection)
+def emul_KNN(mode: emulation.Mode.MULTIPROCESS):
+    def proc_KNN(X_train, y_train, X_test, n_classes, n_neighbors=5, weights='uniform'):
+        knn_model = KNNClassifer(
+            n_neighbors=n_neighbors, weights=weights, n_classes=n_classes
+        )
+        knn_model.fit(X_train, y_train)
+
+        return knn_model.predict(X_test)
+
+    try:
+        # bandwidth and latency only work for docker mode
+        emulator = emulation.Emulator(
+            emulation.CLUSTER_ABY3_3PC, mode, bandwidth=300, latency=20
+        )
+        emulator.up()
+        # 假设有一组样本和对应的标签
+        X_train = jnp.array([[1, 2], [2, 3], [3, 4], [5, 1]])
+        y_train = jnp.array([0, 3, 3, 0])
+
+        # 假设有一组测试样本
+        X_test = jnp.array([[4, 2], [1, 3]])
+
+        # 获取样本的类别数
+        n_classes = len(set(y_train.tolist()))
+        n_neighbors = 3
+
+        # 将y_train映射为从0开始的连续数组
+        label_to_int = defaultdict(lambda: len(label_to_int))
+        y_train_new = jnp.array([label_to_int[label] for label in y_train.tolist()])
+
+        (
+            X_train_,
+            y_train_new_,
+            X_test_,
+        ) = emulator.seal(X_train, y_train_new, X_test)
+
+        result_uniform = emulator.run(proc_KNN, static_argnums=(3, 4, 5))(
+            X_train_, y_train_new_, X_test_, n_classes, n_neighbors, 'uniform'
+        )
+
+        result_distance = emulator.run(proc_KNN, static_argnums=(3, 4, 5))(
+            X_train_, y_train_new_, X_test_, n_classes, n_neighbors, 'distance'
+        )
+
+        # 再从连续数组映射回原来的标签
+        int_to_label = {i: label for label, i in label_to_int.items()}
+        result_uniform_np = np.array(result_uniform)
+        result_distance_np = np.array(result_distance)
+        predictions_uniform = [
+            int_to_label[prediction] for prediction in result_uniform_np
+        ]
+        predictions_distance = [
+            int_to_label[prediction] for prediction in result_distance_np
+        ]
+
+        # 与sklearn的结果进行比较
+        X_train = np.array(X_train)
+        y_train = np.array(y_train)
+        X_test = np.array(X_test)
+
+        neigh_uni = KNeighborsClassifier(n_neighbors=3, weights='uniform')
+        neigh_uni.fit(X_train, y_train)
+
+        sklearn_predictions = neigh_uni.predict(X_test)
+
+        assert np.array_equal(predictions_uniform, sklearn_predictions)
+
+        neigh_dis = KNeighborsClassifier(n_neighbors=3, weights='distance')
+        neigh_dis.fit(X_train, y_train)
+
+        sklearn_predictions = neigh_dis.predict(X_test)
+
+        assert np.array_equal(predictions_distance, sklearn_predictions)
+
+    finally:
+        emulator.down()
+
+
+if __name__ == "__main__":
+    emul_KNN(emulation.Mode.MULTIPROCESS)
diff --git a/sml/neighbors/knn.py b/sml/neighbors/knn.py
new file mode 100644
index 00000000..9dfd1cbc
--- /dev/null
+++ b/sml/neighbors/knn.py
@@ -0,0 +1,158 @@
+# Copyright 2023 Ant Group Co., Ltd.
+#
+# 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
+#
+#      https://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 jax.numpy as jnp
+from jax import lax, vmap
+
+
+class KNNClassifer:
+    def __init__(
+        self,
+        n_neighbors=5,
+        weights='uniform',
+        metric=None,
+        metric_params=None,
+        n_classes=None,
+    ):
+        """
+        K-Nearest Neighbors (KNN) Classifier.
+
+        Parameters:
+        - n_neighbors (int): Number of neighbors to use for prediction.
+        - weights (str): Weight function used in prediction. Possible values:
+            * 'uniform': Uniform weights. All points in each neighborhood are weighted equally.
+            * 'distance': Weight points by the inverse of their distance.
+        - metric (callable, optional): A function that computes distance between two points.
+            It should take two arguments and return a scalar. Defaults to the Minkowski metric
+            with p=2 (i.e., Euclidean distance). If None, the Minkowski metric with p=2 is used.
+        - metric_params (dict, optional): Additional keyword arguments for the metric function.
+        - n_classes (int, optional): Number of classes in the dataset. It needs to be provided
+            for computation.
+
+        Note:
+        If the default metric (Minkowski with p=2) is used and metric_params is None,
+        it will default to {'ord': 2}.
+        """
+
+        # Validate n_neighbors
+        if not isinstance(n_neighbors, int) or n_neighbors <= 0:
+            raise ValueError("n_neighbors should be a positive integer.")
+
+        # Validate weights
+        if weights not in ['uniform', 'distance']:
+            raise ValueError("weights should be either 'uniform' or 'distance'.")
+
+        # Validate n_classes
+        if (n_classes is None) or (not isinstance(n_classes, int)) or (n_classes <= 0):
+            raise ValueError("n_classes should be a positive integer if provided.")
+
+        self.X = None
+        self.y = None
+        self.n_neighbors = n_neighbors
+        self.weights = weights
+        self.n_classes = n_classes
+
+        # If metric is None, default to Minkowski with p=2
+        if metric is None:
+            self.metric = jnp.linalg.norm
+            if metric_params is None:
+                metric_params = {'ord': 2}
+        else:
+            # Check if the metric is callable
+            if not callable(metric):
+                raise ValueError("The metric should be a callable function.")
+            self.metric = metric
+
+        self.metric_params = metric_params
+
+    def fit(self, X, y):
+        """
+        Fit the KNN classifier.
+
+        Parameters:
+        - X (array-like): Training data.
+        - y (array-like): Target values.
+
+        Note:
+        The y here needs to be converted to an integer starting from 0 before passing in, otherwise may occur unexpected errors.
+        """
+        self.X = X
+        self.y = jnp.array(y)
+
+    def _compute_weights(self, distances):
+        """
+        Compute the weights based on the distances and the weight scheme.
+
+        Parameters:
+        - distances (array-like): Distances to the neighbors.
+
+        Returns:
+        - array-like: Weights for each neighbor.
+        """
+        if self.weights == 'uniform':
+            return jnp.ones_like(distances)
+        elif self.weights == 'distance':
+            return 1.0 / distances
+        else:
+            raise ValueError(
+                "Invalid weight setting. Only 'uniform' and 'distance' are supported."
+            )
+
+    def _predict_single_sample(self, x):
+        """
+        Predict the class label for a single sample.
+
+        Parameters:
+        - x (array-like): A single sample.
+
+        Returns:
+        - int: Predicted class label.
+        """
+        distances = self.metric(self.X - x, **self.metric_params, axis=1)
+        sorted_indices = jnp.argsort(distances)
+        k_indices = sorted_indices[: self.n_neighbors]
+        k_distances = distances[k_indices]
+        k_labels = self.y[k_indices]
+
+        def predict_uniform(_):
+            counts = jnp.bincount(k_labels, length=int(self.n_classes))
+            prediction = jnp.argmax(counts)
+            return prediction
+
+        def predict_distance(_):
+            weights_inner = self._compute_weights(k_distances)
+            weighted_counts = jnp.bincount(
+                k_labels, weights=weights_inner, length=int(self.n_classes)
+            )
+            prediction = jnp.argmax(weighted_counts)
+            return prediction
+
+        return lax.cond(
+            self.weights == 'uniform', predict_uniform, predict_distance, operand=None
+        )
+
+    def predict(self, X):
+        """
+        Predict the class labels for the provided data.
+
+        Parameters:
+        - X (array-like): Samples to predict.
+
+        Returns:
+        - array-like: Predicted class labels for each sample.
+        """
+        # Use the vectorized function to predict all samples at once
+        predictions = vmap(self._predict_single_sample)(X)
+
+        return predictions


Test patch:
diff --git a/sml/neighbors/tests/BUILD.bazel b/sml/neighbors/tests/BUILD.bazel
index 7832e732..5c705279 100644
--- a/sml/neighbors/tests/BUILD.bazel
+++ b/sml/neighbors/tests/BUILD.bazel
@@ -11,3 +11,17 @@
 # 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.
+
+load("@rules_python//python:defs.bzl", "py_test")
+
+package(default_visibility = ["//visibility:public"])
+
+py_test(
+    name = "knn_test",
+    srcs = ["knn_test.py"],
+    deps = [
+        "//sml/neighbors:knn",
+        "//spu:init",
+        "//spu/utils:simulation",
+    ],
+)
diff --git a/sml/neighbors/tests/knn_test.py b/sml/neighbors/tests/knn_test.py
new file mode 100644
index 00000000..59ff6dd9
--- /dev/null
+++ b/sml/neighbors/tests/knn_test.py
@@ -0,0 +1,103 @@
+# Copyright 2023 Ant Group Co., Ltd.
+#
+# 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
+#
+#      https://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 os
+import sys
+import unittest
+from collections import defaultdict
+
+import jax.numpy as jnp
+import numpy as np
+from sklearn.neighbors import KNeighborsClassifier
+
+import spu.spu_pb2 as spu_pb2
+import spu.utils.simulation as spsim
+
+# Add the sml directory to the path
+sys.path.append(os.path.join(os.path.dirname(__file__), '../../../'))
+
+from sml.neighbors.knn import KNNClassifer
+
+
+class UnitTests(unittest.TestCase):
+    def test_knn(self):
+        sim = spsim.Simulator.simple(
+            3, spu_pb2.ProtocolKind.ABY3, spu_pb2.FieldType.FM64
+        )
+
+        # Test fit_predict
+        def proc_predict(
+            X_train, y_train, X_test, n_classes, n_neighbors=5, weights='uniform'
+        ):
+            knn_model = KNNClassifer(
+                n_neighbors=n_neighbors, weights=weights, n_classes=n_classes
+            )
+            knn_model.fit(X_train, y_train)
+
+            return knn_model.predict(X_test)
+
+        # 假设有一组样本和对应的标签
+        X_train = jnp.array([[1, 2], [2, 3], [3, 4], [5, 1]])
+        y_train = jnp.array([-1, 3, 3, -1])
+
+        # 假设有一组测试样本
+        X_test = jnp.array([[4, 2], [1, 3]])
+
+        # 获取样本的类别数
+        n_classes = len(set(y_train.tolist()))
+
+        # 将y_train映射为从0开始的连续数组
+        label_to_int = defaultdict(lambda: len(label_to_int))
+        y_train_new = jnp.array([label_to_int[label] for label in y_train.tolist()])
+
+        # 运行模拟器
+        result_uniform = spsim.sim_jax(sim, proc_predict, static_argnums=(3, 4, 5))(
+            X_train, y_train_new, X_test, n_classes, 3, 'uniform'
+        )
+        result_distance = spsim.sim_jax(sim, proc_predict, static_argnums=(3, 4, 5))(
+            X_train, y_train_new, X_test, n_classes, 3, 'distance'
+        )
+
+        # 再从连续数组映射回原来的标签
+        int_to_label = {i: label for label, i in label_to_int.items()}
+        result_uniform_np = np.array(result_uniform)
+        result_distance_np = np.array(result_distance)
+        predictions_uniform = [
+            int_to_label[prediction] for prediction in result_uniform_np
+        ]
+        predictions_distance = [
+            int_to_label[prediction] for prediction in result_distance_np
+        ]
+
+        # 与sklearn的结果进行比较
+        X_train = np.array(X_train)
+        y_train = np.array(y_train)
+        X_test = np.array(X_test)
+
+        neigh_uni = KNeighborsClassifier(n_neighbors=3, weights='uniform')
+        neigh_uni.fit(X_train, y_train)
+
+        sklearn_predictions = neigh_uni.predict(X_test)
+
+        self.assertEqual(predictions_uniform, sklearn_predictions.tolist())
+
+        neigh_dis = KNeighborsClassifier(n_neighbors=3, weights='distance')
+        neigh_dis.fit(X_train, y_train)
+
+        sklearn_predictions = neigh_dis.predict(X_test)
+
+        self.assertEqual(predictions_distance, sklearn_predictions.tolist())
+
+
+if __name__ == "__main__":
+    unittest.main()