cleanlab · huiwengoh · May 14, 2024 · May 3, 2024 · May 3, 2024 · May 3, 2024
diff --git a/cleanlab/data_valuation.py b/cleanlab/data_valuation.py
@@ -19,14 +19,12 @@
 """
 
 
-from typing import Callable, Optional, Union, cast
+from typing import Callable, Optional, Union
 
 import numpy as np
 from scipy.sparse import csr_matrix
-from scipy.spatial.distance import euclidean
-from sklearn.neighbors import NearestNeighbors
-from sklearn.exceptions import NotFittedError
-from sklearn.utils.validation import check_is_fitted
+
+from cleanlab.internal.neighbor.knn_graph import create_knn_graph_and_index
 
 
 def _knn_shapley_score(knn_graph: csr_matrix, labels: np.ndarray, k: int) -> np.ndarray:
@@ -45,29 +43,6 @@ def _knn_shapley_score(knn_graph: csr_matrix, labels: np.ndarray, k: int) -> np.
     return 0.5 * (np.mean(scores / k, axis=0) + 1)
 
 
-def _process_knn_graph_from_features(
-    features: np.ndarray, metric: Optional[Union[str, Callable]], k: int = 10
-) -> csr_matrix:
-    """Calculate the knn graph from the features if it is not provided in the kwargs."""
-    if k > len(features):  # Ensure number of neighbors less than number of examples
-        raise ValueError(
-            f"Number of nearest neighbors k={k} cannot exceed the number of examples N={len(features)} passed into the estimator (knn)."
-        )
-    if metric == None:
-        metric = (
-            "cosine"
-            if features.shape[1] > 3
-            else "euclidean" if features.shape[0] > 100 else euclidean
-        )
-    knn = NearestNeighbors(n_neighbors=k, metric=metric).fit(features)
-    knn_graph = knn.kneighbors_graph(mode="distance")
-    try:
-        check_is_fitted(knn)
-    except NotFittedError:
-        knn.fit(features)
-    return knn_graph
-
-
 def data_shapley_knn(
     labels: np.ndarray,
     *,
@@ -135,6 +110,7 @@ def data_shapley_knn(
     if knn_graph is None and features is None:
         raise ValueError("Either knn_graph or features must be provided.")
 
+    # Use provided knn_graph or compute it from features
     if knn_graph is None:
-        knn_graph = _process_knn_graph_from_features(cast(np.ndarray, features), metric, k)
+        knn_graph, _ = create_knn_graph_and_index(features, n_neighbors=k, metric=metric)
     return _knn_shapley_score(knn_graph, labels, k)
diff --git a/cleanlab/datalab/internal/issue_manager/data_valuation.py b/cleanlab/datalab/internal/issue_manager/data_valuation.py
@@ -25,19 +25,15 @@
     Optional,
     Union,
 )
-import warnings
 
 
 import numpy as np
 import pandas as pd
 from scipy.sparse import csr_matrix
-from scipy.spatial.distance import euclidean
-from sklearn.exceptions import NotFittedError
-from sklearn.neighbors import NearestNeighbors
-from sklearn.utils.validation import check_is_fitted
 
 from cleanlab.data_valuation import data_shapley_knn
 from cleanlab.datalab.internal.issue_manager import IssueManager
+from cleanlab.internal.neighbor.knn_graph import create_knn_graph_and_index
 
 if TYPE_CHECKING:  # pragma: no cover
     import numpy.typing as npt
@@ -139,33 +135,11 @@ def find_issues(
             )
             raise TypeError(error_msg)
         if knn_graph is None or metric_changes:
-            if features is None:
-                raise ValueError(
-                    "If a knn_graph is not provided, features must be provided to fit a new knn."
-                )
-            if self.metric is None:
-                self.metric = (
-                    "cosine"
-                    if features.shape[1] > 3
-                    else "euclidean" if features.shape[0] > 100 else euclidean
-                )
-            knn = NearestNeighbors(n_neighbors=self.k, metric=self.metric).fit(features)
-
-            if self.metric and self.metric != knn.metric:
-                warnings.warn(
-                    f"Metric {self.metric} does not match metric {knn.metric} used to fit knn. "
-                    "Most likely an existing NearestNeighbors object was passed in, but a different "
-                    "metric was specified."
-                )
+            knn_graph, knn = create_knn_graph_and_index(
+                features, n_neighbors=self.k, metric=self.metric
+            )
             self.metric = knn.metric
 
-            try:
-                check_is_fitted(knn)
-            except NotFittedError:
-                knn.fit(features)
-
-            knn_graph = knn.kneighbors_graph(mode="distance")
-
         scores = data_shapley_knn(labels, knn_graph=knn_graph, k=self.k)
 
         self.issues = pd.DataFrame(

diff --git a/cleanlab/datalab/internal/issue_manager/duplicate.py b/cleanlab/datalab/internal/issue_manager/duplicate.py
@@ -21,12 +21,10 @@
 import numpy as np
 import pandas as pd
 from scipy.sparse import csr_matrix
-from scipy.spatial.distance import euclidean
-from sklearn.neighbors import NearestNeighbors
-from sklearn.exceptions import NotFittedError
-from sklearn.utils.validation import check_is_fitted
+
 
 from cleanlab.datalab.internal.issue_manager import IssueManager
+from cleanlab.internal.neighbor.knn_graph import create_knn_graph_and_index
 from cleanlab.internal.constants import EPSILON
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -76,32 +74,10 @@ def find_issues(
         metric_changes = self.metric and self.metric != old_knn_metric
 
         if knn_graph is None or metric_changes:
-            if features is None:
-                raise ValueError(
-                    "If a knn_graph is not provided, features must be provided to fit a new knn."
-                )
-            if self.metric is None:
-                self.metric = (
-                    "cosine"
-                    if features.shape[1] > 3
-                    else "euclidean" if features.shape[0] > 100 else euclidean
-                )
-            knn = NearestNeighbors(n_neighbors=self.k, metric=self.metric)
-
-            if self.metric and self.metric != knn.metric:
-                warnings.warn(
-                    f"Metric {self.metric} does not match metric {knn.metric} used to fit knn. "
-                    "Most likely an existing NearestNeighbors object was passed in, but a different "
-                    "metric was specified."
-                )
+            knn_graph, knn = create_knn_graph_and_index(
+                features, n_neighbors=self.k, metric=self.metric
+            )
             self.metric = knn.metric
-
-            try:
-                check_is_fitted(knn)
-            except NotFittedError:
-                knn.fit(features)
-
-            knn_graph = knn.kneighbors_graph(mode="distance")
         N = knn_graph.shape[0]
         nn_distances = knn_graph.data.reshape(N, -1)[:, 0]
         median_nn_distance = max(np.median(nn_distances), EPSILON)  # avoid threshold = 0

diff --git a/cleanlab/datalab/internal/issue_manager/noniid.py b/cleanlab/datalab/internal/issue_manager/noniid.py
@@ -1,19 +1,16 @@
 from __future__ import annotations
 
 from typing import TYPE_CHECKING, Any, Callable, ClassVar, Dict, Optional, Union, cast
-import warnings
 import itertools
 
 from scipy.stats import gaussian_kde
 import numpy as np
 import pandas as pd
 from scipy.sparse import csr_matrix
-from scipy.spatial.distance import euclidean
 from sklearn.neighbors import NearestNeighbors
-from sklearn.exceptions import NotFittedError
-from sklearn.utils.validation import check_is_fitted
 
 from cleanlab.datalab.internal.issue_manager import IssueManager
+from cleanlab.internal.neighbor.knn_graph import construct_knn_graph_from_index, features_to_knn
 
 if TYPE_CHECKING:  # pragma: no cover
     import numpy.typing as npt
@@ -203,28 +200,8 @@ def _setup_knn(
             return None
         features_to_use = self._determine_features(features, pred_probs)
 
-        if self.metric is None:
-            self.metric = (
-                "cosine"
-                if features_to_use.shape[1] > 3
-                else "euclidean" if features_to_use.shape[0] > 100 else euclidean
-            )
-
-        knn = NearestNeighbors(n_neighbors=self.k, metric=self.metric)
-
-        if self.metric != knn.metric:
-            warnings.warn(
-                f"Metric {self.metric} does not match metric {knn.metric} used to fit knn. "
-                "Most likely an existing NearestNeighbors object was passed in, but a different "
-                "metric was specified."
-            )
-        self.metric = knn.metric
-
-        try:
-            check_is_fitted(knn)
-        except NotFittedError:
-            knn.fit(features_to_use)
-
+        knn = features_to_knn(features_to_use, n_neighbors=self.k, metric=self.metric)
+        self.metric = knn.metric  # Update the metric to the one used in the KNN object.
         return knn
 
     def find_issues(
@@ -305,7 +282,7 @@ def collect_info(
         }
         if knn_graph is None:
             assert knn is not None, "If knn_graph is None, knn must be provided."
-            knn_graph = knn.kneighbors_graph(mode="distance")  # type: ignore[union-attr]
+            knn_graph = construct_knn_graph_from_index(knn)
 
         assert knn_graph is not None, "knn_graph must be provided or computed."
         statistics_dict = self._build_statistics_dictionary(knn_graph=knn_graph)

diff --git a/cleanlab/datalab/internal/issue_manager/outlier.py b/cleanlab/datalab/internal/issue_manager/outlier.py
@@ -23,6 +23,7 @@
 import pandas as pd
 
 from cleanlab.datalab.internal.issue_manager import IssueManager
+from cleanlab.internal.neighbor.knn_graph import construct_knn_graph_from_index
 from cleanlab.outlier import OutOfDistribution, transform_distances_to_scores
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -219,7 +220,7 @@ def _process_knn_graph_from_features(self, kwargs: Dict) -> csr_matrix:
             # If the pre-existing knn graph has fewer neighbors than the knn object,
             # then we need to recompute the knn graph
             assert knn == self.ood.params["knn"]  # type: ignore[union-attr]
-            knn_graph = knn.kneighbors_graph(mode="distance")  # type: ignore[union-attr]
+            knn_graph = construct_knn_graph_from_index(knn)
             self._metric = knn.metric  # type: ignore[union-attr]
 
         return knn_graph

diff --git a/cleanlab/datalab/internal/issue_manager/underperforming_group.py b/cleanlab/datalab/internal/issue_manager/underperforming_group.py
@@ -22,13 +22,10 @@
 import numpy as np
 import pandas as pd
 from scipy.sparse import csr_matrix
-from scipy.spatial.distance import euclidean
-from sklearn.neighbors import NearestNeighbors
-from sklearn.exceptions import NotFittedError
-from sklearn.utils.validation import check_is_fitted
 from sklearn.cluster import DBSCAN
 
 from cleanlab.datalab.internal.issue_manager import IssueManager
+from cleanlab.internal.neighbor.knn_graph import create_knn_graph_and_index
 from cleanlab.rank import get_self_confidence_for_each_label
 
 if TYPE_CHECKING:  # pragma: no cover
@@ -153,31 +150,10 @@ def set_knn_graph(
         metric_changes = self.metric and self.metric != old_knn_metric
 
         if knn_graph is None or metric_changes:
-            if features is None:
-                raise ValueError(
-                    "If a knn_graph is not provided, features must be provided to fit a new knn."
-                )
-            if self.metric is None:
-                self.metric = (
-                    "cosine"
-                    if features.shape[1] > 3
-                    else "euclidean" if features.shape[0] > 100 else euclidean
-                )
-            knn = NearestNeighbors(n_neighbors=self.k, metric=self.metric)
-
-            if self.metric and self.metric != knn.metric:
-                warnings.warn(
-                    f"Metric {self.metric} does not match metric {knn.metric} used to fit knn. "
-                    "Most likely an existing NearestNeighbors object was passed in, but a different "
-                    "metric was specified."
-                )
+            knn_graph, knn = create_knn_graph_and_index(
+                features, n_neighbors=self.k, metric=self.metric
+            )
             self.metric = knn.metric
-
-            try:
-                check_is_fitted(knn)
-            except NotFittedError:
-                knn.fit(features)
-            knn_graph = knn.kneighbors_graph(mode="distance")
         return knn_graph
 
     def perform_clustering(self, knn_graph: csr_matrix) -> npt.NDArray[np.int_]:

diff --git a/cleanlab/internal/neighbor/__init__.py b/cleanlab/internal/neighbor/__init__.py
@@ -0,0 +1 @@
+from .knn_graph import features_to_knn