fixes in metrics logging

Author: manucouto1

framework3/base/base_clases.py

@@ -648,7 +648,7 @@ def _pre_predict(self, x: XYData) -> XYData:
         except Exception:
             raise ValueError("Trainable filter model not trained or loaded")
 
-    def _pre_fit_wrapp(self, x: XYData, y: Optional[XYData]) -> Optional[float]:
+    def _pre_fit_wrapp(self, x: XYData, y: Optional[XYData]) -> Optional[float | dict]:
         """
         Wrapper method for the fit function.
 
@@ -712,7 +712,7 @@ def __setstate__(self, state: Dict[str, Any]):
         self.__dict__["fit"] = self._pre_fit_wrapp
         self.__dict__["predict"] = self._pre_predict_wrapp
 
-    def fit(self, x: XYData, y: Optional[XYData]) -> Optional[float]:
+    def fit(self, x: XYData, y: Optional[XYData]) -> Optional[float | dict]:
         """
         Method for fitting the filter to the data.
 

framework3/base/base_splitter.py

@@ -150,7 +150,7 @@ def evaluate(
         """
         ...
 
-    def _pre_fit_wrapp(self, x: XYData, y: Optional[XYData]) -> float | None:
+    def _pre_fit_wrapp(self, x: XYData, y: Optional[XYData]) -> float | None | dict:
         """
         Wrapper method for pre-fitting.
 

framework3/plugins/filters/cache/cached_filter.py

@@ -137,7 +137,7 @@ def init(self) -> None:
         self.filter.init()
         super().init()
 
-    def _pre_fit_wrapp(self, x: XYData, y: Optional[XYData]) -> float | None:
+    def _pre_fit_wrapp(self, x: XYData, y: Optional[XYData]) -> float | None | dict:
         """
         Wrapper method for the pre-fit stage.
 

framework3/plugins/optimizer/grid_optimizer.py

@@ -183,11 +183,11 @@ def fit(self, x: XYData, y: Optional[XYData]):
             match pipeline.fit(x, y):
                 case None:
                     losses = pipeline.evaluate(x, y, pipeline.predict(x))
-
                     score = losses.get(self.scorer.__class__.__name__, 0.0)
-
                 case float() as score:
                     pass
+                case dict() as losses:
+                    score = losses.get(self.scorer.__class__.__name__, 0.0)
                 case _:
                     raise ValueError("Unexpected return type from pipeline.fit()")
 

framework3/plugins/optimizer/optuna_optimizer.py

@@ -1,8 +1,9 @@
 import optuna
 
 from typing import Any, Callable, Dict, Sequence, Union, cast
+from framework3 import F1
 from framework3.container import Container
-from framework3.base import BasePlugin, XYData
+from framework3.base import BaseMetric, BasePlugin, XYData
 
 from rich import print
 
@@ -69,6 +70,7 @@ def __init__(
         pipeline: BaseFilter | None = None,
         study_name: str | None = None,
         storage: str | None = None,
+        scorer: BaseMetric = F1(),
     ):
         """
         Initialize the OptunaOptimizer.
@@ -90,6 +92,7 @@ def __init__(
         self.n_trials = n_trials
         self.load_if_exists = load_if_exists
         self.reset_study = reset_study
+        self.scorer = scorer
 
     def optimize(self, pipeline: BaseFilter):
         """
@@ -227,17 +230,12 @@ def matcher(k, v):
 
                 match pipeline.fit(x, y):
                     case None:
-                        return float(
-                            next(
-                                iter(
-                                    pipeline.evaluate(
-                                        x, y, pipeline.predict(x)
-                                    ).values()
-                                )
-                            )
-                        )
+                        metrics = pipeline.evaluate(x, y, pipeline.predict(x))
+                        return float(metrics.get(self.scorer.__class__.__name__, 0.0))
                     case float() as loss:
                         return loss
+                    case dict() as losses:
+                        return float(losses.get(self.scorer.__class__.__name__, 0.0))
                     case _:
                         raise ValueError("Unsupported type in pipeline.fit")
 

framework3/plugins/optimizer/wandb_optimizer.py

@@ -160,6 +160,10 @@ def exec(
                 return {self.scorer.__class__.__name__: float(loss)}
             case float() as loss:
                 return {self.scorer.__class__.__name__: loss}
+            case dict() as losses:
+                loss = losses.pop(self.scorer.__class__.__name__, 0.0)
+                wandb.log(dict(losses))  # type: ignore[attr-defined]
+                return {self.scorer.__class__.__name__: loss}
             case _:
                 raise ValueError("Unexpected return type from pipeline.fit()")
 

framework3/plugins/pipelines/sequential/f3_pipeline.py

@@ -134,7 +134,7 @@ def start(
             print(f"Error during pipeline execution: {e}")
             raise e
 
-    def fit(self, x: XYData, y: Optional[XYData]) -> None | float:
+    def fit(self, x: XYData, y: Optional[XYData]) -> None | float | dict:
         """
         Fit the pipeline to the input data.
 

framework3/plugins/splitter/cross_validation_splitter.py

@@ -103,7 +103,7 @@ def split(self, pipeline: BaseFilter):
         self.pipeline = pipeline
         self.pipeline.verbose(False)
 
-    def fit(self, x: XYData, y: XYData | None) -> Optional[float]:
+    def fit(self, x: XYData, y: XYData | None) -> Optional[float | dict]:
         """
         Perform K-Fold cross-validation on the given data.
 
@@ -131,7 +131,7 @@ def fit(self, x: XYData, y: XYData | None) -> Optional[float]:
         if self.pipeline is None:
             raise ValueError("Pipeline must be fitted before splitting")
 
-        losses = []
+        losses: dict = {}
         splits = self._kfold.split(X)
         for train_idx, val_idx in tqdm(
             splits, total=self._kfold.get_n_splits(X), disable=not self._verbose
@@ -151,11 +151,14 @@ def fit(self, x: XYData, y: XYData | None) -> Optional[float]:
             _y = pipeline.predict(X_val)
 
             loss = pipeline.evaluate(X_val, y_val, _y)
-            losses.append(float(next(iter(loss.values()))))
+            for metric, value in loss.items():
+                v = losses.get(metric, [])
+                v.append(value)
+                losses[metric] = v
 
             self.clear_memory()
 
-        return float(np.mean(losses) if losses else 0.0)
+        return dict(map(lambda item: (item[0], np.mean(item[1])), losses.items()))
 
     def start(
         self, x: XYData, y: Optional[XYData], X_: Optional[XYData]

framework3/plugins/splitter/stratified_cross_validation_splitter.py

@@ -84,7 +84,7 @@ def split(self, pipeline: BaseFilter):
         self.pipeline = pipeline
         self.pipeline.verbose(False)
 
-    def fit(self, x: XYData, y: XYData | None) -> Optional[float]:
+    def fit(self, x: XYData, y: XYData | None) -> Optional[float | dict]:
         """
         Perform Stratified K-Fold cross-validation on the given data.
 
@@ -111,7 +111,7 @@ def fit(self, x: XYData, y: XYData | None) -> Optional[float]:
         X = x.value
         Y = y.value
 
-        losses = []
+        losses: dict = {}
         splits = self._skf.split(X, Y)
         for train_idx, val_idx in tqdm(
             splits, total=self._skf.get_n_splits(X, Y), disable=not self._verbose
@@ -131,11 +131,14 @@ def fit(self, x: XYData, y: XYData | None) -> Optional[float]:
             _y = pipeline.predict(X_val)
 
             loss = pipeline.evaluate(X_val, y_val, _y)
-            losses.append(float(next(iter(loss.values()))))
+            for metric, value in loss.items():
+                v = losses.get(metric, [])
+                v.append(value)
+                losses[metric] = v
 
             self.clear_memory()
 
-        return float(np.mean(losses) if losses else 0.0)
+        return dict(map(lambda item: (item[0], np.mean(item[1])), losses.items()))
 
     def start(
         self, x: XYData, y: Optional[XYData], X_: Optional[XYData]