implement refitting the whole pipeline with all train data

Author: voorhs

autointent/_pipeline/_pipeline.py

@@ -122,7 +122,7 @@ def _is_inference(self) -> bool:
         """
         return isinstance(self.nodes[NodeType.scoring], InferenceNode)
 
-    def fit(self, dataset: Dataset, scheme: Literal["ho", "cv"] = "ho") -> Context:
+    def fit(self, dataset: Dataset, scheme: Literal["ho", "cv"] = "ho", refit_after: bool = False) -> Context:
         """
         Optimize the pipeline from dataset.
 
@@ -150,6 +150,9 @@ def fit(self, dataset: Dataset, scheme: Literal["ho", "cv"] = "ho") -> Context:
 
         self.nodes = {node.node_type: node for node in nodes_list}
 
+        if refit_after:
+            self._refit(context)
+
         predictions = self.predict(context.data_handler.test_utterances())
         for metric_name, metric in PREDICTION_METRICS_MULTILABEL.items():
             context.optimization_info.pipeline_metrics[metric_name] = metric(
@@ -210,6 +213,27 @@ def predict(self, utterances: list[str]) -> ListOfGenericLabels:
         scores = scoring_module.predict(utterances)
         return decision_module.predict(scores)
 
+    def _refit(self, context: Context) -> None:
+        """
+        Fit pipeline of already selected modules with all train data.
+
+        :param utterances: list of utterances
+        :return: list of predicted labels
+        """
+        if not self._is_inference():
+            msg = "Pipeline in optimization mode cannot perform inference"
+            raise RuntimeError(msg)
+
+        scoring_module: ScoringModule = self.nodes[NodeType.scoring].module  # type: ignore[assignment,union-attr]
+        decision_module: DecisionModule = self.nodes[NodeType.decision].module  # type: ignore[assignment,union-attr]
+
+        context.data_handler.prepare_for_refit()
+
+        scoring_module.fit(context.data_handler.train_utterances(0), context.data_handler.train_labels(0))
+        scores = scoring_module.predict(context.data_handler.train_utterances(1))
+
+        decision_module.fit(scores, context.data_handler.train_labels(1))
+
     def predict_with_metadata(self, utterances: list[str]) -> InferencePipelineOutput:
         """
         Predict the labels for the utterances with metadata.

autointent/context/data_handler/_data_handler.py

@@ -46,6 +46,7 @@ def __init__(
                             threshold search).
         """
         set_seed(random_seed)
+        self.random_seed = random_seed
 
         self.dataset = dataset
 
@@ -54,9 +55,9 @@ def __init__(
         self.n_folds = n_folds
 
         if scheme == "ho":
-            self._split_ho(random_seed, split_train)
+            self._split_ho(split_train)
         elif scheme == "cv":
-            self._split_cv(random_seed)
+            self._split_cv()
 
         self.regexp_patterns = [
             RegexPatterns(
@@ -185,20 +186,20 @@ def validation_iterator(self) -> Generator[tuple[list[str], ListOfLabels, list[s
             train_labels = [lab for lab in train_labels if lab is not None]
             yield train_utterances, train_labels, val_utterances, val_labels  # type: ignore[misc]
 
-    def _split_ho(self, random_seed: int, split_train: bool) -> None:
+    def _split_ho(self, split_train: bool) -> None:
         has_validation_split = any(split.startswith(Split.VALIDATION) for split in self.dataset)
 
         if split_train and Split.TRAIN in self.dataset:
-            self._split_train(random_seed)
+            self._split_train()
 
         if Split.TEST not in self.dataset:
             test_size = 0.1 if has_validation_split else 0.2
-            self._split_test(test_size, random_seed)
+            self._split_test(test_size)
 
         if not has_validation_split:
-            self._split_validation_from_train(random_seed)
+            self._split_validation_from_train()
         elif Split.VALIDATION in self.dataset:
-            self._split_validation(random_seed)
+            self._split_validation()
 
         for split in self.dataset:
             n_classes_split = self.dataset.get_n_classes(split)
@@ -209,7 +210,7 @@ def _split_ho(self, random_seed: int, split_train: bool) -> None:
                 )
                 raise ValueError(message)
 
-    def _split_train(self, random_seed: int) -> None:
+    def _split_train(self) -> None:
         """
         Split on two sets.
 
@@ -219,12 +220,12 @@ def _split_train(self, random_seed: int) -> None:
             self.dataset,
             split=Split.TRAIN,
             test_size=0.5,
-            random_seed=random_seed,
+            random_seed=self.random_seed,
             allow_oos_in_train=False,  # only train data for decision node should contain OOS
         )
         self.dataset.pop(Split.TRAIN)
 
-    def _split_validation(self, random_seed: int) -> None:
+    def _split_validation(self) -> None:
         """
         Split on two sets.
 
@@ -234,21 +235,21 @@ def _split_validation(self, random_seed: int) -> None:
             self.dataset,
             split=Split.VALIDATION,
             test_size=0.5,
-            random_seed=random_seed,
+            random_seed=self.random_seed,
             allow_oos_in_train=False,  # only val data for decision node should contain OOS
         )
         self.dataset.pop(Split.VALIDATION)
 
-    def _split_validation_from_test(self, random_seed: int) -> None:
+    def _split_validation_from_test(self) -> None:
         self.dataset[Split.TEST], self.dataset[Split.VALIDATION] = split_dataset(
             self.dataset,
             split=Split.TEST,
             test_size=0.5,
-            random_seed=random_seed,
+            random_seed=self.random_seed,
             allow_oos_in_train=True,  # both test and validation splits can contain OOS
         )
 
-    def _split_cv(self, random_seed: int) -> None:
+    def _split_cv(self) -> None:
         extra_splits = [split_name for split_name in self.dataset if split_name not in [Split.TRAIN, Split.TEST]]
         if extra_splits:
             self.dataset[Split.TRAIN] = concatenate_datasets(
@@ -257,26 +258,26 @@ def _split_cv(self, random_seed: int) -> None:
 
         if Split.TEST not in self.dataset:
             self.dataset[Split.TRAIN], self.dataset[Split.TEST] = split_dataset(
-                self.dataset, split=Split.TRAIN, test_size=0.2, random_seed=random_seed, allow_oos_in_train=True
+                self.dataset, split=Split.TRAIN, test_size=0.2, random_seed=self.random_seed, allow_oos_in_train=True
             )
 
         for j in range(self.n_folds - 1):
             self.dataset[Split.TRAIN], self.dataset[f"{Split.TRAIN}_{j}"] = split_dataset(
                 self.dataset,
                 split=Split.TRAIN,
                 test_size=1 / (self.n_folds - j),
-                random_seed=random_seed,
+                random_seed=self.random_seed,
                 allow_oos_in_train=True,
             )
         self.dataset[f"{Split.TRAIN}_{self.n_folds-1}"] = self.dataset.pop(Split.TRAIN)
 
-    def _split_validation_from_train(self, random_seed: int) -> None:
+    def _split_validation_from_train(self) -> None:
         if Split.TRAIN in self.dataset:
             self.dataset[Split.TRAIN], self.dataset[Split.VALIDATION] = split_dataset(
                 self.dataset,
                 split=Split.TRAIN,
                 test_size=0.2,
-                random_seed=random_seed,
+                random_seed=self.random_seed,
                 allow_oos_in_train=True,
             )
         else:
@@ -285,27 +286,46 @@ def _split_validation_from_train(self, random_seed: int) -> None:
                     self.dataset,
                     split=f"{Split.TRAIN}_{idx}",
                     test_size=0.2,
-                    random_seed=random_seed,
+                    random_seed=self.random_seed,
                     allow_oos_in_train=idx == 1,  # for decision node it's ok to have oos in train
                 )
 
-    def _split_test(self, test_size: float, random_seed: int) -> None:
+    def _split_test(self, test_size: float) -> None:
         """Obtain test set from train."""
         self.dataset[f"{Split.TRAIN}_0"], self.dataset[f"{Split.TEST}_0"] = split_dataset(
             self.dataset,
             split=f"{Split.TRAIN}_0",
             test_size=test_size,
-            random_seed=random_seed,
+            random_seed=self.random_seed,
         )
         self.dataset[f"{Split.TRAIN}_1"], self.dataset[f"{Split.TEST}_1"] = split_dataset(
             self.dataset,
             split=f"{Split.TRAIN}_1",
             test_size=test_size,
-            random_seed=random_seed,
+            random_seed=self.random_seed,
             allow_oos_in_train=True,
         )
         self.dataset[Split.TEST] = concatenate_datasets(
             [self.dataset[f"{Split.TEST}_0"], self.dataset[f"{Split.TEST}_1"]],
         )
         self.dataset.pop(f"{Split.TEST}_0")
         self.dataset.pop(f"{Split.TEST}_1")
+
+    def prepare_for_refit(self) -> None:
+        if self.scheme == "ho":
+            return
+
+        train_folds = [split_name for split_name in self.dataset if split_name.startswith("train")]
+        self.dataset[Split.TRAIN] = concatenate_datasets([self.dataset[name] for name in train_folds])
+        for name in train_folds:
+            self.dataset.pop(name)
+
+        self.dataset[f"{Split.TRAIN}_0"], self.dataset[f"{Split.TRAIN}_1"] = split_dataset(
+            self.dataset,
+            split=Split.TRAIN,
+            test_size=0.5,
+            random_seed=self.random_seed,
+            allow_oos_in_train=False,
+        )
+
+        self.dataset.pop(Split.TRAIN)

tests/pipeline/test_optimization.py

@@ -31,7 +31,7 @@ def test_cv(dataset, task_type):
     if task_type == "multilabel":
         dataset = dataset.to_multilabel()
 
-    context = pipeline_optimizer.fit(dataset, scheme="cv")
+    context = pipeline_optimizer.fit(dataset, scheme="cv", refit_after=True)
     context.dump()
 
     assert os.listdir(pipeline_optimizer.logging_config.dump_dir)