fix flaws

Author: teodor-r

pyproject.toml

@@ -85,7 +85,6 @@ 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

@@ -22,15 +22,14 @@
 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
 
 from rectools import ExternalIds
 from rectools.dataset.dataset import Dataset, DatasetSchema, DatasetSchemaDict, IdMap
 from rectools.models.base import ErrorBehaviour, InternalRecoTriplet, ModelBase, ModelConfig
 from rectools.types import InternalIdsArray
-from rectools.utils.misc import get_class_or_function_full_path, import_object
+from rectools.utils.misc import get_class_or_function_full_path, import_object, make_dict_flat, unflatten_dict
 
 from ..item_net import (
     CatFeaturesItemNet,
@@ -605,31 +604,33 @@ def __setstate__(self, state: tp.Dict[str, tp.Any]) -> None:
     def load_from_checkpoint(
         cls,
         checkpoint_path: tp.Union[str, Path],
-        map_location: tp.Union[str, torch.device, None] = None,
-        config_update: tp.Dict[str, tp.Any] = {},
+        map_location: tp.Optional[tp.Union[str, torch.device]] = None,
+        model_params_update: tp.Optional[tp.Dict[str, tp.Any]] = None,
     ) -> tpe.Self:
         """Load model from Lightning checkpoint path.
 
         Parameters
         ----------
         checkpoint_path: Union[str, Path]
             Path to checkpoint location.
-        map_location: Union[str, torch.device, None], default None
+        map_location: Union[str, torch.device], optional
             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 '{}'
-            Contains custom values for checkpoint['hyper_parameters'].
-            Config_update has to be flattened with 'dot' reducer, before passed.
-
+        model_params_update: Dict[str, tp.Any], optional
+            Contains custom values for checkpoint['hyper_parameters']['model_config'].
+            Has to be flattened with 'dot' reducer, before passed.
+            You can use this argument to remove training-specific parameters that are not needed anymore.
+             e.g. 'get_trainer_func'
         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(config_update)
-        checkpoint["hyper_parameters"] = unflatten(prev_config_flatten, splitter="dot")
+        prev_model_config = checkpoint["hyper_parameters"]["model_config"]
+        if model_params_update:
+            prev_config_flatten = make_dict_flat(prev_model_config)
+            prev_config_flatten.update(model_params_update)
+            checkpoint["hyper_parameters"]["model_config"] = unflatten_dict(prev_config_flatten)
         loaded = cls._model_from_checkpoint(checkpoint)
         return loaded
 

tests/models/nn/transformers/test_base.py

@@ -19,7 +19,6 @@
 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
@@ -155,22 +154,25 @@ def test_save_load_for_fitted_model(
     @pytest.mark.parametrize("model_cls", (SASRecModel, BERT4RecModel))
     @pytest.mark.parametrize("map_location", ("cpu", torch.device("cuda:0"), None))
     @pytest.mark.parametrize(
-        "config_update",
+        "model_params_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",
-                }
+                "get_val_mask_func": "tests.models.nn.transformers.utils.leave_one_out_mask",
+                "get_trainer_func": "tests.models.nn.transformers.utils.custom_trainer",
             },
+            {
+                "get_val_mask_func": None,
+                "get_trainer_func": None,
+            },
+            None
         ),
     )
     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],
+        model_params_update: tp.Dict[str, tp.Any],
         request: FixtureRequest,
     ) -> None:
 
@@ -188,9 +190,8 @@ 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)
-        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
+            ckpt_path, map_location=map_location, model_params_update=model_params_update
         )
         assert isinstance(recovered_model, model_cls)