Merge branch 'main' into ahosler-patch-1

Author: ahosler

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

@@ -116,7 +116,10 @@ def _train_model(self, i, s_id, df, model_kwargs):
                 lower_bound=self.get_horizon(forecast["yhat_lower"]).values,
             )
 
-            self.models[s_id] = model
+            self.models[s_id] = {}
+            self.models[s_id]["model"] = model
+            self.models[s_id]["le"] = self.le[s_id]
+            self.models[s_id]["predict_component_cols"] = X_pred.columns
 
             params = vars(model).copy()
             for param in ["arima_res_", "endog_index_"]:
@@ -163,7 +166,7 @@ def _generate_report(self):
             sec5_text = rc.Heading("ARIMA Model Parameters", level=2)
             blocks = [
                 rc.Html(
-                    m.summary().as_html(),
+                    m['model'].summary().as_html(),
                     label=s_id if self.target_cat_col else None,
                 )
                 for i, (s_id, m) in enumerate(self.models.items())
@@ -251,7 +254,7 @@ def _generate_report(self):
     def get_explain_predict_fn(self, series_id):
         def _custom_predict(
             data,
-            model=self.models[series_id],
+            model=self.models[series_id]["model"],
             dt_column_name=self.datasets._datetime_column_name,
             target_col=self.original_target_column,
         ):

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

@@ -56,8 +56,8 @@ def set_kwargs(self):
         )
         return model_kwargs_cleaned, time_budget
 
-    def preprocess(self, data):  # TODO: re-use self.le for explanations
-        _, df_encoded = _label_encode_dataframe(
+    def preprocess(self, data, series_id):  # TODO: re-use self.le for explanations
+        self.le[series_id], df_encoded = _label_encode_dataframe(
             data,
             no_encode={self.spec.datetime_column.name, self.original_target_column},
         )
@@ -124,7 +124,7 @@ def _build_model(self) -> pd.DataFrame:
                 self.forecast_output.init_series_output(
                     series_id=s_id, data_at_series=df
                 )
-                data = self.preprocess(df)
+                data = self.preprocess(df, s_id)
                 data_i = self.drop_horizon(data)
                 X_pred = self.get_horizon(data).drop(target, axis=1)
 
@@ -156,7 +156,9 @@ def _build_model(self) -> pd.DataFrame:
                     target
                 ].values
 
-                self.models[s_id] = model
+                self.models[s_id] = {}
+                self.models[s_id]["model"] = model
+                self.models[s_id]["le"] = self.le[s_id]
 
                 # In case of Naive model, model.forecast function call does not return confidence intervals.
                 if f"{target}_ci_upper" not in summary_frame:
@@ -217,7 +219,8 @@ def _generate_report(self):
         other_sections = []
 
         if len(self.models) > 0:
-            for s_id, m in models.items():
+            for s_id, artifacts in models.items():
+                m = artifacts["model"]
                 selected_models[s_id] = {
                     "series_id": s_id,
                     "selected_model": m.selected_model_,
@@ -323,7 +326,7 @@ def _generate_report(self):
         )
 
     def get_explain_predict_fn(self, series_id):
