small changes

Author: govarsha

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

@@ -79,10 +79,7 @@ def generate_report(self):
 
             # load models if given
             if self.spec.previous_output_dir is not None:
-                try:
-                    self.loaded_models = utils.load_pkl(self.spec.previous_output_dir + "/model.pkl")
-                except:
-                    logger.info("model.pkl is not present")
+                self._load_model()
 
             start_time = time.time()
             result_df = self._build_model()
@@ -592,6 +589,12 @@ def _generate_train_metrics(self) -> pd.DataFrame:
         """
         raise NotImplementedError
 
+    def _load_model(self):
+        try:
+            self.loaded_models = utils.load_pkl(self.spec.previous_output_dir + "/model.pkl")
+        except:
+            logger.info("model.pkl is not present")
+
     def _save_model_specific_files(self, output_dir, storage_options):
         """
         The method that needs to be implemented on the particular model level

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

@@ -73,19 +73,21 @@ def __init__(self, config: ForecastOperatorConfig, datasets: ForecastDatasets):
         self.loaded_trainers = None
         self.trainers = None
 
+    def _load_model(self):
+        try:
+            self.loaded_models = utils.load_pkl(self.spec.previous_output_dir + "/model.pkl")
+            self.loaded_trainers = utils.load_pkl(self.spec.previous_output_dir + "/trainer.pkl")
+        except:
+            logger.info("model.pkl/trainer.pkl is not present")
+
+
     def _build_model(self) -> pd.DataFrame:
         from neuralprophet import NeuralProphet
 
         full_data_dict = self.datasets.full_data_dict
         models = []
         trainers = []
 
-        if self.loaded_models is not None:
-            try:
-                self.loaded_trainers = utils.load_pkl(self.spec.previous_output_dir + "/trainer.pkl")
-            except:
-                logger.info("trainer.pkl is not present")
-
         outputs = dict()
         outputs_legacy = []