Repository: cleanlab/cleanlab

Issue:
allow Data Valuation Issue Manager to operate on provided features
It currently only operates on `knn_graph`. Should simply construct the KNN graph based on the features if one does not already exist in the Datalab object. This is what the Outlier Issue Manager does.


Discussion:

Resolving patch:
diff --git a/cleanlab/datalab/internal/issue_finder.py b/cleanlab/datalab/internal/issue_finder.py
--- a/cleanlab/datalab/internal/issue_finder.py
+++ b/cleanlab/datalab/internal/issue_finder.py
@@ -59,7 +59,7 @@
     "near_duplicate": ["features", "knn_graph"],
     "non_iid": ["pred_probs", "features", "knn_graph"],
     "underperforming_group": ["pred_probs", "features", "knn_graph", "cluster_ids"],
-    "data_valuation": ["knn_graph"],
+    "data_valuation": ["features", "knn_graph"],
     "class_imbalance": [],
     "null": ["features"],
 }
@@ -68,6 +68,7 @@
     "outlier": ["features", "knn_graph"],
     "near_duplicate": ["features", "knn_graph"],
     "non_iid": ["features", "knn_graph"],
+    "data_valuation": ["features", "knn_graph"],
     "null": ["features"],
 }
 
@@ -76,6 +77,7 @@
     "outlier": ["features", "knn_graph"],
     "near_duplicate": ["features", "knn_graph"],
     "non_iid": ["features", "knn_graph"],
+    "data_valuation": ["features", "knn_graph"],
     "null": ["features"],
 }
 
diff --git a/cleanlab/datalab/internal/issue_manager/data_valuation.py b/cleanlab/datalab/internal/issue_manager/data_valuation.py
--- a/cleanlab/datalab/internal/issue_manager/data_valuation.py
+++ b/cleanlab/datalab/internal/issue_manager/data_valuation.py
@@ -24,14 +24,20 @@
     Optional,
     Union,
 )
+import warnings
+
 
 import numpy as np
 import pandas as pd
 from scipy.sparse import csr_matrix
+from sklearn.exceptions import NotFittedError
+from sklearn.neighbors import NearestNeighbors
+from sklearn.utils.validation import check_is_fitted
 
 from cleanlab.datalab.internal.issue_manager import IssueManager
 
 if TYPE_CHECKING:  # pragma: no cover
+    import numpy.typing as npt
     import pandas as pd
     from cleanlab.datalab.datalab import Datalab
 
@@ -94,16 +100,19 @@ class DataValuationIssueManager(IssueManager):
     def __init__(
         self,
         datalab: Datalab,
+        metric: Optional[str] = None,
         threshold: Optional[float] = None,
         k: int = 10,
         **kwargs,
     ):
         super().__init__(datalab)
+        self.metric = metric
         self.k = k
         self.threshold = threshold if threshold is not None else self.DEFAULT_THRESHOLD
 
     def find_issues(
         self,
+        features: Optional[npt.NDArray] = None,
         **kwargs,
     ) -> None:
         """Calculate the data valuation score with a provided or existing knn graph.
@@ -117,6 +126,8 @@ def find_issues(
         """
         self.k = kwargs.get("k", self.k)
         knn_graph = self._process_knn_graph_from_inputs(kwargs)
+        old_knn_metric = self.datalab.get_info("statistics").get("knn_metric")
+        metric_changes = self.metric and self.metric != old_knn_metric
         labels = self.datalab.labels
         if not isinstance(labels, np.ndarray):
             error_msg = (
@@ -124,12 +135,29 @@ def find_issues(
                 f"but got {type(labels)} instead."
             )
             raise TypeError(error_msg)
-        if knn_graph is None:
-            raise ValueError(
-                "knn_graph must be provided in kwargs or already stored in the Datalab instance\n"
-                "It should be calculated by other issue managers if it is not provided via "
-                "`Datalab.find_issues(knn_graph=knn_graph, ...)`"
-            )
+        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"
+            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."
+                )
+            self.metric = knn.metric
+
+            try:
+                check_is_fitted(knn)
+            except NotFittedError:
+                knn.fit(features)
+
+            knn_graph = knn.kneighbors_graph(mode="distance")
         if labels is None:
             raise ValueError("labels must be provided to run data valuation")
 
@@ -143,7 +171,7 @@ def find_issues(
         )
         self.summary = self.make_summary(score=scores.mean())
 
-        self.info = self.collect_info(self.issues)
+        self.info = self.collect_info(issues=self.issues, knn_graph=knn_graph)
 
     def _process_knn_graph_from_inputs(self, kwargs: Dict[str, Any]) -> Union[csr_matrix, None]:
         """Determine if a knn_graph is provided in the kwargs or if one is already stored in the associated Datalab instance."""
