Feature/sasrec torch rank

rectools/models/nn/bert4rec.py

@@ -268,14 +268,16 @@ class BERT4RecModel(TransformerModelBase[BERT4RecModelConfig]):
         When set to ``None``, "cuda" will be used if it is available, "cpu" otherwise.
         If you want to change this parameter after model is initialized,
         you can manually assign new value to model `recommend_device` attribute.
+    recommend_use_torch_ranking : bool, default ``True``
+        Use `TorchRanker` for items ranking while preparing recommendations.
+        If set to ``False``, use `ImplicitRanker` instead.
+        If you want to change this parameter after model is initialized,
+        you can manually assign new value to model `recommend_use_torch_ranking` attribute.
     recommend_n_threads : int, default 0
-        Number of threads to use in ranker if GPU ranking is turned off or unavailable.
+        Number of threads to use for `ImplicitRanker`. Omitted if `recommend_use_torch_ranking` is
+        set to ``True`` (default).
         If you want to change this parameter after model is initialized,
         you can manually assign new value to model `recommend_n_threads` attribute.
-    recommend_use_gpu_ranking : bool, default ``True``
-        If ``True`` and HAS_CUDA ``True``, set use_gpu=True in ImplicitRanker.rank.
-        If you want to change this parameter after model is initialized,
-        you can manually assign new value to model `recommend_use_gpu_ranking` attribute.
     """
 
     config_class = BERT4RecModelConfig
@@ -311,8 +313,8 @@ def __init__(  # pylint: disable=too-many-arguments, too-many-locals
         get_trainer_func: tp.Optional[TrainerCallable] = None,
         recommend_batch_size: int = 256,
         recommend_device: tp.Optional[str] = None,
+        recommend_use_torch_ranking: bool = True,
         recommend_n_threads: int = 0,
-        recommend_use_gpu_ranking: bool = True,  # TODO: remove after TorchRanker
     ):
         self.mask_prob = mask_prob
 
@@ -339,7 +341,7 @@ def __init__(  # pylint: disable=too-many-arguments, too-many-locals
             recommend_batch_size=recommend_batch_size,
             recommend_device=recommend_device,
             recommend_n_threads=recommend_n_threads,
-            recommend_use_gpu_ranking=recommend_use_gpu_ranking,
+            recommend_use_torch_ranking=recommend_use_torch_ranking,
             train_min_user_interactions=train_min_user_interactions,
             item_net_block_types=item_net_block_types,
             item_net_constructor_type=item_net_constructor_type,

rectools/models/nn/sasrec.py

@@ -352,14 +352,16 @@ class SASRecModel(TransformerModelBase[SASRecModelConfig]):
         When set to ``None``, "cuda" will be used if it is available, "cpu" otherwise.
         If you want to change this parameter after model is initialized,
         you can manually assign new value to model `recommend_device` attribute.
+    recommend_use_torch_ranking : bool, default ``True``
+        Use `TorchRanker` for items ranking while preparing recommendations.
+        If set to ``False``, use `ImplicitRanker` instead.
+        If you want to change this parameter after model is initialized,
+        you can manually assign new value to model `recommend_use_torch_ranking` attribute.
     recommend_n_threads : int, default 0
-        Number of threads to use in ranker if GPU ranking is turned off or unavailable.
+        Number of threads to use for `ImplicitRanker`. Omitted if `recommend_use_torch_ranking` is
+        set to ``True`` (default).
         If you want to change this parameter after model is initialized,
         you can manually assign new value to model `recommend_n_threads` attribute.
-    recommend_use_gpu_ranking : bool, default ``True``
-        If ``True`` and HAS_CUDA ``True``, set use_gpu=True in ImplicitRanker.rank.
-        If you want to change this parameter after model is initialized,
-        you can manually assign new value to model `recommend_use_gpu_ranking` attribute.
     """
 
     config_class = SASRecModelConfig
