oracle-ads 2.10.0__py3-none-any.whl → 2.11.0__py3-none-any.whl

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

@@ -4,6 +4,7 @@
 # Copyright (c) 2023 Oracle and/or its affiliates.
 # Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/
 
+import json
 import os
 import tempfile
 import time
@@ -15,15 +16,39 @@ import fsspec
 import numpy as np
 import pandas as pd
 
-from ads.opctl.operator.lowcode.forecast.utils import default_signer
+from ads.opctl.operator.lowcode.forecast.utils import (
+    default_signer,
+    evaluate_train_metrics,
+    get_forecast_plots,
+    _build_metrics_df,
+    _build_metrics_per_horizon,
+    load_pkl,
+    write_pkl,
+    _label_encode_dataframe,
+)
 from ads.common.object_storage_details import ObjectStorageDetails
 from ads.opctl import logger
 
-from .. import utils
-from ..const import SUMMARY_METRICS_HORIZON_LIMIT, SupportedMetrics, SupportedModels
+from ads.opctl.operator.lowcode.common.utils import (
+    human_time_friendly,
+    enable_print,
+    disable_print,
+    write_data,
+    merged_category_column_name,
+    datetime_to_seconds,
+    seconds_to_datetime,
+    find_output_dirname,
+)
+from ..const import (
+    SUMMARY_METRICS_HORIZON_LIMIT,
+    SupportedMetrics,
+    SupportedModels,
+    SpeedAccuracyMode,
+)
 from ..operator_config import ForecastOperatorConfig, ForecastOperatorSpec
 from ads.common.decorator.runtime_dependency import runtime_dependency
 from .forecast_datasets import ForecastDatasets, ForecastOutput
+from ads.opctl.operator.lowcode.forecast.model.forecast_datasets import TestData
 
 
 class ForecastOperatorBaseModel(ABC):
@@ -41,23 +66,27 @@ class ForecastOperatorBaseModel(ABC):
         self.spec: ForecastOperatorSpec = config.spec
         self.datasets: ForecastDatasets = datasets
 
-        self.original_user_data = datasets.original_user_data
-        self.original_total_data = datasets.original_total_data
-        self.original_additional_data = datasets.original_additional_data
-        self.full_data_dict = datasets.full_data_dict
-        self.target_columns = datasets.target_columns
-        self.categories = datasets.categories
+        self.full_data_dict = datasets.get_data_by_series()
 
         self.test_eval_metrics = None
         self.original_target_column = self.spec.target_column
+        self.dt_column_name = self.spec.datetime_column.name
+
+        self.model_parameters = dict()
+        self.loaded_models = None
 
         # these fields are populated in the _build_model() method
         self.models = 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
         self.forecast_output = None
+        self.errors_dict = dict()
+        self.le = dict()
+
+        self.formatted_global_explanation = None
+        self.formatted_local_explanation = None
 
-        self.train_metrics = False
         self.forecast_col_name = "yhat"
         self.perform_tuning = self.spec.tuning != None
 
@@ -73,10 +102,14 @@ class ForecastOperatorBaseModel(ABC):
             warnings.simplefilter(action="ignore", category=ConvergenceWarning)
             import datapane as dp
 
-            # load data and build models
+            # load models if given
+            if self.spec.previous_output_dir is not None:
+                self._load_model()
+
             start_time = time.time()
             result_df = self._build_model()
             elapsed_time = time.time() - start_time
+            logger.info("Building the models completed in %s seconds", elapsed_time)
 
             # Generate metrics
             summary_metrics = None
@@ -84,21 +117,7 @@ class ForecastOperatorBaseModel(ABC):
             self.eval_metrics = None
 
             if self.spec.generate_report or self.spec.generate_metrics:
