Bug: mlforecast not taking datetime column from the spec (#860)

oracle · May 30, 2024 · 79ab368 · 79ab368
2 parents dc583b3 + 5934b73
commit 79ab368
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/ads/opctl/operator/lowcode/forecast/model/ml_forecast.py b/ads/opctl/operator/lowcode/forecast/model/ml_forecast.py
@@ -24,6 +24,7 @@ def __init__(self, config: ForecastOperatorConfig, datasets: ForecastDatasets):
         self.local_explanation = {}
         self.formatted_global_explanation = None
         self.formatted_local_explanation = None
+        self.date_col = config.spec.datetime_column.name
 
     def set_kwargs(self):
         """
@@ -77,8 +78,8 @@ def _train_model(self, data_train, data_test, model_kwargs):
                         alpha=model_kwargs["lower_quantile"],
                     ),
                 },
-                freq=pd.infer_freq(data_train["Date"].drop_duplicates())
-                or pd.infer_freq(data_train["Date"].drop_duplicates()[-5:]),
+                freq=pd.infer_freq(data_train[self.date_col].drop_duplicates())
+                or pd.infer_freq(data_train[self.date_col].drop_duplicates()[-5:]),
                 target_transforms=[Differences([12])],
                 lags=model_kwargs.get(
                     "lags",
@@ -108,7 +109,7 @@ def _train_model(self, data_train, data_test, model_kwargs):
                 data_train[self.model_columns],
                 static_features=model_kwargs.get("static_features", []),
                 id_col=ForecastOutputColumns.SERIES,
-                time_col=self.spec.datetime_column.name,
+                time_col=self.date_col,
                 target_col=self.spec.target_column,
                 fitted=True,
                 max_horizon=None if num_models is False else self.spec.horizon,
@@ -173,7 +174,7 @@ def _build_model(self) -> pd.DataFrame:
             confidence_interval_width=self.spec.confidence_interval_width,
             horizon=self.spec.horizon,
             target_column=self.original_target_column,
-            dt_column=self.spec.datetime_column.name,
+            dt_column=self.date_col,
         )
         self._train_model(data_train, data_test, model_kwargs)
         return self.forecast_output.get_forecast_long()