Repository: cleanlab/cleanlab

Issue:
add class_imbalance issue type among the Datalab defaults
I'd test it on a couple different types of datasets (including w extreme imbalance -- only 1 example of minority class) before adding it to the defaults.

Make sure to also comment on it in the issue guide and tutorial once this is added to defaults.


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
@@ -55,6 +55,7 @@
     "non_iid": ["pred_probs", "features", "knn_graph"],
     "underperforming_group": ["pred_probs", "features", "knn_graph", "cluster_ids"],
     "data_valuation": ["knn_graph"],
+    "class_imbalance": [],
 }
 _REGRESSION_ARGS_DICT = {
     "label": ["features", "predictions"],
@@ -69,9 +70,15 @@ def _resolve_required_args_for_classification(**kwargs):
         for issue_type in initial_args_dict
     }
 
+    # Some issue types (like class-imbalance) have no required args.
+    # This conditional lambda is used to include them in args dict.
+    keep_empty_argument = lambda k: not len(_CLASSIFICATION_ARGS_DICT[k])
+
     # Remove None values from argument list, rely on default values in IssueManager
     args_dict = {
-        k: {k2: v2 for k2, v2 in v.items() if v2 is not None} for k, v in args_dict.items() if v
+        k: {k2: v2 for k2, v2 in v.items() if v2 is not None}
+        for k, v in args_dict.items()
+        if (v or keep_empty_argument(k))
     }
 
     # Prefer `knn_graph` over `features` if both are provided.
@@ -91,7 +98,8 @@ def _resolve_required_args_for_classification(**kwargs):
             )
 
     # Only keep issue types that have at least one argument
-    args_dict = {k: v for k, v in args_dict.items() if v}
+    # or those that require no arguments.
+    args_dict = {k: v for k, v in args_dict.items() if (v or keep_empty_argument(k))}
 
     return args_dict
 
@@ -103,12 +111,15 @@ def _resolve_required_args_for_regression(**kwargs):
         issue_type: {arg: kwargs.get(arg, None) for arg in initial_args_dict[issue_type]}
         for issue_type in initial_args_dict
     }
+    # Some issue types have no required args.
+    # This conditional lambda is used to include them in args dict.
+    keep_empty_argument = lambda k: not len(_REGRESSION_ARGS_DICT[k])
 
     # Remove None values from argument list, rely on default values in IssueManager
     args_dict = {
         k: {k2: v2 for k2, v2 in v.items() if v2 is not None}
         for k, v in args_dict.items()
-        if v or k == "label"  # Allow label issues to require no arguments
+        if v or k == "label" or keep_empty_argument(k)  # Allow label issues to require no arguments
     }
 
     return args_dict
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
@@ -203,10 +203,5 @@ def list_default_issue_types(task: str) -> List[str]:
     if task == "regression":
         default_issue_types = ["label"]
     else:
-        default_issue_types = [
-            "label",
-            "outlier",
-            "near_duplicate",
-            "non_iid",
-        ]
+        default_issue_types = ["label", "outlier", "near_duplicate", "non_iid", "class_imbalance"]
     return default_issue_types


Test patch:
diff --git a/tests/datalab/test_cleanvision_integration.py b/tests/datalab/test_cleanvision_integration.py
--- a/tests/datalab/test_cleanvision_integration.py
+++ b/tests/datalab/test_cleanvision_integration.py
@@ -32,7 +32,7 @@ def num_imagelab_issues(self):
 
     @pytest.fixture
     def num_datalab_issues(self):
-        return 3
+        return 4
 
     @pytest.fixture
     def pred_probs(self, image_dataset):
@@ -67,6 +67,7 @@ def test_imagelab_issues_checked(
             "label",
             "outlier",
             "near_duplicate",
+            "class_imbalance"
             # "non_iid",
         ]
 
@@ -90,8 +91,9 @@ def test_imagelab_issues_checked(
                     "label",
                     "outlier",
                     "near_duplicate",
+                    "class_imbalance",
                 ],
-                "num_issues": [1, 1, 0, 1, 1, 1, 1, 0, 0, 0],
+                "num_issues": [1, 1, 0, 1, 1, 1, 1, 0, 0, 0, 0],
             }
         )
         expected_count = df.sort_values(by="issue_type")["num_issues"].tolist()