@@ -163,18 +191,47 @@ def _process_knn_graph_from_inputs(self, kwargs: Dict[str, Any]) -> Union[csr_ma
             )
         return knn_graph
 
-    def collect_info(self, issues: pd.DataFrame) -> dict:
+    def collect_info(self, issues: pd.DataFrame, knn_graph: csr_matrix) -> dict:
         issues_info = {
             "num_low_valuation_issues": sum(issues[f"is_{self.issue_name}_issue"]),
             "average_data_valuation": issues[self.issue_score_key].mean(),
         }
 
+        params_dict = {
+            "metric": self.metric,
+            "k": self.k,
+            "threshold": self.threshold,
+        }
+
+        statistics_dict = self._build_statistics_dictionary(knn_graph=knn_graph)
+
         info_dict = {
             **issues_info,
+            **params_dict,
+            **statistics_dict,
         }
 
         return info_dict
 
+    def _build_statistics_dictionary(self, knn_graph: csr_matrix) -> Dict[str, Dict[str, Any]]:
+        statistics_dict: Dict[str, Dict[str, Any]] = {"statistics": {}}
+
+        # Add the knn graph as a statistic if necessary
+        graph_key = "weighted_knn_graph"
+        old_knn_graph = self.datalab.get_info("statistics").get(graph_key, None)
+        old_graph_exists = old_knn_graph is not None
+        prefer_new_graph = (
+            not old_graph_exists
+            or knn_graph.nnz > old_knn_graph.nnz
+            or self.metric != self.datalab.get_info("statistics").get("knn_metric", None)
+        )
+        if prefer_new_graph:
+            statistics_dict["statistics"][graph_key] = knn_graph
+            if self.metric is not None:
+                statistics_dict["statistics"]["knn_metric"] = self.metric
+
+        return statistics_dict
+
 
 def _knn_shapley_score(knn_graph: csr_matrix, labels: np.ndarray, k: int) -> np.ndarray:
     """Compute the Shapley values of data points based on a knn graph."""
diff --git a/cleanlab/datalab/internal/issue_manager_factory.py b/cleanlab/datalab/internal/issue_manager_factory.py
--- a/cleanlab/datalab/internal/issue_manager_factory.py
+++ b/cleanlab/datalab/internal/issue_manager_factory.py
@@ -74,6 +74,7 @@
         "outlier": OutlierIssueManager,
         "near_duplicate": NearDuplicateIssueManager,
         "non_iid": NonIIDIssueManager,
+        "data_valuation": DataValuationIssueManager,
         "null": NullIssueManager,
     },
     Task.MULTILABEL: {
@@ -81,6 +82,7 @@
         "outlier": OutlierIssueManager,
         "near_duplicate": NearDuplicateIssueManager,
         "non_iid": NonIIDIssueManager,
+        "data_valuation": DataValuationIssueManager,
         "null": NullIssueManager,
     },
 }


Test patch:
diff --git a/tests/datalab/datalab/test_datalab.py b/tests/datalab/datalab/test_datalab.py
--- a/tests/datalab/datalab/test_datalab.py
+++ b/tests/datalab/datalab/test_datalab.py
@@ -1495,6 +1495,122 @@ def test_report(self, embeddings_with_null):
         ), "Report should not contain a tip to address partial null examples"
 
 