-                if self.train_metrics:
-                    self.eval_metrics = utils.evaluate_train_metrics(
-                        self.target_columns,
-                        self.datasets,
-                        self.forecast_output,
-                        self.spec.datetime_column.name,
-                        target_col=self.forecast_col_name,
-                    )
-                else:
-                    try:
-                        self.eval_metrics = self._generate_train_metrics()
-                    except NotImplementedError:
-                        logger.warn(
-                            f"Training Metrics are not available for model type {self.spec.model}"
-                        )
+                self.eval_metrics = self.generate_train_metrics()
 
                 if self.spec.test_data:
                     try:
@@ -107,10 +126,6 @@ class ForecastOperatorBaseModel(ABC):
                             summary_metrics,
                             test_data,
                         ) = self._test_evaluate_metrics(
-                            target_columns=self.target_columns,
-                            test_filename=self.spec.test_data.url,
-                            output=self.forecast_output,
-                            target_col=self.forecast_col_name,
                             elapsed_time=elapsed_time,
                         )
                     except Exception as e:
@@ -125,94 +140,114 @@ class ForecastOperatorBaseModel(ABC):
                     other_sections,
                 ) = self._generate_report()
 
-                ds_column_series = self.datasets.get_longest_datetime_column()
-
                 title_text = dp.Text("# Forecast Report")
 
-                md_columns = " * ".join([f"{x} \n" for x in self.target_columns])
+                md_columns = " * ".join(
+                    [f"{s_id} \n" for s_id in self.datasets.list_series_ids()]
+                )
+
+                header_section = dp.Blocks(
+                    blocks=[
+                        dp.Text(f"You selected the **`{self.spec.model}`** model."),
+                        model_description,
+                        dp.Text(
+                            "Based on your dataset, you could have also selected "
+                            f"any of the models: `{'`, `'.join(SupportedModels.keys())}`."
+                        ),
+                        dp.Group(
+                            dp.BigNumber(
+                                heading="Analysis was completed in ",
+                                value=human_time_friendly(elapsed_time),
+                            ),
+                            dp.BigNumber(
+                                heading="Starting time index",
+                                value=self.datasets.get_earliest_timestamp().strftime(
+                                    "%B %d, %Y"
+                                ),
+                            ),
+                            dp.BigNumber(
+                                heading="Ending time index",
+                                value=self.datasets.get_latest_timestamp().strftime(
+                                    "%B %d, %Y"
+                                ),
+                            ),
+                            dp.BigNumber(
+                                heading="Num series",
+                                value=len(self.datasets.list_series_ids()),
+                            ),
+                            columns=4,
+                        ),
+                    ]
+                )
+
                 first_10_rows_blocks = [
                     dp.DataTable(
-                        df.head(10).rename({col: self.spec.target_column}, axis=1),
+                        df.head(10),
                         caption="Start",
-                        label=col,
+                        label=s_id,
                     )
-                    for col, df in self.full_data_dict.items()
+                    for s_id, df in self.full_data_dict.items()
                 ]
 
                 last_10_rows_blocks = [
                     dp.DataTable(
-                        df.tail(10).rename({col: self.spec.target_column}, axis=1),
+                        df.tail(10),
                         caption="End",
-                        label=col,
+                        label=s_id,
                     )
-                    for col, df in self.full_data_dict.items()
+                    for s_id, df in self.full_data_dict.items()
                 ]
 
                 data_summary_blocks = [
                     dp.DataTable(
-                        df.rename({col: self.spec.target_column}, axis=1).describe(),
+                        df.describe(),
                         caption="Summary Statistics",
-                        label=col,
+                        label=s_id,
                     )
-                    for col, df in self.full_data_dict.items()
+                    for s_id, df in self.full_data_dict.items()
                 ]