@@ -394,8 +396,8 @@ def __init__(  # pylint: disable=too-many-arguments, too-many-locals
         get_trainer_func: tp.Optional[TrainerCallable] = None,
         recommend_batch_size: int = 256,
         recommend_device: tp.Optional[str] = None,
+        recommend_use_torch_ranking: bool = True,
         recommend_n_threads: int = 0,
-        recommend_use_gpu_ranking: bool = True,  # TODO: remove after TorchRanker
     ):
         super().__init__(
             transformer_layers_type=transformer_layers_type,
@@ -420,7 +422,7 @@ def __init__(  # pylint: disable=too-many-arguments, too-many-locals
             recommend_batch_size=recommend_batch_size,
             recommend_device=recommend_device,
             recommend_n_threads=recommend_n_threads,
-            recommend_use_gpu_ranking=recommend_use_gpu_ranking,
+            recommend_use_torch_ranking=recommend_use_torch_ranking,
             train_min_user_interactions=train_min_user_interactions,
             item_net_block_types=item_net_block_types,
             item_net_constructor_type=item_net_constructor_type,

rectools/models/nn/transformer_base.py

@@ -176,7 +176,7 @@ class TransformerModelConfig(ModelConfig):
     recommend_batch_size: int = 256
     recommend_device: tp.Optional[str] = None
     recommend_n_threads: int = 0
-    recommend_use_gpu_ranking: bool = True  # TODO: remove after TorchRanker
+    recommend_use_torch_ranking: bool = True
     train_min_user_interactions: int = 2
     item_net_block_types: ItemNetBlockTypes = (IdEmbeddingsItemNet, CatFeaturesItemNet)
     item_net_constructor_type: ItemNetConstructorType = SumOfEmbeddingsConstructor
@@ -228,7 +228,7 @@ def __init__(  # pylint: disable=too-many-arguments, too-many-locals
         recommend_batch_size: int = 256,
         recommend_device: tp.Optional[str] = None,
         recommend_n_threads: int = 0,
-        recommend_use_gpu_ranking: bool = True,  # TODO: remove after TorchRanker
+        recommend_use_torch_ranking: bool = True,
         train_min_user_interactions: int = 2,
         item_net_block_types: tp.Sequence[tp.Type[ItemNetBase]] = (IdEmbeddingsItemNet, CatFeaturesItemNet),
         item_net_constructor_type: tp.Type[ItemNetConstructorBase] = SumOfEmbeddingsConstructor,
@@ -260,7 +260,7 @@ def __init__(  # pylint: disable=too-many-arguments, too-many-locals
         self.recommend_batch_size = recommend_batch_size
         self.recommend_device = recommend_device
         self.recommend_n_threads = recommend_n_threads
-        self.recommend_use_gpu_ranking = recommend_use_gpu_ranking
+        self.recommend_use_torch_ranking = recommend_use_torch_ranking
         self.train_min_user_interactions = train_min_user_interactions
         self.item_net_block_types = item_net_block_types
         self.item_net_constructor_type = item_net_constructor_type
@@ -402,15 +402,15 @@ def _recommend_u2i(
             sorted_item_ids_to_recommend = self.data_preparator.get_known_items_sorted_internal_ids()  # model internal
 
         recommend_dataloader = self.data_preparator.get_dataloader_recommend(dataset, self.recommend_batch_size)
-        return self.lightning_model.recommend_u2i(
+        return self.lightning_model._recommend_u2i(  # pylint: disable=protected-access
             user_ids=user_ids,
             recommend_dataloader=recommend_dataloader,
             sorted_item_ids_to_recommend=sorted_item_ids_to_recommend,
             k=k,
             filter_viewed=filter_viewed,
             dataset=dataset,
             recommend_n_threads=self.recommend_n_threads,
-            recommend_use_gpu_ranking=self.recommend_use_gpu_ranking,
+            recommend_use_torch_ranking=self.recommend_use_torch_ranking,
             recommend_device=self.recommend_device,
         )
 
@@ -424,12 +424,12 @@ def _recommend_i2i(
         if sorted_item_ids_to_recommend is None:
             sorted_item_ids_to_recommend = self.data_preparator.get_known_items_sorted_internal_ids()
 
-        return self.lightning_model.recommend_i2i(
+        return self.lightning_model._recommend_i2i(  # pylint: disable=protected-access
             target_ids=target_ids,
             sorted_item_ids_to_recommend=sorted_item_ids_to_recommend,
             k=k,
             recommend_n_threads=self.recommend_n_threads,
-            recommend_use_gpu_ranking=self.recommend_use_gpu_ranking,
+            recommend_use_torch_ranking=self.recommend_use_torch_ranking,
             recommend_device=self.recommend_device,
         )
 

rectools/models/nn/transformer_lightning.py

@@ -16,14 +16,13 @@
 
 import numpy as np
 import torch
-from implicit.gpu import HAS_CUDA
 from pytorch_lightning import LightningModule
 from torch.utils.data import DataLoader
 
 from rectools import ExternalIds
 from rectools.dataset.dataset import Dataset, DatasetSchemaDict
 from rectools.models.base import InternalRecoTriplet
-from rectools.models.rank import Distance, ImplicitRanker
+from rectools.models.rank import Distance, ImplicitRanker, Ranker, TorchRanker
 from rectools.types import InternalIdsArray
 
 from .transformer_backbone import TransformerTorchBackbone
@@ -102,27 +101,30 @@ def validation_step(self, batch: tp.Dict[str, torch.Tensor], batch_idx: int) ->
         """Validate step."""
         raise NotImplementedError()
 
-    def recommend_u2i(
+    def _recommend_u2i(
         self,
         user_ids: InternalIdsArray,
         recommend_dataloader: DataLoader,
         sorted_item_ids_to_recommend: InternalIdsArray,
         k: int,
         dataset: Dataset,  # [n_rec_users x n_items + n_item_extra_tokens]
         filter_viewed: bool,
+        recommend_use_torch_ranking: bool,
+        recommend_n_threads: int,
+        recommend_device: tp.Optional[str],
         *args: tp.Any,
         **kwargs: tp.Any,
     ) -> InternalRecoTriplet:
         """Recommending to users."""
         raise NotImplementedError()
 
-    def recommend_i2i(
+    def _recommend_i2i(
         self,
         target_ids: InternalIdsArray,
         sorted_item_ids_to_recommend: InternalIdsArray,
         k: int,
+        recommend_use_torch_ranking: bool,
         recommend_n_threads: int,
-        recommend_use_gpu_ranking: bool,
         recommend_device: tp.Optional[str],
         *args: tp.Any,
         **kwargs: tp.Any,
@@ -296,12 +298,15 @@ def _prepare_for_inference(self, recommend_device: tp.Optional[str]) -> None:
         self.torch_model.to(device)
         self.torch_model.eval()
 
-    def get_user_item_embeddings(
+    def _get_user_item_embeddings(
         self,
         recommend_dataloader: DataLoader,
         recommend_device: tp.Optional[str],
     ) -> tp.Tuple[torch.Tensor, torch.Tensor]:
-        """TODO."""
+        """
+        Prepare user embeddings for all user interaction sequences in `recommend_dataloader`.
+        Prepare item embeddings for full items catalog.
+        """
         self._prepare_for_inference(recommend_device)
         device = self.torch_model.item_model.device
 
@@ -314,71 +319,86 @@ def get_user_item_embeddings(
 
         return torch.cat(user_embs), item_embs
 
-    def recommend_u2i(
+    def _recommend_u2i(
         self,
         user_ids: InternalIdsArray,
         recommend_dataloader: DataLoader,
         sorted_item_ids_to_recommend: InternalIdsArray,
         k: int,
         dataset: Dataset,  # [n_rec_users x n_items + n_item_extra_tokens]
         filter_viewed: bool,
+        recommend_use_torch_ranking: bool,
         recommend_n_threads: int,
-        recommend_use_gpu_ranking: bool,
         recommend_device: tp.Optional[str],
     ) -> InternalRecoTriplet:
         """Recommend to users."""
         ui_csr_for_filter = None
         if filter_viewed:
-            ui_csr_for_filter = dataset.get_user_item_matrix(include_weights=False)[user_ids]
-
-        user_embs, item_embs = self.get_user_item_embeddings(recommend_dataloader, recommend_device)
-        user_embs_np = user_embs.detach().cpu().numpy()
-        item_embs_np = item_embs.detach().cpu().numpy()
-        del user_embs, item_embs
-        torch.cuda.empty_cache()
-
-        ranker = ImplicitRanker(
-            Distance.DOT,
-            user_embs_np[user_ids],  # [n_rec_users, n_factors]
-            item_embs_np,  # [n_items + n_item_extra_tokens, n_factors]
-            num_threads=recommend_n_threads,
-            use_gpu=recommend_use_gpu_ranking and HAS_CUDA,
-        )
+            ui_csr_for_filter = dataset.get_user_item_matrix(include_weights=False, include_warm_items=True)[user_ids]
+
+        user_embs, item_embs = self._get_user_item_embeddings(recommend_dataloader, recommend_device)
+
+        ranker: Ranker
+        if recommend_use_torch_ranking:
+            ranker = TorchRanker(
+                distance=Distance.DOT,
+                device=item_embs.device,
+                subjects_factors=user_embs[user_ids],
+                objects_factors=item_embs,
+            )
+        else:
+            user_embs_np = user_embs.detach().cpu().numpy()
+            item_embs_np = item_embs.detach().cpu().numpy()
+            del user_embs, item_embs
+            torch.cuda.empty_cache()
+            ranker = ImplicitRanker(
+                Distance.DOT,
+                user_embs_np[user_ids],  # [n_rec_users, n_factors]
+                item_embs_np,  # [n_items + n_item_extra_tokens, n_factors]
+                num_threads=recommend_n_threads,
+                use_gpu=False,
+            )
 
-        # TODO: We should test if torch `topk`` is faster when `filter_viewed`` is ``False``
         user_ids_indices, all_reco_ids, all_scores = ranker.rank(
-            subject_ids=np.arange(user_embs_np.shape[0]),  # n_rec_users
+            subject_ids=np.arange(len(user_ids)),  # n_rec_users
             k=k,
             filter_pairs_csr=ui_csr_for_filter,  # [n_rec_users x n_items + n_item_extra_tokens]
             sorted_object_whitelist=sorted_item_ids_to_recommend,  # model_internal
         )
         all_user_ids = user_ids[user_ids_indices]
         return all_user_ids, all_reco_ids, all_scores
 
-    def recommend_i2i(
+    def _recommend_i2i(
         self,
         target_ids: InternalIdsArray,
         sorted_item_ids_to_recommend: InternalIdsArray,
         k: int,
+        recommend_use_torch_ranking: bool,
         recommend_n_threads: int,
-        recommend_use_gpu_ranking: bool,
         recommend_device: tp.Optional[str],
     ) -> InternalRecoTriplet:
         """Recommend to items."""
         self._prepare_for_inference(recommend_device)
-
         with torch.no_grad():
-            item_embs = self.torch_model.item_model.get_all_embeddings().detach().cpu().numpy()
-        # TODO: i2i recommendations do not need filtering viewed and user most of the times has GPU
-        # We should test if torch `topk`` is faster
-
-        ranker = ImplicitRanker(
-            self.i2i_dist,
-            item_embs,  # [n_items + n_item_extra_tokens, n_factors]
-            item_embs,  # [n_items + n_item_extra_tokens, n_factors]
-            num_threads=recommend_n_threads,
-            use_gpu=recommend_use_gpu_ranking and HAS_CUDA,
-        )
+            item_embs = self.torch_model.item_model.get_all_embeddings()
+
+        ranker: Ranker
+        if recommend_use_torch_ranking:
+            ranker = TorchRanker(
+                distance=self.i2i_dist, device=item_embs.device, subjects_factors=item_embs, objects_factors=item_embs
+            )
+        else:
+            item_embs_np = item_embs.detach().cpu().numpy()
+            del item_embs
+            ranker = ImplicitRanker(
+                self.i2i_dist,
+                item_embs_np,  # [n_items + n_item_extra_tokens, n_factors]
+                item_embs_np,  # [n_items + n_item_extra_tokens, n_factors]
+                num_threads=recommend_n_threads,
+                use_gpu=False,
+            )
+
+        torch.cuda.empty_cache()
         return ranker.rank(
             subject_ids=target_ids,  # model internal
             k=k,

rectools/models/rank/rank_implicit.py

@@ -167,7 +167,10 @@ def _rank_on_gpu(
             object_norms = convert_arr_to_implicit_gpu_matrix(object_norms)
 
         if filter_query_items is not None:
-            filter_query_items = implicit.gpu.COOMatrix(filter_query_items.tocoo())
+            if filter_query_items.count_nonzero() > 0:
+                filter_query_items = implicit.gpu.COOMatrix(filter_query_items.tocoo())
+            else:  # can't create `implicit.gpu.COOMatrix` for all zeroes
+                filter_query_items = None
 
         ids, scores = implicit.gpu.KnnQuery().topk(  # pylint: disable=c-extension-no-member
             items=object_factors,

rectools/models/rank/rank_torch.py

@@ -138,7 +138,7 @@ def rank(
                         torch.from_numpy(filter_pairs_csr[cur_user_emb_inds].toarray()[:, sorted_object_whitelist]).to(
                             scores.device
                         )
-                        == 1
+                        != 0
                     )
                     scores = torch.masked_fill(scores, mask, mask_values)
 

tests/dataset/test_dataset.py

@@ -518,7 +518,7 @@ class TestSerializeFeatureName:
     def test_basic(self, feature_name: AnyFeatureName, expected: Hashable) -> None:
         assert _serialize_feature_name(feature_name) == expected
 
-    @pytest.mark.parametrize("feature_name", (np.array([1]), [1], np.array(["name"]), np.array([True])))
+    @pytest.mark.parametrize("feature_name", (datetime.now(), np.array([1]), [1], np.array(["name"]), np.array([True])))
     def test_raises_on_incorrect_input(self, feature_name: tp.Any) -> None:
         with pytest.raises(TypeError):
             _serialize_feature_name(feature_name)