Fix unleved bug with bert

scikit-learn-contrib · Dec 2, 2024 · a4a006e · a4a006e
1 parent a344014
commit a4a006e
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 1 deletion.
diff --git a/hiclass/HierarchicalClassifier.py b/hiclass/HierarchicalClassifier.py
@@ -161,7 +161,7 @@ def _pre_fit(self, X, y, sample_weight):
             )
         else:
             self.X_ = np.array(X)
-            self.y_ = np.array(y)
+            self.y_ = np.array(make_leveled(y))
 
         if sample_weight is not None:
             self.sample_weight_ = _check_sample_weight(sample_weight, X)

diff --git a/tests/test_LocalClassifiers.py b/tests/test_LocalClassifiers.py
@@ -143,3 +143,14 @@ def test_tmp_dir(classifier):
         assert expected_name == name
         check_is_fitted(classifier)
         clf.fit(x, y)
+
+
+@pytest.mark.parametrize("classifier", classifiers)
+def test_bert_unleveled(classifier):
+    clf = classifier(
+        local_classifier=LogisticRegression(),
+        bert=True,
+    )
+    x = [[0, 1], [2, 3]]
+    y = [["a"], ["b", "c"]]
+    clf.fit(x, y)