-                summary = dp.Blocks(
-                    dp.Select(
+
+                series_name = merged_category_column_name(
+                    self.spec.target_category_columns
+                )
+                series_subtext = dp.Text(f"Indexed by {series_name}")
+                first_10_title = dp.Text("### First 10 Rows of Data")
+                last_10_title = dp.Text("### Last 10 Rows of Data")
+                summary_title = dp.Text("### Data Summary Statistics")
+
+                if series_name is not None and len(self.datasets.list_series_ids()) > 1:
+                    data_summary_sec = dp.Blocks(
                         blocks=[
-                            dp.Group(
-                                dp.Text(
-                                    f"You selected the **`{self.spec.model}`** model."
-                                ),
-                                model_description,
-                                dp.Text(
-                                    "Based on your dataset, you could have also selected "
-                                    f"any of the models: `{'`, `'.join(SupportedModels.keys())}`."
-                                ),
-                                dp.Group(
-                                    dp.BigNumber(
-                                        heading="Analysis was completed in ",
-                                        value=utils.human_time_friendly(elapsed_time),
-                                    ),
-                                    dp.BigNumber(
-                                        heading="Starting time index",
-                                        value=ds_column_series.min().strftime(
-                                            "%B %d, %Y"
-                                        ),
-                                    ),
-                                    dp.BigNumber(
-                                        heading="Ending time index",
-                                        value=ds_column_series.max().strftime(
-                                            "%B %d, %Y"
-                                        ),
-                                    ),
-                                    dp.BigNumber(
-                                        heading="Num series",
-                                        value=len(self.target_columns),
-                                    ),
-                                    columns=4,
-                                ),
-                                dp.Text("### First 10 Rows of Data"),
-                                dp.Select(blocks=first_10_rows_blocks)
-                                if len(first_10_rows_blocks) > 1
-                                else first_10_rows_blocks[0],
-                                dp.Text("----"),
-                                dp.Text("### Last 10 Rows of Data"),
-                                dp.Select(blocks=last_10_rows_blocks)
-                                if len(last_10_rows_blocks) > 1
-                                else last_10_rows_blocks[0],
-                                dp.Text("### Data Summary Statistics"),
-                                dp.Select(blocks=data_summary_blocks)
-                                if len(data_summary_blocks) > 1
-                                else data_summary_blocks[0],
-                                label="Summary",
-                            ),
-                            dp.Text(
-                                "The following report compares a variety of metrics and plots "
-                                f"for your target columns: \n {md_columns}.\n",
-                                label="Target Columns",
-                            ),
+                            first_10_title,
+                            series_subtext,
+                            dp.Select(blocks=first_10_rows_blocks),
+                            last_10_title,
+                            series_subtext,
+                            dp.Select(blocks=last_10_rows_blocks),
+                            summary_title,
+                            series_subtext,
+                            dp.Select(blocks=data_summary_blocks),
+                            dp.Text("----"),
                         ]
-                    ),
+                    )
+                else:
+                    data_summary_sec = dp.Blocks(
+                        blocks=[
+                            first_10_title,
+                            first_10_rows_blocks[0],
+                            last_10_title,
+                            last_10_rows_blocks[0],
+                            summary_title,
+                            data_summary_blocks[0],
+                            dp.Text("----"),
+                        ]
+                    )
+
+                summary = dp.Group(
+                    blocks=[
+                        header_section,
+                        data_summary_sec,
+                    ]
                 )
 
                 test_metrics_sections = []
@@ -236,19 +271,22 @@ class ForecastOperatorBaseModel(ABC):
                     train_metrics_sections = [sec9_text, sec9]
 
                 forecast_text = dp.Text(f"## Forecasted Data Overlaying Historical")
-                forecast_sec = utils.get_forecast_plots(
+                forecast_sec = get_forecast_plots(
                     self.forecast_output,
-                    self.target_columns,
                     horizon=self.spec.horizon,
                     test_data=test_data,
                     ci_interval_width=self.spec.confidence_interval_width,
                 )
-                forecast_plots = [forecast_text, forecast_sec]
+                if series_name is not None and len(self.datasets.list_series_ids()) > 1:
+                    forecast_plots = [forecast_text, series_subtext, forecast_sec]
+                else:
+                    forecast_plots = [forecast_text, forecast_sec]
 
                 yaml_appendix_title = dp.Text(f"## Reference: YAML File")
                 yaml_appendix = dp.Code(code=self.config.to_yaml(), language="yaml")
                 report_sections = (
-                    [title_text, summary]
+                    [title_text]
+                    + [summary]
                     + forecast_plots
                     + other_sections
                     + test_metrics_sections
@@ -264,69 +302,44 @@ class ForecastOperatorBaseModel(ABC):
                 test_metrics_df=self.test_eval_metrics,
             )
 
-    def _test_evaluate_metrics(
-        self, target_columns, test_filename, output, target_col="yhat", elapsed_time=0
-    ):
+    def _test_evaluate_metrics(self, elapsed_time=0):
         total_metrics = pd.DataFrame()
         summary_metrics = pd.DataFrame()
-        data = None
-        try:
-            storage_options = (
-                default_signer()
-                if ObjectStorageDetails.is_oci_path(test_filename)
-                else {}
-            )
-            data = utils._load_data(
-                filename=test_filename,
-                format=self.spec.test_data.format,
-                storage_options=storage_options,
-                columns=self.spec.test_data.columns,
-            )
-        except pd.errors.EmptyDataError:
-            logger.warn("Empty testdata file")
-            return total_metrics, summary_metrics, None
-
-        if data.empty:
-            return total_metrics, summary_metrics, None
+        data = TestData(self.spec)
 
-        data = self._preprocess(
-            data, self.spec.datetime_column.name, self.spec.datetime_column.format
-        )
-        data, confirm_targ_columns = utils._clean_data(
-            data=data,
-            target_column=self.original_target_column,
-            target_category_columns=self.spec.target_category_columns,
-            datetime_column="ds",
-        )
-
-        # Calculating Test Metrics
-        for cat in self.forecast_output.list_categories():
-            target_column_i = self.forecast_output.category_to_target[cat]
-            output_forecast_i = self.forecast_output.get_category(cat)
-            # Only columns present in test file will be used to generate test error
-            if target_column_i in data:
-                # Assuming that predictions have all forecast values
-                dates = output_forecast_i["Date"]
-                # Filling zeros for any date missing in test data to maintain consistency in metric calculation as in all other missing values cases it comes as 0
-                y_true = [
-                    data.loc[data["ds"] == date, target_column_i].values[0]
-                    if date in data["ds"].values
-                    else 0
-                    for date in dates
+        # Generate y_pred and y_true for each series
+        for s_id in self.forecast_output.list_series_ids():
+            try:
+                y_true = data.get_data_for_series(s_id)[data.target_name].values[
+                    -self.spec.horizon :
                 ]
-                y_pred_i = output_forecast_i["forecast_value"].values
-                y_pred = np.asarray(y_pred_i[-len(y_true) :])
-
-                metrics_df = utils._build_metrics_df(
-                    y_true=y_true[-self.spec.horizon :],
-                    y_pred=y_pred[-self.spec.horizon :],
-                    column_name=target_column_i,
-                )
-                total_metrics = pd.concat([total_metrics, metrics_df], axis=1)
-            else:
+            except KeyError as ke:
                 logger.warn(
-                    f"Error Generating Metrics: Unable to find {target_column_i} in the test data."
+                    f"Error Generating Metrics: Unable to find {s_id} in the test data. Error: {ke.args}"
                 )
+            y_pred = self.forecast_output.get_forecast(s_id)["forecast_value"].values[
+                -self.spec.horizon :
+            ]
+
+            drop_na_mask = ~np.isnan(y_true) & ~np.isnan(y_pred)
+            if not drop_na_mask.all():  # There is a missing value
+                if drop_na_mask.any():  # All values are missing
+                    logger.debug(
+                        f"No values in the test data for series: {s_id}. This will affect the test metrics."
+                    )
+                    continue
+                logger.debug(
+                    f"Missing values in the test data for series: {s_id}. This will affect the test metrics."
+                )
+                y_true = y_true[drop_na_mask]
+                y_pred = y_pred[drop_na_mask]
+
+            metrics_df = _build_metrics_df(
+                y_true=y_true,
+                y_pred=y_pred,
+                series_id=s_id,
+            )
+            total_metrics = pd.concat([total_metrics, metrics_df], axis=1)
 
         if total_metrics.empty:
             return total_metrics, summary_metrics, data
@@ -370,20 +383,10 @@ class ForecastOperatorBaseModel(ABC):
 
         """Calculates Mean sMAPE, Median sMAPE, Mean MAPE, Median MAPE, Mean wMAPE, Median wMAPE values for each horizon
         if horizon <= 10."""
-        target_columns_in_output = set(target_columns).intersection(data.columns)
         if self.spec.horizon <= SUMMARY_METRICS_HORIZON_LIMIT:
-            if set(self.forecast_output.list_target_category_columns()) != set(
-                target_columns_in_output
-            ):
-                logger.warn(
-                    f"Column Mismatch between Forecast Output and Target Columns"
-                )
-            metrics_per_horizon = utils._build_metrics_per_horizon(
-                data=data,
+            metrics_per_horizon = _build_metrics_per_horizon(
+                test_data=data,
                 output=self.forecast_output,
-                target_columns=target_columns,
-                target_col=target_col,
-                horizon_periods=self.spec.horizon,
             )
             if not metrics_per_horizon.empty:
                 summary_metrics = pd.concat([summary_metrics, metrics_per_horizon])
@@ -417,17 +420,9 @@ class ForecastOperatorBaseModel(ABC):
         """Saves resulting reports to the given folder."""
         import datapane as dp
 
-        if self.spec.output_directory:
-            output_dir = self.spec.output_directory.url
-        else:
-            output_dir = "tmp_fc_operator_result"
-            logger.warn(
-                "Since the output directory was not specified, the output will be saved to {} directory.".format(
-                    output_dir
-                )
-            )
+        unique_output_dir = find_output_dirname(self.spec.output_directory)
 
-        if ObjectStorageDetails.is_oci_path(output_dir):
+        if ObjectStorageDetails.is_oci_path(unique_output_dir):
             storage_options = default_signer()
         else:
             storage_options = dict()
@@ -437,11 +432,11 @@ class ForecastOperatorBaseModel(ABC):
             # datapane html report
             with tempfile.TemporaryDirectory() as temp_dir:
                 report_local_path = os.path.join(temp_dir, "___report.html")
-                utils.block_print()
+                disable_print()
                 dp.save_report(report_sections, report_local_path)
-                utils.enable_print()
+                enable_print()
 
-                report_path = os.path.join(output_dir, self.spec.report_filename)
+                report_path = os.path.join(unique_output_dir, self.spec.report_filename)
                 with open(report_local_path) as f1:
                     with fsspec.open(
                         report_path,
@@ -451,19 +446,28 @@ class ForecastOperatorBaseModel(ABC):
                         f2.write(f1.read())
 
         # forecast csv report
-        utils._write_data(
+        write_data(
             data=result_df,
-            filename=os.path.join(output_dir, self.spec.forecast_filename),
+            filename=os.path.join(unique_output_dir, self.spec.forecast_filename),
             format="csv",
             storage_options=storage_options,
         )
 
         # metrics csv report
         if self.spec.generate_metrics:
+            metrics_col_name = (
+                self.original_target_column
+                if self.datasets.has_artificial_series()
+                else "Series 1"
+            )
             if metrics_df is not None:
-                utils._write_data(
-                    data=metrics_df.rename_axis("metrics").reset_index(),
-                    filename=os.path.join(output_dir, self.spec.metrics_filename),
+                write_data(
+                    data=metrics_df.reset_index().rename(
+                        {"index": "metrics", "Series 1": metrics_col_name}, axis=1
+                    ),
+                    filename=os.path.join(
+                        unique_output_dir, self.spec.metrics_filename
+                    ),
                     format="csv",
                     storage_options=storage_options,
                     index=False,
@@ -476,10 +480,12 @@ class ForecastOperatorBaseModel(ABC):
             # test_metrics csv report
             if self.spec.test_data is not None:
                 if test_metrics_df is not None:
-                    utils._write_data(
-                        data=test_metrics_df.rename_axis("metrics").reset_index(),
+                    write_data(
+                        data=test_metrics_df.reset_index().rename(
+                            {"index": "metrics", "Series 1": metrics_col_name}, axis=1
+                        ),
                         filename=os.path.join(
-                            output_dir, self.spec.test_metrics_filename
+                            unique_output_dir, self.spec.test_metrics_filename
                         ),
                         format="csv",
                         storage_options=storage_options,
@@ -493,10 +499,10 @@ class ForecastOperatorBaseModel(ABC):
         if self.spec.generate_explanations:
             try:
                 if self.formatted_global_explanation is not None:
-                    utils._write_data(
+                    write_data(
                         data=self.formatted_global_explanation,
                         filename=os.path.join(
-                            output_dir, self.spec.global_explanation_filename
+                            unique_output_dir, self.spec.global_explanation_filename
                         ),
                         format="csv",
                         storage_options=storage_options,
@@ -508,10 +514,10 @@ class ForecastOperatorBaseModel(ABC):
                     )
 
                 if self.formatted_local_explanation is not None:
-                    utils._write_data(
+                    write_data(
                         data=self.formatted_local_explanation,
                         filename=os.path.join(
-                            output_dir, self.spec.local_explanation_filename
+                            unique_output_dir, self.spec.local_explanation_filename
                         ),
                         format="csv",
                         storage_options=storage_options,
@@ -525,17 +531,52 @@ class ForecastOperatorBaseModel(ABC):
                 logger.warn(
                     "Unable to generate explanations for this model type or for this dataset."
                 )
+                logger.debug(f"Got error: {e.args}")
+
+        if self.spec.generate_model_parameters:
+            # model params
+            write_data(
+                data=pd.DataFrame.from_dict(self.model_parameters),
+                filename=os.path.join(unique_output_dir, "model_params.json"),
+                format="json",
+                storage_options=storage_options,
+                index=True,
+                indent=4,
+            )
+
+        # model pickle
+        if self.spec.generate_model_pickle:
+            self._save_model(unique_output_dir, storage_options)
+
         logger.info(
             f"The outputs have been successfully "
-            f"generated and placed into the directory: {output_dir}."
+            f"generated and placed into the directory: {unique_output_dir}."
         )
+        print(
+            f"The outputs have been successfully generated and placed into the directory: {unique_output_dir}."
+        )
+        if self.errors_dict:
+            write_data(
+                data=pd.DataFrame(self.errors_dict.items(), columns=["model", "error"]),
+                filename=os.path.join(
+                    unique_output_dir, self.spec.errors_dict_filename
+                ),
+                format="csv",
+                storage_options=storage_options,
+                index=True,
+            )
+        else:
+            logger.info(f"All modeling completed successfully.")
 
-    def _preprocess(self, data, ds_column, datetime_format):
+    def preprocess(self, df, series_id):
         """The method that needs to be implemented on the particular model level."""
-        data["ds"] = pd.to_datetime(data[ds_column], format=datetime_format)
-        if ds_column != "ds":
-            data.drop([ds_column], axis=1, inplace=True)
-        return data
+        data = df.rename(
+            {self.dt_column_name: "ds", self.original_target_column: "y"}, axis=1
+        )
+        self.le[series_id], df_encoded = _label_encode_dataframe(
+            data, no_encode={"ds", "y"}
+        )
+        return df_encoded
 
     @abstractmethod
     def _generate_report(self):
@@ -551,20 +592,40 @@ class ForecastOperatorBaseModel(ABC):
         The method that needs to be implemented on the particular model level.
         """
 
-    def _generate_train_metrics(self) -> pd.DataFrame:
+    def drop_horizon(self, df: pd.DataFrame) -> pd.DataFrame:
+        return df.iloc[: -self.spec.horizon]
+
+    def get_horizon(self, df: pd.DataFrame) -> pd.DataFrame:
+        return df.iloc[-self.spec.horizon :]
+
+    def generate_train_metrics(self) -> pd.DataFrame:
         """
         Generate Training Metrics when fitted data is not available.
         The method that needs to be implemented on the particular model level.
         """
-        raise NotImplementedError
+        return evaluate_train_metrics(self.forecast_output)
+
+    def _load_model(self):
+        try:
+            self.loaded_models = load_pkl(self.spec.previous_output_dir + "/model.pkl")
+        except:
+            logger.info("model.pkl is not present")
+
+    def _save_model(self, output_dir, storage_options):
+        write_pkl(
+            obj=self.models,
+            filename="model.pkl",
+            output_dir=output_dir,
+            storage_options=storage_options,
+        )
 
     @runtime_dependency(
         module="shap",
         err_msg=(
-            "Please run `pip3 install shap` to install the required dependencies for model explanation."
+            "Please run `python3 -m pip install shap` to install the required dependencies for model explanation."
         ),
     )
-    def explain_model(self, datetime_col_name, explain_predict_fn) -> dict:
+    def explain_model(self):
         """
         Generates an explanation for the model by using the SHAP (Shapley Additive exPlanations) library.
         This function calculates the SHAP values for each feature in the dataset and stores the results in the `global_explanation` dictionary.
@@ -574,65 +635,61 @@ class ForecastOperatorBaseModel(ABC):
             dict: A dictionary containing the global explanation for each feature in the dataset.
                     The keys are the feature names and the values are the average absolute SHAP values.
         """
-        from shap import KernelExplainer
-
-        for series_id in self.target_columns:
-            self.series_id = series_id
-            if self.spec.model == SupportedModels.AutoTS:
-                self.dataset_cols = (
-                    self.full_data_long.loc[
-                        self.full_data_long.series_id == self.series_id
-                    ]
-                    .set_index(datetime_col_name)
-                    .columns
-                )
+        from shap import PermutationExplainer
 
-                self.bg_data = self.full_data_long.loc[
-                    self.full_data_long.series_id == self.series_id
-                ].set_index(datetime_col_name)
+        datetime_col_name = self.datasets._datetime_column_name
 
-            else:
-                self.dataset_cols = (
-                    self.full_data_dict.get(series_id)
-                    .set_index(datetime_col_name)
-                    .drop(series_id, axis=1)
-                    .columns
-                )
+        exp_start_time = time.time()
+        global_ex_time = 0
+        local_ex_time = 0
+        logger.info(
+            f"Calculating explanations using {self.spec.explanations_accuracy_mode} mode"
+        )
+        ratio = SpeedAccuracyMode.ratio[self.spec.explanations_accuracy_mode]
 
-                self.bg_data = self.full_data_dict.get(series_id).set_index(
-                    datetime_col_name
-                )
+        for s_id, data_i in self.datasets.get_data_by_series(
+            include_horizon=False
+        ).items():
+            explain_predict_fn = self.get_explain_predict_fn(series_id=s_id)
 
-            kernel_explnr = KernelExplainer(
-                model=explain_predict_fn,
-                data=self.bg_data[list(self.dataset_cols)][: -self.spec.horizon][
-                    list(self.dataset_cols)
-                ],
-                keep_index=False
-                if self.spec.model == SupportedModels.AutoMLX
-                else True,
+            data_trimmed = data_i.tail(max(int(len(data_i) * ratio), 5)).reset_index(
+                drop=True
+            )
+            data_trimmed[datetime_col_name] = data_trimmed[datetime_col_name].apply(
+                lambda x: x.timestamp()
             )
 
-            kernel_explnr_vals = kernel_explnr.shap_values(
-                self.bg_data[: -self.spec.horizon][list(self.dataset_cols)],
-                nsamples=50,
+            kernel_explnr = PermutationExplainer(
+                model=explain_predict_fn, masker=data_trimmed
             )
+            kernel_explnr_vals = kernel_explnr.shap_values(data_trimmed)
+
+            exp_end_time = time.time()
+            global_ex_time = global_ex_time + exp_end_time - exp_start_time
+
+            self.local_explainer(
+                kernel_explnr, series_id=s_id, datetime_col_name=datetime_col_name
+            )
+            local_ex_time = local_ex_time + time.time() - exp_end_time
 
             if not len(kernel_explnr_vals):
                 logger.warn(
                     f"No explanations generated. Ensure that additional data has been provided."
                 )
             else:
-                self.global_explanation[series_id] = dict(
+                self.global_explanation[s_id] = dict(
                     zip(
-                        self.dataset_cols,
-                        np.average(np.absolute(kernel_explnr_vals), axis=0),
+                        data_trimmed.columns[1:],
+                        np.average(np.absolute(kernel_explnr_vals[:, 1:]), axis=0),
                     )
                 )
 
-            self.local_explainer(
-                kernel_explnr, series_id=series_id, datetime_col_name=datetime_col_name
-            )
+        logger.info(
+            "Global explanations generation completed in %s seconds", global_ex_time
+        )
+        logger.info(
+            "Local explanations generation completed in %s seconds", local_ex_time
+        )
 
     def local_explainer(self, kernel_explainer, series_id, datetime_col_name) -> None:
         """
@@ -642,24 +699,34 @@ class ForecastOperatorBaseModel(ABC):
         ----------
             kernel_explainer: The kernel explainer object to use for generating explanations.
         """
-        # Get the data for the series ID and select the relevant columns
-        # data = self.full_data_dict.get(series_id).set_index(datetime_col_name)
-        data = self.bg_data[-self.spec.horizon :][list(self.dataset_cols)]
+        data = self.datasets.get_horizon_at_series(s_id=series_id)
 
+        data[datetime_col_name] = datetime_to_seconds(data[datetime_col_name])
+        data = data.reset_index(drop=True)
         # Generate local SHAP values using the kernel explainer
-        local_kernel_explnr_vals = kernel_explainer.shap_values(data, nsamples=50)
+        local_kernel_explnr_vals = kernel_explainer.shap_values(data)
 
         # Convert the SHAP values into a DataFrame
         local_kernel_explnr_df = pd.DataFrame(
-            local_kernel_explnr_vals, columns=self.dataset_cols
+            local_kernel_explnr_vals, columns=data.columns
         )
+        self.local_explanation[series_id] = local_kernel_explnr_df
 
-        # set the index of the DataFrame to the datetime column
-        local_kernel_explnr_df.index = data.index
-
-        if self.spec.model == SupportedModels.AutoTS:
-            local_kernel_explnr_df.drop(
-                ["series_id", self.spec.target_column], axis=1, inplace=True
+    def get_explain_predict_fn(self, series_id, fcst_col_name="yhat"):
+        def _custom_predict(
+            data,
+            model=self.models[series_id],
+            dt_column_name=self.datasets._datetime_column_name,
+        ):
+            """
+            data: ForecastDatasets.get_data_at_series(s_id)
+            """
+            data[dt_column_name] = seconds_to_datetime(
+                data[dt_column_name], dt_format=self.spec.datetime_column.format
             )
+            data = self.preprocess(df=data, series_id=series_id)
+            data[self.original_target_column] = None
+            fcst = model.predict(data)[fcst_col_name]
+            return fcst
 
-        self.local_explanation[series_id] = local_kernel_explnr_df
+        return _custom_predict