FIX: Add logic to load custom models when using ReduceLROnPlateau (#1340)

Co-authored-by: Olivier Sprangers <45119856+elephaint@users.noreply.github.com>

Author: marcopeix

nbs/core.ipynb

@@ -1481,7 +1481,14 @@
     "                continue\n",
     "\n",
     "            model_name = repr(model)\n",
-    "            model_class_name = model.__class__.__name__.lower()\n",
+    "            if model.__class__.__name__.lower() in MODEL_FILENAME_DICT:\n",
+    "                model_class_name = model.__class__.__name__.lower()\n",
+    "            elif model.__class__.__base__.__name__.lower() in MODEL_FILENAME_DICT:\n",
+    "                model_class_name = model.__class__.__base__.__name__.lower()\n",
+    "            else:\n",
+    "                raise ValueError(\n",
+    "                    f\"Model {model.__class__.__name__} is not supported for saving.\"\n",
+    "                )\n",
     "            alias_to_model[model_name] = model_class_name\n",
     "            count_names[model_name] = count_names.get(model_name, -1) + 1\n",
     "            model.save(f\"{path}/{model_name}_{count_names[model_name]}.ckpt\")\n",
@@ -1529,7 +1536,7 @@
     "\n",
     "        with fsspec.open(f\"{path}/configuration.pkl\", \"wb\") as f:\n",
     "            pickle.dump(config_dict, f)\n",
-    "\n",
+    "    \n",
     "    @staticmethod\n",
     "    def load(path, verbose=False, **kwargs):\n",
     "        \"\"\"Load NeuralForecast\n",
@@ -1550,50 +1557,60 @@
     "            Instantiated `NeuralForecast` class.\n",
     "        \"\"\"\n",
     "        # Standarize path without '/'\n",
-    "        if path[-1] == '/':\n",
+    "        if path[-1] == \"/\":\n",
     "            path = path[:-1]\n",
-    "        \n",
+    "\n",
     "        fs, _, _ = fsspec.get_fs_token_paths(path)\n",
-    "        files = [f.split('/')[-1] for f in fs.ls(path) if fs.isfile(f)]\n",
+    "        files = [f.split(\"/\")[-1] for f in fs.ls(path) if fs.isfile(f)]\n",
     "\n",
     "        # Load models\n",
-    "        models_ckpt = [f for f in files if f.endswith('.ckpt')]\n",
+    "        models_ckpt = [f for f in files if f.endswith(\".ckpt\")]\n",
     "        if len(models_ckpt) == 0:\n",
-    "            raise Exception('No model found in directory.') \n",
-    "        \n",
-    "        if verbose: print(10 * '-' + ' Loading models ' + 10 * '-')\n",
+    "            raise Exception(\"No model found in directory.\")\n",
+    "\n",
+    "        if verbose:\n",
+    "            print(10 * \"-\" + \" Loading models \" + 10 * \"-\")\n",
     "        models = []\n",
     "        try:\n",
-    "            with fsspec.open(f'{path}/alias_to_model.pkl', 'rb') as f:\n",
+    "            with fsspec.open(f\"{path}/alias_to_model.pkl\", \"rb\") as f:\n",
     "                alias_to_model = pickle.load(f)\n",
     "        except FileNotFoundError:\n",
     "            alias_to_model = {}\n",
+    "        \n",
     "        for model in models_ckpt:\n",
-    "            model_name = '_'.join(model.split('_')[:-1])\n",
+    "            model_name = \"_\".join(model.split(\"_\")[:-1])\n",
     "            model_class_name = alias_to_model.get(model_name, model_name)\n",
-    "            loaded_model = MODEL_FILENAME_DICT[model_class_name].load(f'{path}/{model}', **kwargs)\n",
+    "            loaded_model = MODEL_FILENAME_DICT[model_class_name].load(\n",
+    "                f\"{path}/{model}\", **kwargs\n",
+    "            )\n",
     "            loaded_model.alias = model_name\n",
     "            models.append(loaded_model)\n",
-    "            if verbose: print(f\"Model {model_name} loaded.\")\n",
+    "            if verbose:\n",
+    "                print(f\"Model {model_name} loaded.\")\n",
     "\n",
-    "        if verbose: print(10*'-' + ' Loading dataset ' + 10*'-')\n",
+    "        if verbose:\n",
+    "            print(10 * \"-\" + \" Loading dataset \" + 10 * \"-\")\n",
     "        # Load dataset\n",
     "        try:\n",
     "            with fsspec.open(f\"{path}/dataset.pkl\", \"rb\") as f:\n",
     "                dataset = pickle.load(f)\n",
-    "            if verbose: print('Dataset loaded.')\n",
+    "            if verbose:\n",
+    "                print(\"Dataset loaded.\")\n",
     "        except FileNotFoundError:\n",
     "            dataset = None\n",
-    "            if verbose: print('No dataset found in directory.')\n",
+    "            if verbose:\n",
+    "                print(\"No dataset found in directory.\")\n",
     "\n",
-    "        if verbose: print(10*'-' + ' Loading configuration ' + 10*'-')\n",
+    "        if verbose:\n",
+    "            print(10 * \"-\" + \" Loading configuration \" + 10 * \"-\")\n",
     "        # Load configuration\n",
     "        try:\n",
     "            with fsspec.open(f\"{path}/configuration.pkl\", \"rb\") as f:\n",
     "                config_dict = pickle.load(f)\n",
-    "            if verbose: print('Configuration loaded.')\n",
+    "            if verbose:\n",
+    "                print(\"Configuration loaded.\")\n",
     "        except FileNotFoundError:\n",
-    "            raise Exception('No configuration found in directory.')\n",
+    "            raise Exception(\"No configuration found in directory.\")\n",
     "\n",
     "        # in 1.6.4, `local_scaler_type` / `scalers_` lived on the dataset.\n",
     "        # in order to preserve backwards-compatibility, we check to see if these are found on the dataset\n",
@@ -1604,34 +1621,30 @@
     "        # Create NeuralForecast object\n",
     "        neuralforecast = NeuralForecast(\n",
     "            models=models,\n",
-    "            freq=config_dict['freq'],\n",
+    "            freq=config_dict[\"freq\"],\n",
     "            local_scaler_type=config_dict.get(\"local_scaler_type\", default_scalar_type),\n",
     "        )\n",
     "\n",
-    "        attr_to_default = {\n",
-    "            \"id_col\": \"unique_id\",\n",
-    "            \"time_col\": \"ds\",\n",
-    "            \"target_col\": \"y\"\n",
-    "        }\n",
+    "        attr_to_default = {\"id_col\": \"unique_id\", \"time_col\": \"ds\", \"target_col\": \"y\"}\n",
     "        for attr, default in attr_to_default.items():\n",
     "            setattr(neuralforecast, attr, config_dict.get(attr, default))\n",
     "        # only restore attribute if available\n",
-    "        for attr in ['prediction_intervals', '_cs_df']:\n",
+    "        for attr in [\"prediction_intervals\", \"_cs_df\"]:\n",
     "            setattr(neuralforecast, attr, config_dict.get(attr, None))\n",
     "\n",
     "        # Dataset\n",
     "        if dataset is not None:\n",
     "            neuralforecast.dataset = dataset\n",
     "            restore_attrs = [\n",
-    "                'uids',\n",
-    "                'last_dates',\n",
-    "                'ds',\n",
+    "                \"uids\",\n",
+    "                \"last_dates\",\n",
+    "                \"ds\",\n",
     "            ]\n",
     "            for attr in restore_attrs:\n",
     "                setattr(neuralforecast, attr, config_dict[attr])\n",
     "\n",
     "        # Fitted flag\n",
-    "        neuralforecast._fitted = config_dict['_fitted']\n",
+    "        neuralforecast._fitted = config_dict[\"_fitted\"]\n",
     "\n",
     "        neuralforecast.scalers_ = config_dict.get(\"scalers_\", default_scalars_)\n",
     "\n",

neuralforecast/core.py

@@ -1488,7 +1488,14 @@ def save(
                 continue
 
             model_name = repr(model)
-            model_class_name = model.__class__.__name__.lower()
+            if model.__class__.__name__.lower() in MODEL_FILENAME_DICT:
+                model_class_name = model.__class__.__name__.lower()
+            elif model.__class__.__base__.__name__.lower() in MODEL_FILENAME_DICT:
+                model_class_name = model.__class__.__base__.__name__.lower()
+            else:
+                raise ValueError(
+                    f"Model {model.__class__.__name__} is not supported for saving."
+                )
             alias_to_model[model_name] = model_class_name
             count_names[model_name] = count_names.get(model_name, -1) + 1
             model.save(f"{path}/{model_name}_{count_names[model_name]}.ckpt")
@@ -1577,6 +1584,7 @@ def load(path, verbose=False, **kwargs):
                 alias_to_model = pickle.load(f)
         except FileNotFoundError:
             alias_to_model = {}
+
         for model in models_ckpt:
             model_name = "_".join(model.split("_")[:-1])
             model_class_name = alias_to_model.get(model_name, model_name)