Feature/vector rank params

rectools/models/ease.py

@@ -18,6 +18,7 @@
 
 import numpy as np
 import typing_extensions as tpe
+from implicit.gpu import HAS_CUDA
 from scipy import sparse
 
 from rectools import InternalIds
@@ -34,6 +35,7 @@
 
     regularization: float = 500.0
     num_threads: int = 1
+    recommend_use_gpu_ranking: tp.Optional[bool] = None
 
 
 class EASEModel(ModelBase[EASEModelConfig]):
@@ -54,7 +56,12 @@
     verbose : int, default 0
         Degree of verbose output. If 0, no output will be provided.
     num_threads: int, default 1
-        Number of threads used for `recommend` method.
+        Number of threads used for recommendation ranking on cpu.
+    recommend_use_gpu_ranking: Optional[bool], default ``None``
+        Flag to use gpu for recommendation ranking. If ``None``, `implicit.gpu.HAS_CUDA` will be
+        checked before inference.
+        This attribute can be changed manually before calling model `recommend` method if you
+        want to change ranking behaviour.
     """
 
     recommends_for_warm = False
@@ -65,21 +72,33 @@
     def __init__(
         self,
         regularization: float = 500.0,
-        num_threads: int = 1,
+        num_threads: int = 1,  # TODO: decide. We already have it. But this is actually recommend_cpu_n_threads
         verbose: int = 0,
+        recommend_use_gpu_ranking: tp.Optional[bool] = None,
     ):
 
         super().__init__(verbose=verbose)
         self.weight: np.ndarray
         self.regularization = regularization
         self.num_threads = num_threads
+        self.recommend_use_gpu_ranking = recommend_use_gpu_ranking
 
     def _get_config(self) -> EASEModelConfig:
-        return EASEModelConfig(regularization=self.regularization, num_threads=self.num_threads, verbose=self.verbose)
+        return EASEModelConfig(
+            regularization=self.regularization,
+            num_threads=self.num_threads,
+            recommend_use_gpu_ranking=self.recommend_use_gpu_ranking,
+            verbose=self.verbose,
+        )
 
     @classmethod
     def _from_config(cls, config: EASEModelConfig) -> tpe.Self:
-        return cls(regularization=config.regularization, num_threads=config.num_threads, verbose=config.verbose)
+        return cls(
+            regularization=config.regularization,
+            recommend_use_gpu_ranking=config.recommend_use_gpu_ranking,
+            num_threads=config.num_threads,
+            verbose=config.verbose,
+        )
 
     def _fit(self, dataset: Dataset) -> None:  # type: ignore
         ui_csr = dataset.get_user_item_matrix(include_weights=True)
@@ -93,6 +112,13 @@
         self.weight = np.array(gram_matrix_inv / (-np.diag(gram_matrix_inv)))
         np.fill_diagonal(self.weight, 0.0)
 
+    @property
+    def _recommend_use_gpu_ranking(self) -> bool:
+        use_gpu = HAS_CUDA
+        if self.recommend_use_gpu_ranking is False:
+            use_gpu = False
+        return use_gpu
+
     def _recommend_u2i(
         self,
         user_ids: InternalIdsArray,
@@ -116,6 +142,7 @@
             filter_pairs_csr=ui_csr_for_filter,
             sorted_object_whitelist=sorted_item_ids_to_recommend,
             num_threads=self.num_threads,
+            use_gpu=self._recommend_use_gpu_ranking,
         )
 
         return all_user_ids, all_reco_ids, all_scores

rectools/models/implicit_als.py

@@ -104,6 +104,8 @@ class ImplicitALSWrapperModelConfig(ModelConfig):
 
     model: AlternatingLeastSquaresConfig
     fit_features_together: bool = False
+    recommend_cpu_n_threads: tp.Optional[int] = None
+    recommend_use_gpu_ranking: tp.Optional[bool] = None
 
 
 class ImplicitALSWrapperModel(VectorModel[ImplicitALSWrapperModelConfig]):
@@ -123,6 +125,17 @@ class ImplicitALSWrapperModel(VectorModel[ImplicitALSWrapperModelConfig]):
         Whether fit explicit features together with latent features or not.
         Used only if explicit features are present in dataset.
         See documentations linked above for details.
+    recommend_cpu_n_threads: Optional[int], default ``None``
+        Number of threads to use for recommendation ranking on cpu.
+        If ``None``, then number of threads will be set same as `model.num_threads`.
+        This attribute can be changed manually before calling model `recommend` method if you
+        want to change ranking behaviour.
+    recommend_use_gpu_ranking: Optional[bool], default ``None``
+        Flag to use gpu for recommendation ranking. If ``None``, then will be set same as
+        `model.use_gpu`.
+        `implicit.gpu.HAS_CUDA` will also be checked before inference.
+        This attribute can be changed manually before calling model `recommend` method if you
+        want to change ranking behaviour.
     """
 
     recommends_for_warm = False
