refactor metrics

tryolabs · Mar 25, 2024 · c32b1ab · c32b1ab
1 parent 38279c0
commit c32b1ab
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/pipeline_lib/core/steps/calculate_metrics.py b/pipeline_lib/core/steps/calculate_metrics.py
@@ -19,15 +19,14 @@ def __init__(self, mape_threshold: float = 0.01) -> None:
 
     def execute(self, data: DataContainer) -> DataContainer:
         self.logger.debug("Starting metric calculation")
-        model_output = data.flow
 
         target_column_name = data.target
 
         if target_column_name is None:
-            raise ValueError("Target column not found on any configuration.")
+            raise ValueError("Target column nsot found on any configuration.")
 
-        true_values = model_output[target_column_name]
-        predictions = model_output["predictions"]
+        true_values = data.flow[target_column_name]
+        predictions = data.predictions
 
         mae = mean_absolute_error(true_values, predictions)
         rmse = np.sqrt(mean_squared_error(true_values, predictions))