small changes

Author: govarsha

ads/opctl/operator/lowcode/forecast/model/arima.py

@@ -127,8 +127,9 @@ def _build_model(self) -> pd.DataFrame:
             outputs[target] = forecast
 
             params = vars(model).copy()
-            params.pop("arima_res_")
-            params.pop("endog_index_")
+            for param in ['arima_res', 'endog_index_']:
+                if param in params:
+                    params.pop(param)
             self.model_parameters[target] = {
                 "framework": SupportedModels.Arima,
                 **params,
@@ -195,7 +196,7 @@ def _generate_report(self):
                 global_explanation_df = pd.DataFrame(self.global_explanation)
 
                 self.formatted_global_explanation = (
-                    global_explanation_df / global_explanation_df.sum(axis=0) * 100
+                        global_explanation_df / global_explanation_df.sum(axis=0) * 100
                 )
 
                 # Create a markdown section for the global explainability

ads/opctl/operator/lowcode/forecast/model/autots.py

@@ -218,7 +218,8 @@ def _build_model(self) -> pd.DataFrame:
                 "validation_indexes",
                 "best_model",
             ]:
-                params.pop(param)
+                if param in params:
+                    params.pop(param)
             self.model_parameters[cat_target] = {
                 "framework": SupportedModels.AutoTS,
                 **params,

ads/opctl/operator/lowcode/forecast/model/base_model.py

@@ -56,9 +56,6 @@ def __init__(self, config: ForecastOperatorConfig, datasets: ForecastDatasets):
 
         # these fields are populated in the _build_model() method
         self.models = None
-        # Relevant only for neuralprophet
-        self.loaded_trainers = None
-        self.trainers = None
 
         # "outputs" is a list of outputs generated by the models. These should only be generated when the framework requires the original output for plotting
         self.outputs = None
@@ -559,13 +556,9 @@ def _save_report(
                 output_dir=output_dir,
                 storage_options=storage_options,
             )
-            if self.trainers is not None:
-                utils.write_pkl(
-                    obj=self.trainers,
-                    filename="trainer.pkl",
-                    output_dir=output_dir,
-                    storage_options=storage_options,
-                )
+
+        self._save_model_specific_files(output_dir, storage_options)
+
         logger.info(
             f"The outputs have been successfully "
             f"generated and placed into the directory: {output_dir}."
@@ -599,6 +592,13 @@ def _generate_train_metrics(self) -> pd.DataFrame:
         """
         raise NotImplementedError
 
+    def _save_model_specific_files(self, output_dir, storage_options):
+        """
+        The method that needs to be implemented on the particular model level
+        """
+        pass
+
+
     @runtime_dependency(
         module="shap",
         err_msg=(

ads/opctl/operator/lowcode/forecast/model/neuralprophet.py

@@ -70,6 +70,8 @@ def __init__(self, config: ForecastOperatorConfig, datasets: ForecastDatasets):
         super().__init__(config=config, datasets=datasets)
         self.train_metrics = True
         self.forecast_col_name = "yhat1"
+        self.loaded_trainers = None
+        self.trainers = None
 
     def _build_model(self) -> pd.DataFrame:
         from neuralprophet import NeuralProphet
@@ -319,6 +321,15 @@ def objective(trial):
 
         return output_col
 
+    def _save_model_specific_files(self, output_dir, storage_options):
+        if self.spec.generate_model_pickle and self.trainers is not None:
+            utils.write_pkl(
+                obj=self.trainers,
+                filename="trainer.pkl",
+                output_dir=output_dir,
+                storage_options=storage_options,
+            )
+
     def _generate_report(self):
         import datapane as dp
 

ads/opctl/operator/lowcode/forecast/model/prophet.py

@@ -213,7 +213,8 @@ def objective(trial):
 
             params = vars(model).copy()
             for param in ["history", "history_dates", "stan_fit"]:
-                params.pop(param)
+                if param in params:
+                    params.pop(param)
             self.model_parameters[target] = {
                 "framework": SupportedModels.Prophet,
                 **params,