@@ -133,22 +146,41 @@ class ImplicitALSWrapperModel(VectorModel[ImplicitALSWrapperModelConfig]):
 
     config_class = ImplicitALSWrapperModelConfig
 
-    def __init__(self, model: AnyAlternatingLeastSquares, verbose: int = 0, fit_features_together: bool = False):
-        self._config = self._make_config(model, verbose, fit_features_together)
+    def __init__(
+        self,
+        model: AnyAlternatingLeastSquares,
+        verbose: int = 0,
+        fit_features_together: bool = False,
+        recommend_cpu_n_threads: tp.Optional[int] = None,
+        recommend_use_gpu_ranking: tp.Optional[bool] = None,
+    ):
+        self._config = self._make_config(
+            model, verbose, fit_features_together, recommend_cpu_n_threads, recommend_use_gpu_ranking
+        )
 
         super().__init__(verbose=verbose)
 
         self.model: AnyAlternatingLeastSquares
         self._model = model  # for refit
 
         self.fit_features_together = fit_features_together
-        self.use_gpu = isinstance(model, GPUAlternatingLeastSquares)
-        if not self.use_gpu:
-            self.n_threads = model.num_threads
+
+        if recommend_cpu_n_threads is None and isinstance(model, CPUAlternatingLeastSquares):
+            recommend_cpu_n_threads = model.num_threads
+        self.recommend_cpu_n_threads = recommend_cpu_n_threads
+
+        if recommend_use_gpu_ranking is None:
+            recommend_use_gpu_ranking = isinstance(model, GPUAlternatingLeastSquares)
+        self.recommend_use_gpu_ranking = recommend_use_gpu_ranking
 
     @classmethod
     def _make_config(
-        cls, model: AnyAlternatingLeastSquares, verbose: int, fit_features_together: bool
+        cls,
+        model: AnyAlternatingLeastSquares,
+        verbose: int,
+        fit_features_together: bool,
+        recommend_cpu_n_threads: tp.Optional[int] = None,
+        recommend_use_gpu_ranking: tp.Optional[bool] = None,
     ) -> ImplicitALSWrapperModelConfig:
         params = {
             "factors": model.factors,
@@ -183,6 +215,8 @@ def _make_config(
             ),
             verbose=verbose,
             fit_features_together=fit_features_together,
+            recommend_cpu_n_threads=recommend_cpu_n_threads,
+            recommend_use_gpu_ranking=recommend_use_gpu_ranking,
         )
 
     def _get_config(self) -> ImplicitALSWrapperModelConfig:
@@ -195,7 +229,13 @@ def _from_config(cls, config: ImplicitALSWrapperModelConfig) -> tpe.Self:
         else:
             model_cls = config.model.cls
         model = model_cls(**config.model.params)
-        return cls(model=model, verbose=config.verbose, fit_features_together=config.fit_features_together)
+        return cls(
+            model=model,
+            verbose=config.verbose,
+            fit_features_together=config.fit_features_together,
+            recommend_cpu_n_threads=config.recommend_cpu_n_threads,
+            recommend_use_gpu_ranking=config.recommend_use_gpu_ranking,
+        )
 
     def _fit(self, dataset: Dataset) -> None:
         self.model = deepcopy(self._model)

rectools/models/lightfm.py

@@ -93,6 +93,7 @@ class LightFMWrapperModelConfig(ModelConfig):
     model: LightFMConfig
     epochs: int = 1
     num_threads: int = 1
+    recommend_use_gpu_ranking: tp.Optional[bool] = None
 
 
 class LightFMWrapperModel(FixedColdRecoModelMixin, VectorModel[LightFMWrapperModelConfig]):
