move automlx explaination to subclass

Author: codeloop

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

@@ -370,3 +370,75 @@ def _custom_predict_automlx(self, data):
         return self.models.get(self.series_id).forecast(
             X=data_temp, periods=data_temp.shape[0]
         )[self.series_id]
+
+    @runtime_dependency(
+        module="automlx",
+        err_msg=(
+            "Please run `python3 -m pip install automlx` to install the required dependencies for model explanation."
+        ),
+    )
+    def explain_model(self):
+        """
+        Generates explanations for the model using the AutoMLx library.
+
+        Parameters
+        ----------
+        None
+
+        Returns
+        -------
+        None
+
+        Notes
+        -----
+        This function works by generating local explanations for each series in the dataset.
+        It uses the ``MLExplainer`` class from the AutoMLx library to generate feature attributions
+        for each series. The feature attributions are then stored in the ``self.local_explanation`` dictionary.
+
+        If the accuracy mode is set to AutoMLX, it uses the AutoMLx library to generate explanations.
+        Otherwise, it falls back to the default explanation generation method.
+        """
+        import automlx
+
+        # Loop through each series in the dataset
+        for s_id, data_i in self.datasets.get_data_by_series(
+            include_horizon=False
+        ).items():
+            if self.spec.explanations_accuracy_mode == SpeedAccuracyMode.AUTOMLX:
+                # Use the MLExplainer class from AutoMLx to generate explanations
+                explainer = automlx.MLExplainer(
+                    self.models[s_id],
+                    self.datasets.additional_data.get_data_for_series(series_id=s_id)
+                    .drop(self.spec.datetime_column.name, axis=1)
+                    .head(-self.spec.horizon)
+                    if self.spec.additional_data
+                    else None,
+                    pd.DataFrame(data_i[self.spec.target_column]),
+                    task="forecasting",
+                )
+
+                # Generate explanations for the forecast
+                explanations = explainer.explain_prediction(
+                    X=self.datasets.additional_data.get_data_for_series(series_id=s_id)
+                    .drop(self.spec.datetime_column.name, axis=1)
+                    .tail(self.spec.horizon)
+                    if self.spec.additional_data
+                    else None,
+                    forecast_timepoints=list(range(self.spec.horizon + 1)),
+                )
+
+                # Convert the explanations to a DataFrame
+                explanations_df = pd.concat(
+                    [exp.to_dataframe() for exp in explanations]
+                )
+                explanations_df["row"] = explanations_df.groupby("Feature").cumcount()
+                explanations_df = explanations_df.pivot(
+                    index="row", columns="Feature", values="Attribution"
+                )
+                explanations_df = explanations_df.reset_index(drop=True)
+
+                # Store the explanations in the local_explanation dictionary
+                self.local_explanation[s_id] = explanations_df
+            else:
+                # Fall back to the default explanation generation method
+                super().explain_model()

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

@@ -649,18 +649,19 @@ def _save_model(self, output_dir, storage_options):
             storage_options=storage_options,
         )
 
+    def _validate_automlx_explanation_mode(self):
+        if self.spec.model != SupportedModels.AutoMLX and self.spec.explanations_accuracy_mode == SpeedAccuracyMode.AUTOMLX:
+            raise ValueError(
+                "AUTOMLX explanation accuracy mode is only supported for AutoMLX models. "
+                "Please select mode other than AUTOMLX from the available explanations_accuracy_mode options"
+            )
+
     @runtime_dependency(
         module="shap",
         err_msg=(
             "Please run `python3 -m pip install shap` to install the required dependencies for model explanation."
         ),
     )
-    @runtime_dependency(
-        module="automlx",
-        err_msg=(
-            "Please run `python3 -m pip install automlx` to install the required dependencies for model explanation."
-        ),
-    )
     def explain_model(self):
         """
         Generates an explanation for the model by using the SHAP (Shapley Additive exPlanations) library.
@@ -683,53 +684,13 @@ def explain_model(self):
         )
         ratio = SpeedAccuracyMode.ratio[self.spec.explanations_accuracy_mode]
 
+        # validate the automlx mode is use for automlx model
+        self._validate_automlx_explanation_mode()
+
         for s_id, data_i in self.datasets.get_data_by_series(
             include_horizon=False
         ).items():
-            if (
-                self.spec.model == SupportedModels.AutoMLX
-                and self.spec.explanations_accuracy_mode == SpeedAccuracyMode.AUTOMLX
-            ):
-                import automlx
-
-                explainer = automlx.MLExplainer(
-                    self.models[s_id],
-                    self.datasets.additional_data.get_data_for_series(series_id=s_id)
-                    .drop(self.spec.datetime_column.name, axis=1)
-                    .head(-self.spec.horizon)
-                    if self.spec.additional_data
-                    else None,
-                    pd.DataFrame(data_i[self.spec.target_column]),
-                    task="forecasting",
-                )
-
-                explanations = explainer.explain_prediction(
-                    X=self.datasets.additional_data.get_data_for_series(series_id=s_id)
-                    .drop(self.spec.datetime_column.name, axis=1)
-                    .tail(self.spec.horizon)
-                    if self.spec.additional_data
-                    else None,
-                    forecast_timepoints=list(range(self.spec.horizon + 1)),
-                )
-
-                explanations_df = pd.concat(
-                    [exp.to_dataframe() for exp in explanations]
-                )
-                explanations_df["row"] = explanations_df.groupby("Feature").cumcount()
-                explanations_df = explanations_df.pivot(
-                    index="row", columns="Feature", values="Attribution"
-                )
-                explanations_df = explanations_df.reset_index(drop=True)
-                self.local_explanation[s_id] = explanations_df
-            elif (
-                self.spec.explanations_accuracy_mode == SpeedAccuracyMode.AUTOMLX
-                and self.spec.model != SupportedModels.AutoMLX
-            ):
-                raise ValueError(
-                    "AUTOMLX explanation accuracy mode is only supported for AutoMLX models. "
-                    "Please select mode other than AUTOMLX from the available explanations_accuracy_mode options"
-                )
-            elif s_id in self.models:
+            if s_id in self.models:
                 explain_predict_fn = self.get_explain_predict_fn(series_id=s_id)
                 data_trimmed = data_i.tail(
                     max(int(len(data_i) * ratio), 5)