Add final pipeline metrics to logs (#83)

Add final pipeline metrics to logs

Author: truff4ut

autointent/_pipeline/_pipeline.py

@@ -12,6 +12,7 @@
 from autointent import Context, Dataset
 from autointent.configs import EmbedderConfig, InferenceNodeConfig, LoggingConfig, VectorIndexConfig
 from autointent.custom_types import NodeType
+from autointent.metrics import PREDICTION_METRICS_MULTILABEL
 from autointent.nodes import InferenceNode, NodeOptimizer
 from autointent.utils import load_default_search_space, load_search_space
 
@@ -103,7 +104,7 @@ def _is_inference(self) -> bool:
         """
         return isinstance(self.nodes[NodeType.scoring], InferenceNode)
 
-    def fit(self, dataset: Dataset, force_multilabel: bool = False, init_for_inference: bool = True) -> Context:
+    def fit(self, dataset: Dataset, force_multilabel: bool = False) -> Context:
         """
         Optimize the pipeline from dataset.
 
@@ -122,15 +123,20 @@ def fit(self, dataset: Dataset, force_multilabel: bool = False, init_for_inferen
 
         self._fit(context)
 
-        if init_for_inference:
-            if context.is_ram_to_clear():
-                nodes_configs = context.optimization_info.get_inference_nodes_config()
-                nodes_list = [InferenceNode.from_config(cfg) for cfg in nodes_configs]
-            else:
-                modules_dict = context.optimization_info.get_best_modules()
-                nodes_list = [InferenceNode(module, node_type) for node_type, module in modules_dict.items()]
+        if context.is_ram_to_clear():
+            nodes_configs = context.optimization_info.get_inference_nodes_config()
+            nodes_list = [InferenceNode.from_config(cfg) for cfg in nodes_configs]
+        else:
+            modules_dict = context.optimization_info.get_best_modules()
+            nodes_list = [InferenceNode(module, node_type) for node_type, module in modules_dict.items()]
+
+        self.nodes = {node.node_type: node for node in nodes_list}
 
-            self.nodes = {node.node_type: node for node in nodes_list}
+        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(
+                context.data_handler.test_labels(), predictions,
+            )
 
         return context
 

autointent/context/optimization_info/_optimization_info.py

@@ -65,6 +65,7 @@ def __init__(self) -> None:
         self.trials = Trials()
         self._trials_best_ids = TrialsIds()
         self.modules = ModulesList()
+        self.pipeline_metrics: dict[str, float] = {}
 
     def log_module_optimization(
         self,
@@ -196,6 +197,7 @@ def dump_evaluation_results(self) -> dict[str, Any]:
         """
         node_wise_metrics = {node_type: self._get_metrics_values(node_type) for node_type in NodeType}
         return {
+            "pipeline_metrics": self.pipeline_metrics,
             "metrics": node_wise_metrics,
             "configs": self.trials.model_dump(),
         }

tests/pipeline/test_optimization.py

@@ -44,7 +44,7 @@ def test_no_context_optimization(dataset, task_type):
     pipeline_optimizer.set_config(VectorIndexConfig(db_dir=Path(db_dir).resolve()))
     pipeline_optimizer.set_config(EmbedderConfig(batch_size=16, max_length=32, device="cpu"))
 
-    context = pipeline_optimizer.fit(dataset, force_multilabel=(task_type == "multilabel"), init_for_inference=False)
+    context = pipeline_optimizer.fit(dataset, force_multilabel=(task_type == "multilabel"))
     context.dump()
 
 
@@ -62,7 +62,7 @@ def test_save_db(dataset, task_type):
     pipeline_optimizer.set_config(VectorIndexConfig(db_dir=Path(db_dir).resolve(), save_db=True))
     pipeline_optimizer.set_config(EmbedderConfig(batch_size=16, max_length=32, device="cpu"))
 
-    context = pipeline_optimizer.fit(dataset, force_multilabel=(task_type == "multilabel"), init_for_inference=False)
+    context = pipeline_optimizer.fit(dataset, force_multilabel=(task_type == "multilabel"))
     context.dump()
 
     assert os.listdir(db_dir)
@@ -82,7 +82,7 @@ def test_dump_modules(dataset, task_type):
     pipeline_optimizer.set_config(VectorIndexConfig(db_dir=Path(db_dir).resolve()))
     pipeline_optimizer.set_config(EmbedderConfig(batch_size=16, max_length=32, device="cpu"))
 
-    context = pipeline_optimizer.fit(dataset, force_multilabel=(task_type == "multilabel"), init_for_inference=False)
+    context = pipeline_optimizer.fit(dataset, force_multilabel=(task_type == "multilabel"))
     context.dump()
 
     assert os.listdir(dump_dir)