@@ -115,6 +116,12 @@ class LightFMWrapperModel(FixedColdRecoModelMixin, VectorModel[LightFMWrapperMod
         Will be used as `num_threads` parameter for `LightFM.fit`.
     verbose : int, default 0
         Degree of verbose output. If 0, no output will be provided.
+    recommend_use_gpu_ranking: Optional[bool], default ``None``
+        Flag to use gpu for recommendation ranking. If ``None``, then will be set same as
+        `model.use_gpu`.
+        `implicit.gpu.HAS_CUDA` will also be checked before inference.
+        This attribute can be changed manually before calling model `recommend` method if you
+        want to change ranking behaviour.
     """
 
     recommends_for_warm = True
@@ -129,15 +136,18 @@ def __init__(
         self,
         model: LightFM,
         epochs: int = 1,
-        num_threads: int = 1,
+        num_threads: int = 1,  # TODO: decide. this is used for both fit n_threads and ranker n_threads
         verbose: int = 0,
+        recommend_use_gpu_ranking: tp.Optional[bool] = None,
     ):
         super().__init__(verbose=verbose)
 
         self.model: LightFM
         self._model = model
         self.n_epochs = epochs
         self.n_threads = num_threads
+        self.recommend_cpu_n_threads = num_threads  # TODO: not consistent with VectorModel parent behaviour
+        self.recommend_use_gpu_ranking = recommend_use_gpu_ranking
 
     def _get_config(self) -> LightFMWrapperModelConfig:
         inner_model = self._model
@@ -164,6 +174,7 @@ def _get_config(self) -> LightFMWrapperModelConfig:
             epochs=self.n_epochs,
             num_threads=self.n_threads,
             verbose=self.verbose,
+            recommend_use_gpu_ranking=self.recommend_use_gpu_ranking,
         )
 
     @classmethod

rectools/models/pure_svd.py

@@ -34,6 +34,8 @@ class PureSVDModelConfig(ModelConfig):
     tol: float = 0
     maxiter: tp.Optional[int] = None
     random_state: tp.Optional[int] = None
+    recommend_cpu_n_threads: tp.Optional[int] = None
+    recommend_use_gpu_ranking: tp.Optional[bool] = None
 
 
 class PureSVDModel(VectorModel[PureSVDModelConfig]):
@@ -54,6 +56,17 @@ class PureSVDModel(VectorModel[PureSVDModelConfig]):
         Pseudorandom number generator state used to generate resamples.
     verbose : int, default ``0``
         Degree of verbose output. If ``0``, no output will be provided.
+    recommend_cpu_n_threads: Optional[int], default ``None``
+        Number of threads to use for recommendation ranking on cpu.
+        If ``None``, then number of threads will be set same as `model.num_threads`.
+        This attribute can be changed manually before calling model `recommend` method if you
+        want to change ranking behaviour.
+    recommend_use_gpu_ranking: Optional[bool], default ``None``
+        Flag to use gpu for recommendation ranking. If ``None``, then will be set same as
+        `model.use_gpu`.
+        `implicit.gpu.HAS_CUDA` will also be checked before inference.
+        This attribute can be changed manually before calling model `recommend` method if you
+        want to change ranking behaviour.
     """
 
     recommends_for_warm = False
@@ -71,13 +84,17 @@ def __init__(
         maxiter: tp.Optional[int] = None,
         random_state: tp.Optional[int] = None,
         verbose: int = 0,
+        recommend_cpu_n_threads: tp.Optional[int] = None,
+        recommend_use_gpu_ranking: tp.Optional[bool] = None,
     ):
         super().__init__(verbose=verbose)
 
         self.factors = factors
         self.tol = tol
         self.maxiter = maxiter
         self.random_state = random_state
+        self.recommend_cpu_n_threads = recommend_cpu_n_threads
+        self.recommend_use_gpu_ranking = recommend_use_gpu_ranking
 
         self.user_factors: np.ndarray
         self.item_factors: np.ndarray
@@ -89,6 +106,8 @@ def _get_config(self) -> PureSVDModelConfig:
             maxiter=self.maxiter,
             random_state=self.random_state,
             verbose=self.verbose,
+            recommend_cpu_n_threads=self.recommend_cpu_n_threads,
+            recommend_use_gpu_ranking=self.recommend_use_gpu_ranking,
         )
 
     @classmethod
@@ -99,6 +118,8 @@ def _from_config(cls, config: PureSVDModelConfig) -> tpe.Self:
             maxiter=config.maxiter,
             random_state=config.random_state,
             verbose=config.verbose,
+            recommend_cpu_n_threads=config.recommend_cpu_n_threads,
+            recommend_use_gpu_ranking=config.recommend_use_gpu_ranking,
         )
 
     def _fit(self, dataset: Dataset) -> None:  # type: ignore

