added test& edited pyproject

Author: teodor-r

pyproject.toml

@@ -85,6 +85,7 @@ ipywidgets = {version = ">=7.7,<8.2", optional = true}
 plotly = {version="^5.22.0", optional = true}
 nbformat = {version = ">=4.2.0", optional = true}
 cupy-cuda12x = {version = "^13.3.0", python = "<3.13", optional = true}
+flatten-dict = "^0.4.2"
 
 
 [tool.poetry.extras]

rectools/models/nn/transformers/base.py

@@ -18,11 +18,11 @@
 from copy import deepcopy
 from pathlib import Path
 from tempfile import NamedTemporaryFile
-from flatten_dict import flatten, unflatten
 
 import numpy as np
 import torch
 import typing_extensions as tpe
+from flatten_dict import flatten, unflatten
 from pydantic import BeforeValidator, PlainSerializer
 from pytorch_lightning import Trainer
 
@@ -608,8 +608,7 @@ def load_from_checkpoint(
         map_location: tp.Union[str, torch.device, None] = None,
         config_update: tp.Dict[str, tp.Any] = {},
     ) -> tpe.Self:
-        """
-        Load model from Lightning checkpoint path.
+        """Load model from Lightning checkpoint path.
 
         Parameters
         ----------
@@ -619,16 +618,18 @@ def load_from_checkpoint(
             Target device to load the checkpoint (e.g., 'cpu', 'cuda:0').
             If None, will use the device the checkpoint was saved on.
         config_update: tp.Dict[str, tp.Any], default '{}'
-            Сontains custom values for checkpoint['hyper_parameters']
+            Contains custom values for checkpoint['hyper_parameters'].
+            Config_update has to be flattened with 'dot' reducer, before passed.
+
         Returns
         -------
         Model instance.
         """
         checkpoint = torch.load(checkpoint_path, map_location=map_location, weights_only=False)
         prev_config = checkpoint["hyper_parameters"]
-        prev_config_flatten = flatten(prev_config, reducer='dot')
-        prev_config_flatten.update(flatten(config_update, reducer='dot'))
-        checkpoint["hyper_parameters"] = unflatten(prev_config_flatten, splitter='dot')
+        prev_config_flatten = flatten(prev_config, reducer="dot")
+        prev_config_flatten.update(config_update)
+        checkpoint["hyper_parameters"] = unflatten(prev_config_flatten, splitter="dot")
         loaded = cls._model_from_checkpoint(checkpoint)
         return loaded
 

tests/models/nn/transformers/test_base.py

@@ -19,6 +19,7 @@
 import pandas as pd
 import pytest
 import torch
+from flatten_dict import flatten
 from pytest import FixtureRequest
 from pytorch_lightning import Trainer, seed_everything
 from pytorch_lightning.loggers import CSVLogger
@@ -152,10 +153,24 @@ def test_save_load_for_fitted_model(
 
     @pytest.mark.parametrize("test_dataset", ("dataset", "dataset_item_features"))
     @pytest.mark.parametrize("model_cls", (SASRecModel, BERT4RecModel))
+    @pytest.mark.parametrize("map_location", ("cpu", torch.device("cuda:0"), None))
+    @pytest.mark.parametrize(
+        "config_update",
+        (
+            {
+                "model_config": {
+                    "get_val_mask_func": "tests.models.nn.transformers.utils.leave_one_out_mask",
+                    "get_trainer_func": "tests.models.nn.transformers.utils.custom_trainer",
+                }
+            },
+        ),
+    )
     def test_load_from_checkpoint(
         self,
         model_cls: tp.Type[TransformerModelBase],
         test_dataset: str,
+        map_location: tp.Union[str, torch.device, None],
+        config_update: tp.Dict[str, tp.Any],
         request: FixtureRequest,
     ) -> None:
 
@@ -173,7 +188,10 @@ def test_load_from_checkpoint(
             raise ValueError("No log dir")
         ckpt_path = os.path.join(model.fit_trainer.log_dir, "checkpoints", "last_epoch.ckpt")
         assert os.path.isfile(ckpt_path)
-        recovered_model = model_cls.load_from_checkpoint(ckpt_path)
+        config_update_flatten = flatten(config_update, reducer="dot")
+        recovered_model = model_cls.load_from_checkpoint(
+            ckpt_path, map_location=map_location, config_update=config_update_flatten
+        )
         assert isinstance(recovered_model, model_cls)
 
         self._assert_same_reco(model, recovered_model, dataset)