bug fix

Author: voorhs

autointent/_pipeline/_pipeline.py

@@ -229,10 +229,10 @@ def _refit(self, context: Context) -> None:
 
         context.data_handler.prepare_for_refit()
 
-        scoring_module.fit(context.data_handler.train_utterances(0), context.data_handler.train_labels(0))  # type: ignore[arg-type]
+        scoring_module.fit(*scoring_module.get_train_data(context))  # type: ignore[arg-type]
         scores = scoring_module.predict(context.data_handler.train_utterances(1))
 
-        decision_module.fit(scores, context.data_handler.train_labels(1))
+        decision_module.fit(scores, context.data_handler.train_labels(1), context.data_handler.tags)
 
     def predict_with_metadata(self, utterances: list[str]) -> InferencePipelineOutput:
         """

tests/data/test_stratificaiton.py

@@ -38,6 +38,6 @@ def test_multilabel_train_test_split(dataset_unsplitted):
 
     assert Split.TRAIN in dataset
     assert Split.TEST in dataset
-    assert dataset[Split.TRAIN].num_rows == 17
-    assert dataset[Split.TEST].num_rows == 19
+    assert dataset[Split.TRAIN].num_rows == 19
+    assert dataset[Split.TEST].num_rows == 17
     assert dataset.get_n_classes(Split.TRAIN) == dataset.get_n_classes(Split.TEST)