rectools/models/vector.py

@@ -18,6 +18,7 @@
 
 import attr
 import numpy as np
+from implicit.gpu import HAS_CUDA
 
 from rectools import InternalIds
 from rectools.dataset import Dataset
@@ -40,7 +41,20 @@ class VectorModel(ModelBase[ModelConfig_T]):
 
     u2i_dist: Distance = NotImplemented
     i2i_dist: Distance = NotImplemented
-    n_threads: int = 0  # TODO: decide how to pass it correctly for all models
+    recommend_cpu_n_threads: tp.Optional[int] = None  # TODO: decide. Lightfm has num_threads
+    recommend_use_gpu_ranking: tp.Optional[bool] = None
+
+    @property
+    def _recommend_use_gpu_ranking(self) -> bool:
+        use_gpu = HAS_CUDA
+        if self.recommend_use_gpu_ranking is False:
+            use_gpu = False
+        return use_gpu
+
+    @property
+    def _recommend_cpu_n_threads(self) -> int:
+        num_threads = 0 if self.recommend_cpu_n_threads is None else self.recommend_cpu_n_threads
+        return num_threads
 
     def _recommend_u2i(
         self,
@@ -65,7 +79,8 @@ def _recommend_u2i(
             k=k,
             filter_pairs_csr=ui_csr_for_filter,
             sorted_object_whitelist=sorted_item_ids_to_recommend,
-            num_threads=self.n_threads,
+            num_threads=self._recommend_cpu_n_threads,
+            use_gpu=self._recommend_use_gpu_ranking,
         )
 
     def _recommend_i2i(
@@ -84,7 +99,8 @@ def _recommend_i2i(
             k=k,
             filter_pairs_csr=None,
             sorted_object_whitelist=sorted_item_ids_to_recommend,
-            num_threads=self.n_threads,
+            num_threads=self._recommend_cpu_n_threads,
+            use_gpu=self._recommend_use_gpu_ranking,
         )
 
     def _process_biases_to_vectors(

tests/models/test_ease.py

@@ -249,22 +249,15 @@ def test_get_config(self) -> None:
             regularization=500,
             num_threads=1,
             verbose=1,
+            recommend_use_gpu_ranking=None,
         )
         config = model.get_config()
-        expected = {
-            "regularization": 500,
-            "num_threads": 1,
-            "verbose": 1,
-        }
+        expected = {"regularization": 500, "num_threads": 1, "verbose": 1, "recommend_use_gpu_ranking": None}
         assert config == expected
 
     @pytest.mark.parametrize("simple_types", (False, True))
     def test_get_config_and_from_config_compatibility(self, simple_types: bool) -> None:
-        initial_config = {
-            "regularization": 500,
-            "num_threads": 1,
-            "verbose": 1,
-        }
+        initial_config = {"regularization": 500, "num_threads": 1, "verbose": 1, "recommend_use_gpu_ranking": True}
         assert_get_config_and_from_config_compatibility(EASEModel, DATASET, initial_config, simple_types)
 
     def test_default_config_and_default_model_params_are_the_same(self) -> None:

tests/models/test_implicit_als.py

@@ -496,6 +496,8 @@ def test_to_config(self, use_gpu: bool, random_state: tp.Optional[int], simple_t
             },
             "fit_features_together": True,
             "verbose": 1,
+            "recommend_use_gpu_ranking": None,
+            "recommend_cpu_n_threads": None,
         }
         assert config == expected
 

tests/models/test_lightfm.py

@@ -408,6 +408,7 @@ def test_to_config(self, random_state: tp.Optional[int], simple_types: bool) ->
             "epochs": 2,
             "num_threads": 3,
             "verbose": 1,
+            "recommend_use_gpu_ranking": None,
         }
         assert config == expected