@@ -143,12 +145,7 @@ def test_imagelab_issues_not_checked(
         assert len(datalab.issues.columns) == num_datalab_issues * 2
         assert len(datalab.issue_summary) == num_datalab_issues
 
-        all_keys = [
-            "statistics",
-            "label",
-            "outlier",
-            "near_duplicate",
-        ]
+        all_keys = ["statistics", "label", "outlier", "near_duplicate", "class_imbalance"]
 
         assert set(all_keys) == set(datalab.info.keys())
         datalab.report()
diff --git a/tests/datalab/test_datalab.py b/tests/datalab/test_datalab.py
--- a/tests/datalab/test_datalab.py
+++ b/tests/datalab/test_datalab.py
@@ -87,6 +87,7 @@ def test_list_default_issue_types(self):
             "outlier",
             "near_duplicate",
             "non_iid",
+            "class_imbalance",
         ]
 
     def tmp_path(self):
@@ -585,13 +586,14 @@ def test_features_and_knn_graph(self, data_tuple):
         assert lab.get_info("statistics").get("knn_metric") == "cosine"
 
     def test_without_features_or_knn_graph(self, data_tuple):
-        """Test that the `knn_graph` argument to `find_issues` is used instead of computing a new
-        one from the `features` argument."""
+        """Test that only the class_imbalance issue is run
+        when no features, knn_graph or pred_probs are passed."""
         lab, _, _ = data_tuple
 
         # Test that a warning is raised
         lab.find_issues()
-        assert lab.issues.empty  # No columns should be added to the issues dataframe
+        # Only class_imbalance issue columns should be present
+        assert list(lab.issues.columns) == ["is_class_imbalance_issue", "class_imbalance_score"]
 
     def test_data_valuation_issue_with_knn_graph(self, data_tuple):
         lab, knn_graph, features = data_tuple
@@ -856,10 +858,10 @@ def test_find_non_iid_issues_sorted_using_pred_probs(self, lab, sorted_embedding
     def test_incremental_search(self, lab, sorted_embeddings):
         lab.find_issues(features=sorted_embeddings)
         summary = lab.get_issue_summary()
-        assert len(summary) == 3
+        assert len(summary) == 4
         lab.find_issues(features=sorted_embeddings, issue_types={"non_iid": {}})
         summary = lab.get_issue_summary()
-        assert len(summary) == 3
+        assert len(summary) == 4
         assert "non_iid" in summary["issue_type"].values
         non_iid_summary = lab.get_issue_summary("non_iid")
         assert non_iid_summary["score"].values[0] == 0
@@ -929,11 +931,11 @@ def test_incremental_search(self, pred_probs, random_embeddings):
         lab = Datalab(data=data, label_name="labels")
         lab.find_issues(features=random_embeddings)
         summary = lab.get_issue_summary()
-        assert len(summary) == 4
+        assert len(summary) == 5
         assert "label" in summary["issue_type"].values
         lab.find_issues(pred_probs=pred_probs, issue_types={"label": {}})
         summary = lab.get_issue_summary()
-        assert len(summary) == 4
+        assert len(summary) == 5
         assert "label" in summary["issue_type"].values
         label_summary = lab.get_issue_summary("label")
         assert label_summary["num_issues"].values[0] > 0
@@ -1314,8 +1316,9 @@ def test_find_issues_features_works_with_and_without_labels(self, features, labe
         issues_with_labels = lab_with_labels.issues
         issues_without_label_name = lab_without_label_name.issues
 
-        # issues_with_labels should have two additional columns about label issues
-        assert len(issues_without_labels.columns) + 2 == len(issues_with_labels.columns)
+        # issues_with_labels should have four additional columns, which include label issues
+        # and class_imbalance issues
+        assert len(issues_without_labels.columns) + 4 == len(issues_with_labels.columns)
         pd.testing.assert_frame_equal(issues_without_labels, issues_without_label_name)
 
 
diff --git a/tests/datalab/test_issue_finder.py b/tests/datalab/test_issue_finder.py
--- a/tests/datalab/test_issue_finder.py
+++ b/tests/datalab/test_issue_finder.py
@@ -25,13 +25,13 @@ def test_init(self, issue_finder):
         assert issue_finder.verbosity == 1
 
     def test_get_available_issue_types(self, issue_finder):
-        expected_issue_types = {}
+        expected_issue_types = {"class_imbalance": {}}
         # Test with no kwargs, no issue type expected to be returned
         for key in ["pred_probs", "features", "knn_graph"]:
             issue_types = issue_finder.get_available_issue_types(**{key: None})
             assert (
                 issue_types == expected_issue_types
-            ), "Every issue type for classification requires some kwargs, expected empty dict"
+            ), "Only class_imbalance issue type for classification requires no kwargs"
 
         # Test with only issue_types, input should be
         issue_types_dicts = [