refactor: artifacts are added only for best module

Author: voorhs

autointent/context/optimization_info/_optimization_info.py

@@ -74,7 +74,6 @@ def log_module_optimization(
         module_params: dict[str, Any],
         metric_value: float,
         metric_name: str,
-        artifact: Artifact,
         module_dump_dir: str | None,
         module: "Module | None" = None,
     ) -> None:
@@ -103,13 +102,11 @@ def log_module_optimization(
         if module:
             self.modules.add_module(node_type, module)
 
-        self.artifacts.add_artifact(node_type, artifact)
-
     def _get_metrics_values(self, node_type: str) -> list[float]:
         """Retrieve all metric values for a specific node type."""
         return [trial.metric_value for trial in self.trials.get_trials(node_type)]
 
-    def _get_best_trial_idx(self, node_type: str) -> int | None:
+    def get_best_trial_idx(self, node_type: str) -> int | None:
         """
         Retrieve the index of the best trial for a node type.
 
@@ -133,7 +130,7 @@ def _get_best_artifact(self, node_type: str) -> RetrieverArtifact | ScorerArtifa
         :return: The best artifact for the node type.
         :raises ValueError: If no best trial exists for the node type.
         """
-        best_idx = self._get_best_trial_idx(node_type)
+        best_idx = self.get_best_trial_idx(node_type)
         if best_idx is None:
             msg = f"No best trial for {node_type}"
             raise ValueError(msg)
@@ -194,7 +191,7 @@ def get_inference_nodes_config(self, asdict: bool = False) -> list[InferenceNode
 
         :return: List of `InferenceNodeConfig` objects for inference nodes.
         """
-        trial_ids = [self._get_best_trial_idx(node_type) for node_type in NodeType]
+        trial_ids = [self.get_best_trial_idx(node_type) for node_type in NodeType]
         res = []
         for idx, node_type in zip(trial_ids, NodeType, strict=True):
             if idx is None:
@@ -216,7 +213,7 @@ def _get_best_module(self, node_type: str) -> "Module | None":
         :param node_type: Type of the node.
         :return: The best module, or None if no best trial exists.
         """
-        idx = self._get_best_trial_idx(node_type)
+        idx = self.get_best_trial_idx(node_type)
         if idx is not None:
             return self.modules.get(node_type)[idx]
         return None

autointent/modules/abc/_base.py

@@ -44,7 +44,7 @@ def score(self, context: Context, split: Literal["validation", "test"], metrics:
         """
 
     @abstractmethod
-    def get_assets(self) -> Artifact:
+    def get_artifact(self, context: Context) -> Artifact:
         """Return useful assets that represent intermediate data into context."""
 
     @abstractmethod

autointent/modules/abc/_decision.py

@@ -49,13 +49,14 @@ def score(self, context: Context, split: Literal["validation", "test"], metrics:
         :return: Computed metrics value for the test set or error code of metrics
         """
         labels, scores = get_decision_evaluation_data(context, split)
-        self._decisions = self.predict(scores)
+        decisions = self.predict(scores)
         chosen_metrics = {name: fn for name, fn in PREDICTION_METRICS_MULTICLASS.items() if name in metrics}
-        return self.score_metrics((labels, self._decisions), chosen_metrics)
+        return self.score_metrics((labels, decisions), chosen_metrics)
 
-    def get_assets(self) -> DecisionArtifact:
+    def get_artifact(self, context: Context) -> DecisionArtifact:
         """Return useful assets that represent intermediate data into context."""
-        return DecisionArtifact(labels=self._decisions)
+        _, scores = get_decision_evaluation_data(context, split="test")
+        return DecisionArtifact(labels=self.predict(scores))
 
     def clear_cache(self) -> None:
         """Clear cache."""

autointent/modules/abc/_scoring.py

@@ -41,24 +41,24 @@ def score(self, context: Context, split: Literal["validation", "test"], metrics:
 
         scores = self.predict(utterances)
 
-        self._train_scores = self.predict(context.data_handler.train_utterances(1))
-        self._validation_scores = self.predict(context.data_handler.validation_utterances(1))
-        self._test_scores = self.predict(context.data_handler.test_utterances())
-
         metrics_dict = SCORING_METRICS_MULTILABEL if context.is_multilabel() else SCORING_METRICS_MULTICLASS
         chosen_metrics = {name: fn for name, fn in metrics_dict.items() if name in metrics}
         return self.score_metrics((labels, scores), chosen_metrics)
 
-    def get_assets(self) -> ScorerArtifact:
+    def get_artifact(self, context: Context) -> ScorerArtifact:
         """
         Retrieve assets generated during scoring.
 
         :return: ScorerArtifact containing test, validation and test scores.
         """
+        train_scores = self.predict(context.data_handler.train_utterances(1))
+        validation_scores = self.predict(context.data_handler.validation_utterances(1))
+        test_scores = self.predict(context.data_handler.test_utterances())
+
         return ScorerArtifact(
-            train_scores=self._train_scores,
-            validation_scores=self._validation_scores,
-            test_scores=self._test_scores,
+            train_scores=train_scores,
+            validation_scores=validation_scores,
+            test_scores=test_scores,
         )
 
     @abstractmethod

autointent/modules/embedding/_logreg.py

@@ -152,7 +152,7 @@ def score(self, context: Context, split: Literal["validation", "test"], metrics:
         chosen_metrics = {name: fn for name, fn in metrics_dict.items() if name in metrics}
         return self.score_metrics((labels, probas), chosen_metrics)
 
-    def get_assets(self) -> RetrieverArtifact:
+    def get_artifact(self, context: Context) -> RetrieverArtifact:
         """
         Get the classifier artifacts for this module.
 

autointent/modules/embedding/_retrieval.py

@@ -132,7 +132,7 @@ def score(self, context: Context, split: Literal["validation", "test"], metrics:
         chosen_metrics = {name: fn for name, fn in metrics_dict.items() if name in metrics}
         return self.score_metrics((labels, predictions), chosen_metrics)
 
-    def get_assets(self) -> RetrieverArtifact:
+    def get_artifact(self, context: Context) -> RetrieverArtifact:
         """
         Get the retriever artifacts for this module.
 

autointent/modules/regexp/_regexp.py

@@ -133,7 +133,7 @@ def clear_cache(self) -> None:
         """Clear cache."""
         del self.regexp_patterns
 
-    def get_assets(self) -> Artifact:
+    def get_artifact(self) -> Artifact:
         """Get assets."""
         return Artifact()
 

autointent/nodes/_optimization/_node_optimizer.py

@@ -51,7 +51,7 @@ def fit(self, context: Context) -> None:
         :param context: Context
         """
         self._logger.info("starting %s node optimization...", self.node_info.node_type)
-
+        scored_modules = []
         for search_space in deepcopy(self.modules_search_spaces):
             module_name = search_space.pop("module_name")
 
@@ -62,7 +62,10 @@ def fit(self, context: Context) -> None:
                 context.callback_handler.start_module(
                     module_name=module_name, num=j_combination, module_kwargs=module_kwargs
                 )
-                module = self.node_info.modules_available[module_name].from_context(context, **module_kwargs)
+                module_type = self.node_info.modules_available[module_name]
+                module = module_type.from_context(context, **module_kwargs)
+
+                scored_modules.append((module_type, module_kwargs))
 
                 embedder_name = module.get_embedder_name()
                 if embedder_name is not None:
@@ -92,7 +95,6 @@ def fit(self, context: Context) -> None:
                     module_kwargs,
                     metric_value,
                     self.target_metric,
-                    module.get_assets(),  # retriever name / scores / predictions
                     module_dump_dir,
                     module=module if not context.is_ram_to_clear() else None,
                 )
@@ -102,7 +104,14 @@ def fit(self, context: Context) -> None:
                     gc.collect()
                     torch.cuda.empty_cache()
 
-        self._logger.info("%s node optimization is finished!", self.node_info.node_type)
+        self._logger.info("%s node optimization is finished! saving best assets", self.node_info.node_type)
+        # TODO refactor the following code (via implementing `autointent.load_module(path)` utility)
+        trial_idx = context.optimization_info.get_best_trial_idx(self.node_type)
+        trial = context.optimization_info.trials.get_trial(self.node_type, trial_idx)
+        module_type, module_kwargs = scored_modules[trial_idx]
+        best_module: Module = module_type(**module_kwargs)
+        best_module.load(trial.module_dump_dir)
+        context.optimization_info.artifacts.add_artifact(self.node_type, best_module.get_artifact(context))
 
     def get_module_dump_dir(self, dump_dir: Path, module_name: str, j_combination: int) -> str:
         """

tests/modules/embedding/test_logreg.py

@@ -1,9 +1,9 @@
 from autointent.modules.embedding import LogregAimedEmbedding
 
 
-def test_get_assets_returns_correct_artifact_for_logreg():
+def test_get_artifact_returns_correct_artifact_for_logreg():
     module = LogregAimedEmbedding(embedder_name="sergeyzh/rubert-tiny-turbo")
-    artifact = module.get_assets()
+    artifact = module.get_artifact()
     assert artifact.embedder_name == "sergeyzh/rubert-tiny-turbo"
 
 

tests/modules/embedding/test_retrieval.py

@@ -4,9 +4,9 @@
 from tests.conftest import setup_environment
 
 
-def test_get_assets_returns_correct_artifact():
+def test_get_artifact_returns_correct_artifact():
     module = RetrievalAimedEmbedding(k=5, embedder_name="sergeyzh/rubert-tiny-turbo")
-    artifact = module.get_assets()
+    artifact = module.get_artifact()
     assert artifact.embedder_name == "sergeyzh/rubert-tiny-turbo"