-        selected_model = self.models[series_id]
+        selected_model = self.models[series_id]["model"]
 
         # If training date, use method below. If future date, use forecast!
         def _custom_predict_fn(
@@ -341,12 +344,12 @@ def _custom_predict_fn(
             data[dt_column_name] = seconds_to_datetime(
                 data[dt_column_name], dt_format=self.spec.datetime_column.format
             )
-            data = self.preprocess(data)
+            data = self.preprocess(data, series_id)
             horizon_data = horizon_data.drop(target_col, axis=1)
             horizon_data[dt_column_name] = seconds_to_datetime(
                 horizon_data[dt_column_name], dt_format=self.spec.datetime_column.format
             )
-            horizon_data = self.preprocess(horizon_data)
+            horizon_data = self.preprocess(horizon_data, series_id)
 
             rows = []
             for i in range(data.shape[0]):
@@ -424,10 +427,8 @@ def explain_model(self):
                 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
-                        )
+                        self.models[s_id]["model"],
+                        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

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

@@ -828,7 +828,7 @@ def local_explainer(self, kernel_explainer, series_id, datetime_col_name) -> Non
     def get_explain_predict_fn(self, series_id, fcst_col_name="yhat"):
         def _custom_predict(
             data,
-            model=self.models[series_id],
+            model=self.models[series_id]["model"],
             dt_column_name=self.datasets._datetime_column_name,
         ):
             """

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

@@ -172,8 +172,10 @@ def _train_model(self, i, s_id, df, model_kwargs):
                 ).values,
             )
 
-            self.models[s_id] = model
             self.trainers[s_id] = model.trainer
+            self.models[s_id] = {}
+            self.models[s_id]["model"] = model
+            self.models[s_id]["le"] = self.le[s_id]
 
             self.model_parameters[s_id] = {
                 "framework": SupportedModels.NeuralProphet,
@@ -355,7 +357,8 @@ def _generate_report(self):
 
             sec5_text = rc.Heading("Neural Prophet Model Parameters", level=2)
             model_states = []
-            for s_id, m in self.models.items():
+            for s_id, artifacts in self.models.items():
+                m = artifacts["model"]
                 model_states.append(
                     pd.Series(
                         m.state_dict(),

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

@@ -112,7 +112,10 @@ def _train_model(self, i, series_id, df, model_kwargs):
                 upper_bound=self.get_horizon(forecast["yhat_upper"]).values,
                 lower_bound=self.get_horizon(forecast["yhat_lower"]).values,
             )
-            self.models[series_id] = model
+
+            self.models[series_id] = {}
+            self.models[series_id]["model"] = model
+            self.models[series_id]["le"] = self.le[series_id]
 
             params = vars(model).copy()
             for param in ["history", "history_dates", "stan_fit"]:
@@ -256,7 +259,7 @@ def _generate_report(self):
         all_sections = []
         if len(series_ids) > 0:
             sec1 = _select_plot_list(
-                lambda s_id: self.models[s_id].plot(
+                lambda s_id: self.models[s_id]["model"].plot(
                     self.outputs[s_id], include_legend=True
                 ),
                 series_ids=series_ids,
@@ -271,7 +274,7 @@ def _generate_report(self):
             )
 
             sec2 = _select_plot_list(
-                lambda s_id: self.models[s_id].plot_components(self.outputs[s_id]),
+                lambda s_id: self.models[s_id]["model"].plot_components(self.outputs[s_id]),
                 series_ids=series_ids,
                 target_category_column=self.target_cat_col,
             )
@@ -280,11 +283,11 @@ def _generate_report(self):
             )
 
             sec3_figs = {
-                s_id: self.models[s_id].plot(self.outputs[s_id]) for s_id in series_ids
+                s_id: self.models[s_id]["model"].plot(self.outputs[s_id]) for s_id in series_ids
             }
             for s_id in series_ids:
                 add_changepoints_to_plot(
-                    sec3_figs[s_id].gca(), self.models[s_id], self.outputs[s_id]
+                    sec3_figs[s_id].gca(), self.models[s_id]["model"], self.outputs[s_id]
                 )
             sec3 = _select_plot_list(
                 lambda s_id: sec3_figs[s_id],
@@ -298,7 +301,7 @@ def _generate_report(self):
             sec5_text = rc.Heading("Prophet Model Seasonality Components", level=2)
             model_states = []
             for s_id in series_ids:
-                m = self.models[s_id]
+                m = self.models[s_id]["model"]
                 model_states.append(
                     pd.Series(
                         m.seasonalities,

ads/opctl/operator/lowcode/forecast/whatifserve/score.py

@@ -151,34 +151,42 @@ def get_forecast(future_df, model_name, series_id, model_object, date_col, targe
         pred_obj = model_object.predict(future_regressor=future_reg)
         return pred_obj.forecast[series_id].tolist()
     elif model_name == SupportedModels.Prophet and series_id in model_object:
-        model = model_object[series_id]
+        model = model_object[series_id]['model']
+        label_encoder = model_object[series_id]['le']
         processed = future_df.rename(columns={date_col_name: 'ds', target_column: 'y'})
-        forecast = model.predict(processed)
+        encoded_df = label_encoder.transform(processed)
+        forecast = model.predict(encoded_df)
         return forecast['yhat'].tolist()
     elif model_name == SupportedModels.NeuralProphet and series_id in model_object:
-        model = model_object[series_id]
+        model = model_object[series_id]['model']
+        label_encoder = model_object[series_id]['le']
         model.restore_trainer()
         accepted_regressors = list(model.config_regressors.regressors.keys())
         data = future_df.rename(columns={date_col_name: 'ds', target_column: 'y'})
-        future = data[accepted_regressors + ["ds"]].reset_index(drop=True)
+        encoded_df = label_encoder.transform(data)
+        future = encoded_df[accepted_regressors + ["ds"]].reset_index(drop=True)
         future["y"] = None
         forecast = model.predict(future)
         return forecast['yhat1'].tolist()
     elif model_name == SupportedModels.Arima and series_id in model_object:
-        model = model_object[series_id]
-        future_df = future_df.set_index(date_col_name)
-        x_pred = future_df.drop(target_cat_col, axis=1)
+        model = model_object[series_id]['model']
+        label_encoder = model_object[series_id]['le']
+        predict_cols = model_object[series_id]["predict_component_cols"]
+        encoded_df = label_encoder.transform(future_df)
+        x_pred = encoded_df.set_index(date_col_name)
+        x_pred = x_pred.drop(target_cat_col, axis=1)
         yhat, conf_int = model.predict(
             n_periods=horizon,
-            X=x_pred,
+            X=x_pred[predict_cols],
             return_conf_int=True
         )
         yhat_clean = pd.DataFrame(yhat, index=yhat.index, columns=["yhat"])
         return yhat_clean['yhat'].tolist()
     elif model_name == SupportedModels.AutoMLX and series_id in model_object:
-        # automlx model
-        model = model_object[series_id]
-        x_pred = future_df.drop(target_cat_col, axis=1)
+        model = model_object[series_id]['model']
+        label_encoder = model_object[series_id]['le']
+        encoded_df = label_encoder.transform(future_df)
+        x_pred = encoded_df.drop(target_cat_col, axis=1)
         x_pred = x_pred.set_index(date_col_name)
         forecast = model.forecast(
             X=x_pred,

tests/operators/forecast/test_errors.py

@@ -835,16 +835,9 @@ def test_what_if_analysis(operator_setup, model):
     historical_data = pd.read_csv(historical_data_path, parse_dates=["Date"])
     historical_filtered = historical_data[historical_data["Date"] > "2013-03-01"]
     additional_data = pd.read_csv(additional_data_path, parse_dates=["Date"])
-    add_filtered = additional_data[additional_data["Date"] > "2013-03-01"]
-    numeric_columns = add_filtered.select_dtypes(
-        include=["number", "object", "datetime64"]
-    )
-    non_constant_columns = numeric_columns.columns[
-        (numeric_columns != numeric_columns.iloc[0]).any()
-    ]
-    df_non_constant = numeric_columns[non_constant_columns.union(["Store"])]
-    df_non_constant.to_csv(f"{additional_test_path}", index=False)
-    historical_filtered.to_csv(f"{historical_test_path}", index=False)
+    add_filtered = additional_data[additional_data['Date'] > "2013-03-01"]
+    add_filtered.to_csv(f'{additional_test_path}', index=False)
+    historical_filtered.to_csv(f'{historical_test_path}', index=False)
 
     yaml_i, output_data_path = populate_yaml(
         tmpdirname=tmpdirname,