Feature/models saving (#206)

- Added models saving and loading

Closes #100

Author: feldlime

CHANGELOG.md

@@ -9,8 +9,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## Unreleased
 
 ### Added
-- `from_config`, `get_config` and `get_params` methods to all models except neural-net-based([#170](https://github.com/MobileTeleSystems/RecTools/pull/170))
+- `from_config`, `get_config` and `get_params` methods to all models except neural-net-based ([#170](https://github.com/MobileTeleSystems/RecTools/pull/170))
 - Optional `epochs` argument to `ImplicitALSWrapperModel.fit` method ([#203](https://github.com/MobileTeleSystems/RecTools/pull/203))
+- `save` and `load` methods to all of the models ([#206](https://github.com/MobileTeleSystems/RecTools/pull/206))
 
 
 ## [0.8.0] - 28.08.2024

README.md

@@ -150,7 +150,7 @@ See [recommender baselines extended tutorial](https://github.com/MobileTeleSyste
 [Contributing guide](CONTRIBUTING.rst)
 
 To install all requirements
-- you must have `python>=3.8` and `poetry>=1.5.0` installed
+- you must have `python3` and `poetry` installed
 - make sure you have no active virtual environments (deactivate conda `base` if applicable)
 - run
 ```

rectools/models/base.py

@@ -14,8 +14,10 @@
 
 """Base model."""
 
+import pickle
 import typing as tp
 import warnings
+from pathlib import Path
 
 import numpy as np
 import pandas as pd
@@ -42,6 +44,10 @@
 
 RecoTriplet_T = tp.TypeVar("RecoTriplet_T", InternalRecoTriplet, SemiInternalRecoTriplet, ExternalRecoTriplet)
 
+FileLike = tp.Union[str, Path, tp.IO[bytes]]
+
+PICKLE_PROTOCOL = 5
+
 
 def _serialize_random_state(rs: tp.Optional[tp.Union[None, int, np.random.RandomState]]) -> tp.Union[None, int]:
     if rs is None or isinstance(rs, int):
@@ -191,6 +197,85 @@ def from_config(cls, config: tp.Union[dict, ModelConfig_T]) -> tpe.Self:
     def _from_config(cls, config: ModelConfig_T) -> tpe.Self:
         raise NotImplementedError()
 
+    def save(self, f: FileLike) -> int:
+        """
+        Save model to file.
+
+        Parameters
+        ----------
+        f : str or Path or file-like object
+            Path to file or file-like object.
+
+        Returns
+        -------
+        int
+            Number of bytes written.
+        """
+        data = self.dumps()
+
+        if isinstance(f, (str, Path)):
+            return Path(f).write_bytes(data)
+
+        return f.write(data)
+
+    def dumps(self) -> bytes:
+        """
+        Serialize model to bytes.
+
+        Returns
+        -------
+        bytes
+            Serialized model.
+        """
+        return pickle.dumps(self, protocol=PICKLE_PROTOCOL)
+
+    @classmethod
+    def load(cls, f: FileLike) -> tpe.Self:
+        """
+        Load model from file.
+
+        Parameters
+        ----------
+        f : str or Path or file-like object
+            Path to file or file-like object.
+
+        Returns
+        -------
+        model
+            Model instance.
+        """
+        if isinstance(f, (str, Path)):
+            data = Path(f).read_bytes()
+        else:
+            data = f.read()
+
+        return cls.loads(data)
+
+    @classmethod
+    def loads(cls, data: bytes) -> tpe.Self:
+        """
+        Load model from bytes.
+
+        Parameters
+        ----------
+        data : bytes
+            Serialized model.
+
+        Returns
+        -------
+        model
+            Model instance.
+
+        Raises
+        ------
+        TypeError
+            If loaded object is not a direct instance of model class.
+        """
+        loaded = pickle.loads(data)
+        if loaded.__class__ is not cls:
+            raise TypeError(f"Loaded object is not a direct instance of `{cls.__name__}`")
+        return loaded
+
     def fit(self: T, dataset: Dataset, *args: tp.Any, **kwargs: tp.Any) -> T:
         """
         Fit model.

tests/models/test_base.py

@@ -17,6 +17,8 @@
 import typing as tp
 import warnings
 from datetime import timedelta
+from pathlib import Path
+from tempfile import NamedTemporaryFile, TemporaryFile
 
 import numpy as np
 import pandas as pd
@@ -539,6 +541,42 @@ class MyModelWithoutConfig(ModelBase):
             MyModelWithoutConfig().get_config()
 
 
+class MyModel(ModelBase):
+    def __init__(self, x: int = 10, verbose: int = 0):
+        super().__init__(verbose=verbose)
+        self.x = x
+
+
+class TestSavingAndLoading:
+
+    @pytest.fixture()
+    def model(self) -> MyModel:
+        return MyModel()
+
+    def test_save_and_load_to_file(self, model: MyModel) -> None:
+        with TemporaryFile() as f:
+            model.save(f)
+            f.seek(0)
+            loaded_model = MyModel.load(f)
+        assert isinstance(loaded_model, MyModel)
+        assert loaded_model.__dict__ == model.__dict__
+
+    @pytest.mark.parametrize("use_str", (False, True))
+    def test_save_and_load_from_path(self, model: MyModel, use_str: bool) -> None:
+        with NamedTemporaryFile() as f:
+            path: tp.Union[Path, str] = Path(f.name) if not use_str else f.name
+            model.save(path)
+            loaded_model = MyModel.load(path)
+        assert isinstance(loaded_model, MyModel)
+        assert loaded_model.__dict__ == model.__dict__
+
+    def test_load_fails_on_incorrect_model_type(self, model: MyModel) -> None:
+        with NamedTemporaryFile() as f:
+            model.save(f.name)
+            with pytest.raises(TypeError, match="Loaded object is not a direct instance of `ModelBase`"):
+                ModelBase.load(f.name)
+
+
 class TestFixedColdRecoModelMixin:
     def test_cold_reco_works(self) -> None:
         class ColdRecoModel(FixedColdRecoModelMixin, ModelBase):

tests/models/test_dssm.py

@@ -25,7 +25,7 @@
 from rectools.models import DSSMModel
 from rectools.models.dssm import DSSM
 from rectools.models.vector import ImplicitRanker
-from tests.models.utils import assert_second_fit_refits_model
+from tests.models.utils import assert_dumps_loads_do_not_change_model, assert_second_fit_refits_model
 
 from .data import INTERACTIONS
 
@@ -338,3 +338,8 @@ def test_raises_when_no_features_in_dataset(self, dataset: Dataset, exclude_feat
     def test_second_fit_refits_model(self, dataset: Dataset) -> None:
         model = DSSMModel(deterministic=True)
         assert_second_fit_refits_model(model, dataset, pre_fit_callback=self._seed_everything)
+
+    def test_dumps_loads(self, dataset: Dataset) -> None:
+        model = DSSMModel()
+        model.fit(dataset)
+        assert_dumps_loads_do_not_change_model(model, dataset, check_configs=False)

tests/models/test_ease.py

@@ -25,6 +25,7 @@
 from .data import DATASET, INTERACTIONS
 from .utils import (
     assert_default_config_and_default_model_params_are_the_same,
+    assert_dumps_loads_do_not_change_model,
     assert_get_config_and_from_config_compatibility,
     assert_second_fit_refits_model,
 )
@@ -225,6 +226,11 @@ def test_i2i_with_warm_and_cold_items(self, item_features: tp.Optional[pd.DataFr
                 k=2,
             )
 
+    def test_dumps_loads(self, dataset: Dataset) -> None:
+        model = EASEModel()
+        model.fit(dataset)
+        assert_dumps_loads_do_not_change_model(model, dataset)
+
 
 class TestEASEModelConfiguration:
     def test_from_config(self) -> None:

tests/models/test_implicit_als.py

@@ -38,6 +38,7 @@
 from .data import DATASET
 from .utils import (
     assert_default_config_and_default_model_params_are_the_same,
+    assert_dumps_loads_do_not_change_model,
     assert_get_config_and_from_config_compatibility,
     assert_second_fit_refits_model,
 )
@@ -397,6 +398,11 @@ def test_per_epoch_fitting_consistent_with_regular_fitting(
         assert np.allclose(get_users_vectors(model_1.model), get_users_vectors(model_2.model))
         assert np.allclose(get_items_vectors(model_1.model), get_items_vectors(model_2.model))
 
+    def test_dumps_loads(self, use_gpu: bool, dataset: Dataset) -> None:
+        model = ImplicitALSWrapperModel(model=AlternatingLeastSquares(use_gpu=use_gpu))
+        model.fit(dataset)
+        assert_dumps_loads_do_not_change_model(model, dataset)
+
 
 class CustomALS(CPUAlternatingLeastSquares):
     pass

tests/models/test_implicit_knn.py

@@ -26,6 +26,7 @@
 from .data import DATASET, INTERACTIONS
 from .utils import (
     assert_default_config_and_default_model_params_are_the_same,
+    assert_dumps_loads_do_not_change_model,
     assert_get_config_and_from_config_compatibility,
     assert_second_fit_refits_model,
 )
@@ -251,6 +252,11 @@ def test_base_class(self, dataset: Dataset) -> None:
         ).astype({Columns.Score: np.float32})
         pd.testing.assert_frame_equal(actual, expected, atol=0.001)
 
+    def test_dumps_loads(self, dataset: Dataset) -> None:
+        model = ImplicitItemKNNWrapperModel(model=TFIDFRecommender())
+        model.fit(dataset)
+        assert_dumps_loads_do_not_change_model(model, dataset)
+
 
 class CustomKNN(ItemItemRecommender):
     pass

tests/models/test_lightfm.py

@@ -33,6 +33,7 @@
 from tests.models.data import DATASET
 from tests.models.utils import (
     assert_default_config_and_default_model_params_are_the_same,
+    assert_dumps_loads_do_not_change_model,
     assert_get_config_and_from_config_compatibility,
     assert_second_fit_refits_model,
 )
@@ -340,6 +341,11 @@ def _get_items_factors(self, dataset: Dataset) -> Factors:
                 filter_viewed=False,
             )
 
+    def test_dumps_loads(self, dataset: Dataset) -> None:
+        model = LightFMWrapperModel(LightFM())
+        model.fit(dataset)
+        assert_dumps_loads_do_not_change_model(model, dataset)
+
 
 class CustomLightFM(LightFM):
     pass

tests/models/test_popular.py

@@ -25,6 +25,7 @@
 from rectools.models.popular import Popularity
 from tests.models.utils import (
     assert_default_config_and_default_model_params_are_the_same,
+    assert_dumps_loads_do_not_change_model,
     assert_get_config_and_from_config_compatibility,
     assert_second_fit_refits_model,
 )
@@ -220,6 +221,11 @@ def test_second_fit_refits_model(self, dataset: Dataset) -> None:
         model = PopularModel()
         assert_second_fit_refits_model(model, dataset)
 
+    def test_dumps_loads(self, dataset: Dataset) -> None:
+        model = PopularModel()
+        model.fit(dataset)
+        assert_dumps_loads_do_not_change_model(model, dataset)
+
 
 class TestPopularModelConfiguration:
     @pytest.mark.parametrize(