+class TestDatalabDataValuation:
+    label_name: str = "y"
+
+    @pytest.fixture
+    def dataset(self):
+        from sklearn.datasets import make_classification
+
+        np.random.seed(SEED)
+
+        # Generate a 10D dataset with 2 classes
+        X, y = make_classification(
+            n_samples=100,
+            n_features=10,
+            n_informative=2,
+            n_redundant=2,
+            n_repeated=0,
+            n_classes=2,
+            n_clusters_per_class=2,
+            weights=None,
+            flip_y=0.1,
+            class_sep=1.0,
+            hypercube=True,
+            shift=0.0,
+            scale=0.1,
+            shuffle=True,
+            random_state=SEED,
+        )
+
+        return {"X": X, self.label_name: y}
+
+    @pytest.fixture
+    def knn_graph(self, dataset):
+        return (
+            NearestNeighbors(n_neighbors=10, metric="cosine")
+            .fit(dataset["X"])
+            .kneighbors_graph(mode="distance")
+        )
+
+    def test_find_issues(self, dataset, knn_graph):
+        """Test that a fresh Datalab instance can check for data_valuation issues with
+        either `features` or a `knn_graph`.
+        """
+
+        datalabs, summaries, scores_list = [], [], []
+        find_issues_input_dicts = [
+            {"features": dataset["X"]},
+            {"knn_graph": knn_graph},
+        ]
+
+        # Make sure that the results work for both input type
+        for kwargs in find_issues_input_dicts:
+            lab = Datalab(data=dataset, label_name=self.label_name)
+            assert lab.issue_summary.empty
+            lab.find_issues(**kwargs, issue_types={"data_valuation": {}})
+            summary = lab.get_issue_summary()
+            assert len(summary) == 1
+            assert "data_valuation" in summary["issue_type"].values
+            scores = lab.get_issues("data_valuation").get(["data_valuation_score"])
+            assert all((scores >= 0) & (scores <= 1))
+
+            datalabs.append(lab)
+            summaries.append(summary)
+            scores_list.append(scores)
+
+        # Check that the results are the same for both input types
+        base_lab = datalabs[0]
+        base_summary = summaries[0]
+        base_scores = scores_list[0]
+        for lab, summary, scores in zip(datalabs, summaries, scores_list):
+            # The knn-graph is either provided or computed from the features, then stored
+            assert np.allclose(
+                knn_graph.toarray(), lab.get_info("statistics")["weighted_knn_graph"].toarray()
+            )
+            # The summary and scores should be the same
+            assert base_summary.equals(summary)
+            assert np.allclose(base_scores, scores)
+
+    def test_find_issues_with_different_metrics(self, dataset, knn_graph):
+        """Test that a fresh Datalab instance can check for data_valuation issues with
+        different metrics.
+        """
+        knn_graph_euclidean = (
+            NearestNeighbors(n_neighbors=10, metric="euclidean")
+            .fit(dataset["X"])
+            .kneighbors_graph(mode="distance")
+        )
+
+        lab = Datalab(data=dataset, label_name=self.label_name)
+        lab.find_issues(features=dataset["X"], issue_types={"data_valuation": {}})
+
+        # The default metric should be "cosine" for "high-dimensional" features
+        assert lab.get_info("statistics")["knn_metric"] == "cosine"
+        assert np.allclose(
+            knn_graph.toarray(), lab.get_info("statistics")["weighted_knn_graph"].toarray()
+        )
+
+        # Test different scenarios of how the metric affects the knn graph
+        scenarios = [
+            {"metric": "cosine", "expected_knn_graph": knn_graph},
+            {"metric": "euclidean", "expected_knn_graph": knn_graph_euclidean},
+        ]
+
+        # Test what happens to the knn graph when the metric is changed
+        for scenario in scenarios:
+            metric = scenario["metric"]
+            expected_knn_graph = scenario["expected_knn_graph"]
+            lab.find_issues(
+                features=dataset["X"], issue_types={"data_valuation": {"metric": metric}}
+            )
+            assert metric == lab.get_info("statistics")["knn_metric"]
+            assert np.allclose(
+                expected_knn_graph.toarray(),
+                lab.get_info("statistics")["weighted_knn_graph"].toarray(),
+            )
+
+
 class TestIssueManagersReuseKnnGraph:
     """
     `outlier`, `underperforming_group` and `near_duplicate` issue managers require
diff --git a/tests/datalab/datalab/test_multilabel.py b/tests/datalab/datalab/test_multilabel.py
--- a/tests/datalab/datalab/test_multilabel.py
+++ b/tests/datalab/datalab/test_multilabel.py
@@ -119,7 +119,7 @@ def test_available_issue_types(self, lab):
             ["label", "near_duplicate", "non_iid", "outlier", "null"]
         )
         assert set(lab.list_possible_issue_types()) == set(
-            ["label", "near_duplicate", "non_iid", "outlier", "null"]
+            ["label", "near_duplicate", "non_iid", "outlier", "null", "data_valuation"]
         )
 
     @pytest.mark.parametrize(
diff --git a/tests/datalab/datalab/test_regression.py b/tests/datalab/datalab/test_regression.py
--- a/tests/datalab/datalab/test_regression.py
+++ b/tests/datalab/datalab/test_regression.py
@@ -154,7 +154,7 @@ def test_available_issue_types(self, lab):
             ["label", "outlier", "near_duplicate", "non_iid", "null"]
         )
         assert set(lab.list_possible_issue_types()) == set(
-            ["label", "outlier", "near_duplicate", "non_iid", "null"]
+            ["label", "outlier", "near_duplicate", "non_iid", "null", "data_valuation"]
         )
 
     def test_regression_with_features_finds_label_issues(self, lab, regression_data):