From b4e5637dda732b7bddcac447e61f4b0cd1d4ea1a Mon Sep 17 00:00:00 2001
From: Maisa Ben Salah <maisabensalah@AminsMBP131.attlocal.net>
Date: Thu, 13 Jun 2024 16:57:34 -0700
Subject: [PATCH 1/3] add io buffer support

---
 tests/test_utils.py | 32 ++++++++++++++++++++++++++++++++
 1 file changed, 32 insertions(+)

diff --git a/tests/test_utils.py b/tests/test_utils.py
index 3e965e03a..93ddc1073 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -3,6 +3,7 @@
 import logging
 import os
 import pathlib
+import io
 
 import pandas as pd
 import pytest
@@ -66,6 +67,37 @@ def test_save_load():
     pd.testing.assert_frame_equal(forecast, forecast2)
     pd.testing.assert_frame_equal(forecast, forecast3)
 
+def test_save_load_io():
+    df = pd.read_csv(PEYTON_FILE, nrows=NROWS)
+    m = NeuralProphet(
+        epochs=EPOCHS,
+        batch_size=BATCH_SIZE,
+        learning_rate=LR,
+        n_lags=6,
+        n_forecasts=3,
+        n_changepoints=0,
+    )
+    _ = m.fit(df, freq="D")
+    future = m.make_future_dataframe(df, periods=3)
+    forecast = m.predict(df=future)
+    
+    # Save the model to an in-memory buffer
+    log.info("testing: save to buffer")
+    buffer = io.BytesIO()
+    save(m, buffer)
+    buffer.seek(0)  # Reset buffer position to the beginning
+    
+    log.info("testing: load from buffer")
+    m2 = load(buffer)
+    forecast2 = m2.predict(df=future)
+
+    buffer.seek(0)  # Reset buffer position to the beginning for another load
+    m3 = load(buffer, map_location="cpu")
+    forecast3 = m3.predict(df=future)
+
+    # Check that the forecasts are the same
+    pd.testing.assert_frame_equal(forecast, forecast2)
+    pd.testing.assert_frame_equal(forecast, forecast3)
 
 # TODO: add functionality to continue training
 # def test_continue_training():

From 2b50d35ce5392670fa109f675b1931083f1e126f Mon Sep 17 00:00:00 2001
From: Maisa Ben Salah <maisabensalah@AminsMBP131.attlocal.net>
Date: Thu, 13 Jun 2024 16:58:21 -0700
Subject: [PATCH 2/3] added tests

---
 neuralprophet/utils.py | 20 ++++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/neuralprophet/utils.py b/neuralprophet/utils.py
index d5fabe203..d0243a652 100644
--- a/neuralprophet/utils.py
+++ b/neuralprophet/utils.py
@@ -5,7 +5,7 @@
 import os
 import sys
 from collections import OrderedDict
-from typing import TYPE_CHECKING, Iterable, Optional, Union
+from typing import TYPE_CHECKING, Iterable, Optional, Union, BinaryIO, IO, TypeAlias
 
 import numpy as np
 import pandas as pd
@@ -21,19 +21,23 @@
 
 log = logging.getLogger("NP.utils")
 
+FILE_LIKE: TypeAlias = Union[str, os.PathLike, BinaryIO, IO[bytes]]
 
-def save(forecaster, path: str):
+def save(forecaster, path: FILE_LIKE):
     """Save a fitted Neural Prophet model to disk.
 
     Parameters:
         forecaster : np.forecaster.NeuralProphet
             input forecaster that is fitted
-        path : str
-            path and filename to be saved. filename could be any but suggested to have extension .np.
+        path : FILE_LIKE
+            Path and filename to be saved, or an in-memory buffer. Filename could be any but suggested to have extension .np.
 
     After you fitted a model, you may save the model to save_test_model.np
         >>> from neuralprophet import save
         >>> save(forecaster, "test_save_model.np")
+        >>> import io
+        >>> buffer = io.BytesIO()
+        >>> save(forecaster, buffer)
     """
     # List of attributes to remove
     attrs_to_remove_forecaster = ["trainer"]
@@ -69,13 +73,13 @@ def save(forecaster, path: str):
                 setattr(forecaster.model, attr, value)
 
 
-def load(path: str, map_location=None):
-    """retrieve a fitted model from a .np file that was saved by save.
+def load(path: FILE_LIKE, map_location=None):
+    """retrieve a fitted model from a .np file or buffer that was saved by save.
 
     Parameters
     ----------
-        path : str
-            path and filename to be saved. filename could be any but suggested to have extension .np.
+        path : FILE_LIKE
+            Path and filename to be saved, or an in-memory buffer. Filename could be any but suggested to have extension .np.
         map_location : str, optional
             specifying the location where the model should be loaded.
             If you are running on a CPU-only machine, set map_location='cpu' to map your storages to the CPU.

From 1aab28f9c14aa5e620b747b9228d09b66c14a7a8 Mon Sep 17 00:00:00 2001
From: Maisa Ben Salah <maiisabensalah@gmail.com>
Date: Thu, 13 Jun 2024 19:39:21 -0700
Subject: [PATCH 3/3] remove typealias

---
 neuralprophet/utils.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/neuralprophet/utils.py b/neuralprophet/utils.py
index d0243a652..168d55f67 100644
--- a/neuralprophet/utils.py
+++ b/neuralprophet/utils.py
@@ -5,7 +5,7 @@
 import os
 import sys
 from collections import OrderedDict
-from typing import TYPE_CHECKING, Iterable, Optional, Union, BinaryIO, IO, TypeAlias
+from typing import TYPE_CHECKING, Iterable, Optional, Union, BinaryIO, IO
 
 import numpy as np
 import pandas as pd
@@ -21,7 +21,7 @@
 
 log = logging.getLogger("NP.utils")
 
-FILE_LIKE: TypeAlias = Union[str, os.PathLike, BinaryIO, IO[bytes]]
+FILE_LIKE = Union[str, os.PathLike, BinaryIO, IO[bytes]]
 
 def save(forecaster, path: FILE_LIKE):
     """Save a fitted Neural Prophet model to disk.