feat: log calculated metric with SplitTrainer

psycop/common/model_training_v2/loggers/base_logger.py

@@ -3,6 +3,7 @@
 import wasabi
 
 from psycop.common.global_utils.config_utils import flatten_nested_dict
+from psycop.common.model_training_v2.metrics.base_metric import CalculatedMetric
 
 
 class BaselineLogger(Protocol):
@@ -15,7 +16,7 @@ def warn(self, message: str) -> None:
     def fail(self, message: str) -> None:
         ...
 
-    def log_metric(self, name: str, value: float) -> None:
+    def log_metric(self, metric: CalculatedMetric) -> None:
         ...
 
     def log_config(self, config: dict[str, Any]) -> None:
@@ -35,9 +36,9 @@ def warn(self, message: str) -> None:
     def fail(self, message: str) -> None:
         self._l.fail(message)
 
-    def log_metric(self, name: str, value: float) -> None:
-        self._l.divider(f"Logging metric {name}")
-        self._l.info(f"{name}: {value}")
+    def log_metric(self, metric: CalculatedMetric) -> None:
+        self._l.divider(f"Logging metric {metric.name}")
+        self._l.info(f"{metric.name}: {metric.value}")
 
     def log_config(self, config: dict[str, Any]) -> None:
         self._l.divider("Logging config")

psycop/common/model_training_v2/pipeline.py

@@ -23,4 +23,4 @@ def train_baseline_model(cfg: BaselineSchema) -> float:
     result = cfg.training_method.train()
     result.eval_dataset.to_disk(path=cfg.experiment_path)
 
-    return result.metric
+    return result.metric.value

psycop/common/model_training_v2/training_method/base_training_method.py

@@ -3,6 +3,8 @@
 from dataclasses import dataclass
 from typing import Protocol
 
+from psycop.common.model_training_v2.metrics.base_metric import CalculatedMetric
+
 from ..loggers.base_logger import BaselineLogger
 from ..presplit_preprocessing.pipeline import PreprocessingPipeline
 from ..presplit_preprocessing.polars_frame import PolarsFrame
@@ -12,7 +14,7 @@
 
 @dataclass(frozen=True)
 class TrainingResult:
-    metric: float
+    metric: CalculatedMetric
     eval_dataset: BaseEvalDataset
 
 
@@ -60,18 +62,20 @@ def train(self) -> TrainingResult:
         training_data_preprocessed = self.preprocessing_pipeline.apply(
             data=self.training_data,
         )
-        self.preprocessing_pipeline.apply(
+        validation_data = self.preprocessing_pipeline.apply(
             data=self.validation_data,
         )
+
         self.problem_type.train(
             x=training_data_preprocessed.drop(self.training_outcome_col_name),
             y=training_data_preprocessed.select(self.training_outcome_col_name),
         )
+
         result = self.problem_type.evaluate(
-            x=self.validation_data.drop(self.validation_outcome_col_name),
-            y=self.validation_data.select(self.validation_outcome_col_name),
+            x=validation_data.drop(self.validation_outcome_col_name),
+            y=validation_data.select(self.validation_outcome_col_name),
         )
 
-        self.logger.log_metric(name="metric", value=result.metric)
+        self.logger.log_metric(result.metric)
 
         return result