Update src/nf.py

src/nf.py

@@ -8,11 +8,10 @@ import pandas as pd
 import pytorch_lightning as pl
 from datasetsforecast.utils import download_file
 from hyperopt import hp
+from neuralforecast.core import NeuralForecast
 from neuralforecast.auto import NHITS as autoNHITS
-from neuralforecast.data.tsdataset import WindowsDataset
-from neuralforecast.data.tsloader import TimeSeriesLoader
-from neuralforecast.models.mqnhits.mqnhits import MQNHITS
-from neuralforecast.models.nhits.nhits import NHITS
+from neuralforecast.tsdataset import TimeSeriesDataset, TimeSeriesLoader
+from neuralforecast.models import NHITS
 
 # GLOBAL PARAMETERS
 DEFAULT_HORIZON = 30
@@ -73,7 +72,6 @@ MODELS = {
     },
 }
 
-
 def download_models():
     for _, meta in MODELS.items():
         if not Path(f'./models/{meta["model"]}.ckpt').is_file():
@@ -82,19 +80,17 @@ def download_models():
                 f'https://nixtla-public.s3.amazonaws.com/transfer/pretrained_models/{meta["model"]}.ckpt',
             )
 
-
 download_models()
 
-
 class StandardScaler:
     """This class helps to standardize a dataframe with multiple time series."""
-
     def __init__(self):
-        self.norm: pd.DataFrame
+        self.norm: pd.DataFrame = None
 
     def fit(self, X: pd.DataFrame) -> "StandardScaler":
         self.norm = X.groupby("unique_id").agg({"y": [np.mean, np.std]})
         self.norm = self.norm.droplevel(0, 1).reset_index()
+        return self
 
     def transform(self, X: pd.DataFrame) -> pd.DataFrame:
         transformed = X.merge(self.norm, how="left", on=["unique_id"])
@@ -109,7 +105,6 @@ class StandardScaler:
             )
         return transformed[["unique_id", "ds"] + cols]
 
-
 def compute_ds_future(Y_df, fh):
     if Y_df["unique_id"].nunique() == 1:
         ds_ = pd.to_datetime(Y_df["ds"].values)
@@ -130,10 +125,7 @@ def compute_ds_future(Y_df, fh):
         )
         return list(ds_future)
 
-
-def forecast_pretrained_model(
-    Y_df: pd.DataFrame, model: str, fh: int, max_steps: int = 0
-):
+def forecast_pretrained_model(Y_df: pd.DataFrame, model: str, fh: int, max_steps: int = 0):
     if "unique_id" not in Y_df:
         Y_df.insert(0, "unique_id", "ts_1")
 
@@ -143,26 +135,12 @@ def forecast_pretrained_model(
 
     # Model
     file_ = f"./models/{model}.ckpt"
-    mqnhits = MQNHITS.load_from_checkpoint(file_)
+    nhits = NHITS.load_from_checkpoint(file_)
 
     # Fit
     if max_steps > 0:
-        train_dataset = WindowsDataset(
-            Y_df=Y_df,
-            X_df=None,
-            S_df=None,
-            mask_df=None,
-            f_cols=[],
-            input_size=mqnhits.n_time_in,
-            output_size=mqnhits.n_time_out,
-            sample_freq=1,
-            complete_windows=True,
-            verbose=False,
-        )
-
-        train_loader = TimeSeriesLoader(
-            dataset=train_dataset, batch_size=1, n_windows=32, shuffle=True
-        )
+        train_dataset = TimeSeriesDataset.from_dataframe(Y_df, input_size=nhits.hparams.n_time_in, output_size=nhits.hparams.n_time_out)
+        train_loader = TimeSeriesLoader(dataset=train_dataset, batch_size=1, n_windows=32, shuffle=True)
 
         trainer = pl.Trainer(
             max_epochs=None,
@@ -174,13 +152,12 @@ def forecast_pretrained_model(
             log_every_n_steps=1,
         )
 
-        trainer.fit(mqnhits, train_loader)
+        trainer.fit(nhits, train_loader)
 
     # Forecast
-    forecast_df = mqnhits.forecast(Y_df=Y_df)
+    forecast_df = nhits.forecast(Y_df=Y_df)
     forecast_df = scaler.inverse_transform(forecast_df, cols=["y_5", "y_50", "y_95"])
 
-    # Foreoast
     n_ts = forecast_df["unique_id"].nunique()
     if fh * n_ts > len(forecast_df):
         forecast_df = (
@@ -194,7 +171,6 @@ def forecast_pretrained_model(
 
     return forecast_df
 
-
 if __name__ == "__main__":
     df = pd.read_csv(
         "https://raw.githubusercontent.com/Nixtla/transfer-learning-time-series/main/datasets/ercot_COAST.csv"
@@ -208,4 +184,4 @@ if __name__ == "__main__":
         assert forecast.shape == (80, 5)
         # test multiple time series
         multi_forecast = forecast_pretrained_model(multi_df, model=meta["model"], fh=80)
-        assert multi_forecast.shape == (80 * 2, 5)
+        assert multi_forecast.shape == (80 * 2, 5)