separate class balance errors

Author: paroma

snorkel/labeling/model/label_model.py

@@ -539,22 +539,25 @@ def _set_class_balance(
         """
         if class_balance is not None:
             self.p = np.array(class_balance)
+            if len(self.p) != self.cardinality:
+                raise ValueError(
+                    f"class_balance has {len(self.p)} entries. Does not match LabelModel cardinality {self.cardinality}."
+                )
         elif Y_dev is not None:
             class_counts = Counter(Y_dev)
             sorted_counts = np.array([v for k, v in sorted(class_counts.items())])
             self.p = sorted_counts / sum(sorted_counts)
+            if len(self.p) != self.cardinality:
+                raise ValueError(
+                    f"Y_dev has {len(self.p)} class(es). Does not match LabelModel cardinality {self.cardinality}."
+                )
         else:
             self.p = (1 / self.cardinality) * np.ones(self.cardinality)
 
         if np.any(self.p == 0):
             raise ValueError(
                 f"Class balance prior is 0 for class(es) {np.where(self.p)[0]}."
             )
-        if len(self.p) != self.cardinality:
-            raise ValueError(
-                f"Y_dev has {len(self.p)} class(es). Does not match LabelModel cardinality {self.cardinality}."
-            )
-
         self.P = torch.diag(torch.from_numpy(self.p)).float()
 
     def _set_constants(self, L: np.ndarray) -> None:

test/labeling/model/test_label_model.py

@@ -69,8 +69,14 @@ def test_class_balance(self):
         with self.assertRaisesRegex(ValueError, "Class balance prior is 0"):
             label_model._set_class_balance(class_balance=class_balance, Y_dev=Y_dev)
 
+        class_balance = np.array([0.0])
+        with self.assertRaisesRegex(ValueError, "class_balance has 1 entries."):
+            label_model._set_class_balance(class_balance=class_balance, Y_dev=Y_dev)
+
         Y_dev_one_class = np.array([0, 0, 0])
-        with self.assertRaisesRegex(ValueError, "Does not match LabelModel cardinality"):
+        with self.assertRaisesRegex(
+            ValueError, "Does not match LabelModel cardinality"
+        ):
             label_model._set_class_balance(class_balance=None, Y_dev=Y_dev_one_class)
 
     def test_generate_O(self):