Feature/torch_ranker (#251)

`TorchRanker` to allow ranking with gpu using pytorch

Author: nsundalov

CHANGELOG.md

@@ -5,13 +5,17 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-
 ## Unreleased
 
 ### Added
 - `use_gpu` for PureSVD ([#229](https://github.com/MobileTeleSystems/RecTools/pull/229))
 - `from_params` method for models and `model_from_params` function ([#252](https://github.com/MobileTeleSystems/RecTools/pull/252))
+- `TorchRanker` ranker which calculates scores using torch. Supports GPU. [#251](https://github.com/MobileTeleSystems/RecTools/pull/251)
+- `Ranker` ranker protocol which unify rankers call. [#251](https://github.com/MobileTeleSystems/RecTools/pull/251)
+
+### Changed
 
+- `ImplicitRanker` `rank` method compatible with `Ranker` protocol. `use_gpu` and `num_threads` params moved from `rank` method to `__init__`. [#251](https://github.com/MobileTeleSystems/RecTools/pull/251)
 
 ## [0.10.0] - 16.01.2025
 

rectools/models/ease.py

@@ -85,7 +85,6 @@ def __init__(
         recommend_use_gpu_ranking: bool = True,
         verbose: int = 0,
     ):
-
         super().__init__(verbose=verbose)
         self.weight: np.ndarray
         self.regularization = regularization
@@ -146,16 +145,17 @@ def _recommend_u2i(
             distance=Distance.DOT,
             subjects_factors=user_items,
             objects_factors=self.weight,
+            use_gpu=self.recommend_use_gpu_ranking and HAS_CUDA,
+            num_threads=self.recommend_n_threads,
         )
+
         ui_csr_for_filter = user_items[user_ids] if filter_viewed else None
 
         all_user_ids, all_reco_ids, all_scores = ranker.rank(
             subject_ids=user_ids,
             k=k,
             filter_pairs_csr=ui_csr_for_filter,
             sorted_object_whitelist=sorted_item_ids_to_recommend,
-            num_threads=self.recommend_n_threads,
-            use_gpu=self.recommend_use_gpu_ranking and HAS_CUDA,
         )
 
         return all_user_ids, all_reco_ids, all_scores

rectools/models/rank/__init__.py

@@ -0,0 +1,43 @@
+#  Copyright 2022-2025 MTS (Mobile Telesystems)
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+# pylint: disable=wrong-import-position
+
+"""
+Recommendation models (:mod:`rectools.models.rank`)
+==============================================
+
+Rankers to build recs from embeddings.
+
+
+Rankers
+------
+`rank.ImplicitRanker`
+`rank.TorchRanker`
+"""
+
+try:
+    from .rank_torch import TorchRanker
+except ImportError:  # pragma: no cover
+    from .compat import TorchRanker  # type: ignore
+
+from rectools.models.rank.rank import Distance, Ranker
+from rectools.models.rank.rank_implicit import ImplicitRanker
+
+__all__ = [
+    "TorchRanker",
+    "ImplicitRanker",
+    "Distance",
+    "Ranker",
+]

rectools/models/rank/compat.py

@@ -0,0 +1,7 @@
+from rectools.compat import RequirementUnavailable
+
+
+class TorchRanker(RequirementUnavailable):
+    """Dummy class, which is returned if there are no dependencies required for the model"""
+
+    requirement = "torch"

rectools/models/rank/rank.py

@@ -0,0 +1,50 @@
+import typing as tp
+from enum import Enum
+
+from scipy import sparse
+
+from rectools import InternalIds
+from rectools.models.base import Scores
+from rectools.types import InternalIdsArray
+
+
+class Distance(Enum):
+    """Distance metric"""
+
+    DOT = 1  # Bigger value means closer vectors
+    COSINE = 2  # Bigger value means closer vectors
+    EUCLIDEAN = 3  # Smaller value means closer vectors
+
+
+class Ranker(tp.Protocol):
+    """Protocol for all rankers"""
+
+    def rank(
+        self,
+        subject_ids: InternalIds,
+        k: tp.Optional[int] = None,
+        filter_pairs_csr: tp.Optional[sparse.csr_matrix] = None,
+        sorted_object_whitelist: tp.Optional[InternalIdsArray] = None,
+    ) -> tp.Tuple[InternalIds, InternalIds, Scores]:  # pragma: no cover
+        """Rank objects by corresponding embeddings.
+
+        Parameters
+        ----------
+        subject_ids : InternalIds
+            Array of ids to recommend for.
+        k : int, optional, default ``None``
+            Derived number of recommendations for every subject id.
+            Return all recs if None.
+        filter_pairs_csr : sparse.csr_matrix, optional, default ``None``
+            Subject-object interactions that should be filtered from recommendations.
+            This is relevant for u2i case.
+        sorted_object_whitelist : sparse.csr_matrix, optional, default ``None``
+            Whitelist of object ids.
+            If given, only these items will be used for recommendations.
+            Otherwise all items from dataset will be used.
+
+        Returns
+        -------
+        (InternalIds, InternalIds, Scores)
+            Array of subject ids, array of recommended items, sorted by score descending and array of scores.
+        """

rectools/models/rank.py rectools/models/rank/rank_implicit.pyrectools/models/rank.py renamed to rectools/models/rank/rank_implicit.py

@@ -16,7 +16,6 @@
 
 import typing as tp
 import warnings
-from enum import Enum
 
 import implicit.cpu
 import implicit.gpu
@@ -27,18 +26,10 @@
 
 from rectools import InternalIds
 from rectools.models.base import Scores
+from rectools.models.rank.rank import Distance
+from rectools.models.utils import convert_arr_to_implicit_gpu_matrix
 from rectools.types import InternalIdsArray
 
-from .utils import convert_arr_to_implicit_gpu_matrix
-
-
-class Distance(Enum):
-    """Distance metric"""
-
-    DOT = 1  # Bigger value means closer vectors
-    COSINE = 2  # Bigger value means closer vectors
-    EUCLIDEAN = 3  # Smaller value means closer vectors
-
 
 class ImplicitRanker:
     """
@@ -58,18 +49,28 @@ class ImplicitRanker:
     objects_factors : np.ndarray
         Array with embeddings of all objects, shape (n_objects, n_factors).
         For item-item similarity models item similarity vectors are viewed as factors.
+    num_threads : int, default 0
+            Will be used as `num_threads` parameter for `implicit.cpu.topk.topk`. Omitted if use_gpu is True
+    use_gpu : bool, default False
+        If True `implicit.gpu.KnnQuery().topk` will be used instead of classic cpu version.
     """
 
     def __init__(
-        self, distance: Distance, subjects_factors: tp.Union[np.ndarray, sparse.csr_matrix], objects_factors: np.ndarray
+        self,
+        distance: Distance,
+        subjects_factors: tp.Union[np.ndarray, sparse.csr_matrix],
+        objects_factors: np.ndarray,
+        num_threads: int = 0,
+        use_gpu: bool = False,
     ) -> None:
-
         if isinstance(subjects_factors, sparse.csr_matrix) and distance != Distance.DOT:
             raise ValueError("To use `sparse.csr_matrix` distance must be `Distance.DOT`")
 
         self.distance = distance
         self.subjects_factors: np.ndarray = subjects_factors.astype(np.float32)
         self.objects_factors: np.ndarray = objects_factors.astype(np.float32)
+        self.num_threads = num_threads
+        self.use_gpu = use_gpu
 
         self.subjects_norms: np.ndarray
         if distance == Distance.COSINE:
@@ -85,7 +86,8 @@ def _get_neginf_score(self) -> float:
         # we're comparing `scores <= neginf_score`
         return float(
             np.asarray(
-                np.asarray(-np.finfo(np.float32).max, dtype=np.float32).view(np.uint32) - 1, dtype=np.uint32
+                np.asarray(-np.finfo(np.float32).max, dtype=np.float32).view(np.uint32) - 1,
+                dtype=np.uint32,
             ).view(np.float32)
         )
 
@@ -118,7 +120,6 @@ def _get_mask_for_correct_scores(self, scores: np.ndarray) -> tp.List[bool]:
     def _process_implicit_scores(
         self, subject_ids: InternalIds, ids: np.ndarray, scores: np.ndarray
     ) -> tp.Tuple[InternalIds, InternalIds, Scores]:
-
         all_target_ids = []
         all_reco_ids: tp.List[np.ndarray] = []
         all_scores: tp.List[np.ndarray] = []
@@ -152,7 +153,6 @@ def _rank_on_gpu(
         object_norms: tp.Optional[np.ndarray],
         filter_query_items: tp.Optional[tp.Union[sparse.csr_matrix, sparse.csr_array]],
     ) -> tp.Tuple[np.ndarray, np.ndarray]:  # pragma: no cover
-
         object_factors = convert_arr_to_implicit_gpu_matrix(object_factors)
 
         if isinstance(subject_factors, sparse.spmatrix):
@@ -184,19 +184,17 @@ def _rank_on_gpu(
     def rank(  # pylint: disable=too-many-branches
         self,
         subject_ids: InternalIds,
-        k: int,
+        k: tp.Optional[int] = None,
         filter_pairs_csr: tp.Optional[sparse.csr_matrix] = None,
         sorted_object_whitelist: tp.Optional[InternalIdsArray] = None,
-        num_threads: int = 0,
-        use_gpu: bool = False,
     ) -> tp.Tuple[InternalIds, InternalIds, Scores]:
         """Rank objects to proceed inference using implicit library topk cpu method.
 
         Parameters
         ----------
         subject_ids : csr_matrix
             Array of ids to recommend for.
-        k : int
+        k : int, optional, default ``None``
             Derived number of recommendations for every subject id.
         filter_pairs_csr : sparse.csr_matrix, optional, default ``None``
             Subject-object interactions that should be filtered from recommendations.
@@ -205,16 +203,16 @@ def rank(  # pylint: disable=too-many-branches
             Whitelist of object ids.
             If given, only these items will be used for recommendations.
             Otherwise all items from dataset will be used.
-        num_threads : int, default 0
-            Will be used as `num_threads` parameter for `implicit.cpu.topk.topk`. Omitted if use_gpu is True
-        use_gpu : bool, default False
-            If True `implicit.gpu.KnnQuery().topk` will be used instead of classic cpu version.
 
         Returns
         -------
         (InternalIds, InternalIds, Scores)
             Array of subject ids, array of recommended items, sorted by score descending and array of scores.
         """
+        if filter_pairs_csr is not None and filter_pairs_csr.shape[0] != len(subject_ids):
+            explanation = "Number of rows in `filter_pairs_csr` must be equal to `len(sublect_ids)`"
+            raise ValueError(explanation)
+
         if sorted_object_whitelist is not None:
             object_factors = self.objects_factors[sorted_object_whitelist]
 
@@ -229,6 +227,9 @@ def rank(  # pylint: disable=too-many-branches
             object_factors = self.objects_factors
             filter_query_items = filter_pairs_csr
 
+        if k is None:
+            k = object_factors.shape[0]
+
         subject_factors = self.subjects_factors[subject_ids]
 
         object_norms = None  # for DOT and EUCLIDEAN distance
@@ -243,6 +244,7 @@ def rank(  # pylint: disable=too-many-branches
 
         real_k = min(k, object_factors.shape[0])
 
+        use_gpu = self.use_gpu
         if use_gpu and not HAS_CUDA:
             warnings.warn("Forced rank() on CPU")
             use_gpu = False
@@ -263,7 +265,7 @@ def rank(  # pylint: disable=too-many-branches
                 item_norms=object_norms,  # query norms for COSINE distance are applied afterwards
                 filter_query_items=filter_query_items,  # queries x objects csr matrix for getting neginf scores
                 filter_items=None,  # rectools doesn't support blacklist for now
-                num_threads=num_threads,
+                num_threads=self.num_threads,
             )
 
         if sorted_object_whitelist is not None: