replay-rec 0.16.0__py3-none-any.whl → 0.17.0__py3-none-any.whl

replay/models/nn/sequential/callbacks/validation_callback.py

@@ -1,13 +1,12 @@
-from typing import Any, List, Optional, Protocol, Tuple, Literal
+from typing import Any, List, Literal, Optional, Protocol, Tuple
 
-import lightning as L
+import lightning
 import torch
 from lightning.pytorch.utilities.rank_zero import rank_zero_only
 
 from replay.metrics.torch_metrics_builder import TorchMetricsBuilder, metrics_to_df
 from replay.models.nn.sequential.postprocessors import BasePostProcessor
 
-
 CallbackMetricName = Literal[
     "recall",
     "precision",
@@ -19,17 +18,17 @@ CallbackMetricName = Literal[
 ]
 
 
-# pylint: disable=too-few-public-methods
 class ValidationBatch(Protocol):
     """
     Validation callback batch
     """
+
     query_id: torch.LongTensor
     ground_truth: torch.LongTensor
     train: torch.LongTensor
 
 
-class ValidationMetricsCallback(L.Callback):
+class ValidationMetricsCallback(lightning.Callback):
     """
     Callback for validation and testing stages.
 
@@ -37,7 +36,6 @@ class ValidationMetricsCallback(L.Callback):
     the suffix of the metric name will contain the serial number of the dataloader.
     """
 
-    # pylint: disable=invalid-name
     def __init__(
         self,
         metrics: Optional[List[CallbackMetricName]] = None,
@@ -63,8 +61,9 @@ class ValidationMetricsCallback(L.Callback):
             return [len(dataloaders)]
         return [len(dataloader) for dataloader in dataloaders]
 
-    # pylint: disable=unused-argument
-    def on_validation_epoch_start(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None:
+    def on_validation_epoch_start(
+        self, trainer: lightning.Trainer, pl_module: lightning.LightningModule  # noqa: ARG002
+    ) -> None:
         self._dataloaders_size = self._get_dataloaders_size(trainer.val_dataloaders)
         self._metrics_builders = [
             TorchMetricsBuilder(self._metrics, self._ks, self._item_count) for _ in self._dataloaders_size
@@ -72,8 +71,11 @@ class ValidationMetricsCallback(L.Callback):
         for builder in self._metrics_builders:
             builder.reset()
 
-    # pylint: disable=unused-argument
-    def on_test_epoch_start(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None:  # pragma: no cover
+    def on_test_epoch_start(
+        self,
+        trainer: lightning.Trainer,
+        pl_module: lightning.LightningModule,  # noqa: ARG002
+    ) -> None:  # pragma: no cover
         self._dataloaders_size = self._get_dataloaders_size(trainer.test_dataloaders)
         self._metrics_builders = [
             TorchMetricsBuilder(self._metrics, self._ks, self._item_count) for _ in self._dataloaders_size
@@ -88,11 +90,10 @@ class ValidationMetricsCallback(L.Callback):
             query_ids, scores, ground_truth = postprocessor.on_validation(query_ids, scores, ground_truth)
         return query_ids, scores, ground_truth
 
-    # pylint: disable=too-many-arguments
     def on_validation_batch_end(
         self,
-        trainer: L.Trainer,
-        pl_module: L.LightningModule,
+        trainer: lightning.Trainer,
+        pl_module: lightning.LightningModule,
         outputs: torch.Tensor,
         batch: ValidationBatch,
         batch_idx: int,
@@ -100,11 +101,10 @@ class ValidationMetricsCallback(L.Callback):
     ) -> None:
         self._batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx)
 
-    # pylint: disable=unused-argument, too-many-arguments
     def on_test_batch_end(
         self,
-        trainer: L.Trainer,
-        pl_module: L.LightningModule,
+        trainer: lightning.Trainer,
+        pl_module: lightning.LightningModule,
         outputs: torch.Tensor,
         batch: ValidationBatch,
         batch_idx: int,
@@ -112,11 +112,10 @@ class ValidationMetricsCallback(L.Callback):
     ) -> None:  # pragma: no cover
         self._batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx)
 
-    # pylint: disable=too-many-arguments
     def _batch_end(
         self,
-        trainer: L.Trainer,
-        pl_module: L.LightningModule,
+        trainer: lightning.Trainer,  # noqa: ARG002
+        pl_module: lightning.LightningModule,
         outputs: torch.Tensor,
         batch: ValidationBatch,
         batch_idx: int,
@@ -131,31 +130,29 @@ class ValidationMetricsCallback(L.Callback):
                 self._metrics_builders[dataloader_idx].get_metrics(),
                 on_epoch=True,
                 sync_dist=True,
-                add_dataloader_idx=True
+                add_dataloader_idx=True,
             )
 
-    # pylint: disable=unused-argument
-    def on_validation_epoch_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None:
+    def on_validation_epoch_end(self, trainer: lightning.Trainer, pl_module: lightning.LightningModule) -> None:
         self._epoch_end(trainer, pl_module)
 
-    # pylint: disable=unused-argument
-    def on_test_epoch_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None:  # pragma: no cover
+    def on_test_epoch_end(
+        self, trainer: lightning.Trainer, pl_module: lightning.LightningModule
+    ) -> None:  # pragma: no cover
         self._epoch_end(trainer, pl_module)
 
-    # pylint: disable=unused-argument
-    def _epoch_end(self, trainer: L.Trainer, pl_module: L.LightningModule) -> None:
-        # pylint: disable=W0212
+    def _epoch_end(self, trainer: lightning.Trainer, pl_module: lightning.LightningModule) -> None:  # noqa: ARG002
         @rank_zero_only
         def print_metrics() -> None:
             metrics = {}
             for name, value in trainer.logged_metrics.items():
-                if '@' in name:
+                if "@" in name:
                     metrics[name] = value.item()
 
             if metrics:
                 metrics_df = metrics_to_df(metrics)
 
-                print(metrics_df)
-                print()
+                print(metrics_df)  # noqa: T201
+                print()  # noqa: T201
 
         print_metrics()

replay/models/nn/sequential/postprocessors/postprocessors.py

@@ -5,6 +5,7 @@ import pandas as pd
 import torch
 
 from replay.data.nn import SequentialDataset
+
 from ._base import BasePostProcessor
 
 
@@ -85,7 +86,6 @@ class SampleItems(BasePostProcessor):
     Generates negative samples to compute sampled metrics
     """
 
-    # pylint: disable=too-many-arguments
     def __init__(
         self,
         grouped_validation_items: pd.DataFrame,

replay/models/nn/sequential/sasrec/dataset.py

@@ -30,7 +30,6 @@ class SasRecTrainingDataset(TorchDataset):
     Dataset that generates samples to train SasRec-like model
     """
 
-    # pylint: disable=too-many-arguments
     def __init__(
         self,
         sequential: SequentialDataset,
@@ -56,13 +55,16 @@ class SasRecTrainingDataset(TorchDataset):
         super().__init__()
         if label_feature_name:
             if label_feature_name not in sequential.schema:
-                raise ValueError("Label feature name not found in provided schema")
+                msg = "Label feature name not found in provided schema"
+                raise ValueError(msg)
 
             if not sequential.schema[label_feature_name].is_cat:
-                raise ValueError("Label feature must be categorical")
+                msg = "Label feature must be categorical"
+                raise ValueError(msg)
 
             if not sequential.schema[label_feature_name].is_seq:
-                raise ValueError("Label feature must be sequential")
+                msg = "Label feature must be sequential"
+                raise ValueError(msg)
 
         self._sequence_shift = sequence_shift
         self._max_sequence_length = max_sequence_length + sequence_shift
@@ -83,8 +85,8 @@ class SasRecTrainingDataset(TorchDataset):
         query_id, padding_mask, features = self._inner[index]
 
         assert self._label_feature_name
-        labels = features[self._label_feature_name][self._sequence_shift :]  # noqa: E203
-        labels_padding_mask = padding_mask[self._sequence_shift :]  # noqa: E203
+        labels = features[self._label_feature_name][self._sequence_shift :]
+        labels_padding_mask = padding_mask[self._sequence_shift :]
 
         output_features: MutableTensorMap = {}
         for feature_name in self._schema:
@@ -165,7 +167,6 @@ class SasRecValidationDataset(TorchDataset):
     Dataset that generates samples to infer and validate SasRec-like model
     """
 
-    # pylint: disable=too-many-arguments
     def __init__(
         self,
         sequential: SequentialDataset,

replay/models/nn/sequential/sasrec/lightning.py

@@ -1,17 +1,17 @@
 import math
-from typing import Any, Optional, Tuple, Union, cast, Dict
+from typing import Any, Dict, Optional, Tuple, Union, cast
 
-import lightning as L
+import lightning
 import torch
 
 from replay.data.nn import TensorMap, TensorSchema
 from replay.models.nn.optimizer_utils import FatOptimizerFactory, LRSchedulerFactory, OptimizerFactory
+
 from .dataset import SasRecPredictionBatch, SasRecTrainingBatch, SasRecValidationBatch
 from .model import SasRecModel
 
 
-# pylint: disable=too-many-instance-attributes
-class SasRec(L.LightningModule):
+class SasRec(lightning.LightningModule):
     """
     SASRec Lightning module.
 
@@ -19,7 +19,6 @@ class SasRec(L.LightningModule):
     for object of SasRec instance.
     """
 
-    # pylint: disable=too-many-arguments, too-many-locals
     def __init__(
         self,
         tensor_schema: TensorSchema,
@@ -94,7 +93,6 @@ class SasRec(L.LightningModule):
         assert item_count
         self._vocab_size = item_count
 
-    # pylint: disable=unused-argument, arguments-differ
     def training_step(self, batch: SasRecTrainingBatch, batch_idx: int) -> torch.Tensor:
         """
         :param batch (SasRecTrainingBatch): Batch of training data.
@@ -108,7 +106,6 @@ class SasRec(L.LightningModule):
         self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, sync_dist=True)
         return loss
 
-    # pylint: disable=arguments-differ
     def forward(self, feature_tensors: TensorMap, padding_mask: torch.BoolTensor) -> torch.Tensor:  # pragma: no cover
         """
         :param feature_tensors: Batch of features.
@@ -118,8 +115,9 @@ class SasRec(L.LightningModule):
         """
         return self._model_predict(feature_tensors, padding_mask)
 
-    # pylint: disable=unused-argument
-    def predict_step(self, batch: SasRecPredictionBatch, batch_idx: int, dataloader_idx: int = 0) -> torch.Tensor:
+    def predict_step(
+        self, batch: SasRecPredictionBatch, batch_idx: int, dataloader_idx: int = 0  # noqa: ARG002
+    ) -> torch.Tensor:
         """
         :param batch: Batch of prediction data.
         :param batch_idx: Batch index.
@@ -130,8 +128,9 @@ class SasRec(L.LightningModule):
         batch = self._prepare_prediction_batch(batch)
         return self._model_predict(batch.features, batch.padding_mask)
 
-    # pylint: disable=unused-argument, arguments-differ
-    def validation_step(self, batch: SasRecValidationBatch, batch_idx: int, dataloader_idx: int = 0) -> torch.Tensor:
+    def validation_step(
+        self, batch: SasRecValidationBatch, batch_idx: int, dataloader_idx: int = 0  # noqa: ARG002
+    ) -> torch.Tensor:
         """
         :param batch (SasRecValidationBatch): Batch of prediction data.
         :param batch_idx (int): Batch index.
@@ -155,57 +154,46 @@ class SasRec(L.LightningModule):
 
     def _prepare_prediction_batch(self, batch: SasRecPredictionBatch) -> SasRecPredictionBatch:
         if batch.padding_mask.shape[1] > self._model.max_len:
-            raise ValueError(
-                f"The length of the submitted sequence \
+            msg = f"The length of the submitted sequence \
                 must not exceed the maximum length of the sequence. \
                 The length of the sequence is given {batch.padding_mask.shape[1]}, \
-                while the maximum length is {self._model.max_len}")
+                while the maximum length is {self._model.max_len}"
+            raise ValueError(msg)
+
         if batch.padding_mask.shape[1] < self._model.max_len:
             query_id, padding_mask, features = batch
             sequence_item_count = padding_mask.shape[1]
             for feature_name, feature_tensor in features.items():
                 if self._schema[feature_name].is_cat:
                     features[feature_name] = torch.nn.functional.pad(
-                        feature_tensor,
-                        (self._model.max_len - sequence_item_count, 0),
-                        value=0
+                        feature_tensor, (self._model.max_len - sequence_item_count, 0), value=0
                     )
                 else:
                     features[feature_name] = torch.nn.functional.pad(
                         feature_tensor.view(feature_tensor.size(0), feature_tensor.size(1)),
                         (self._model.max_len - sequence_item_count, 0),
-                        value=0
+                        value=0,
                     ).unsqueeze(-1)
             padding_mask = torch.nn.functional.pad(
-                padding_mask,
-                (self._model.max_len - sequence_item_count, 0),
-                value=0
+                padding_mask, (self._model.max_len - sequence_item_count, 0), value=0
             )
             batch = SasRecPredictionBatch(query_id, padding_mask, features)
         return batch
 
     def _model_predict(self, feature_tensors: TensorMap, padding_mask: torch.BoolTensor) -> torch.Tensor:
         model: SasRecModel
-        if isinstance(self._model, torch.nn.DataParallel):
-            model = cast(SasRecModel, self._model.module)  # multigpu
-        else:
-            model = self._model
+        model = cast(SasRecModel, self._model.module) if isinstance(self._model, torch.nn.DataParallel) else self._model
         scores = model.predict(feature_tensors, padding_mask)
         return scores
 
     def _compute_loss(self, batch: SasRecTrainingBatch) -> torch.Tensor:
         if self._loss_type == "BCE":
-            if self._loss_sample_count is None:
-                loss_func = self._compute_loss_bce
-            else:
-                loss_func = self._compute_loss_bce_sampled
+            loss_func = self._compute_loss_bce if self._loss_sample_count is None else self._compute_loss_bce_sampled
         elif self._loss_type == "CE":
-            if self._loss_sample_count is None:
-                loss_func = self._compute_loss_ce
-            else:
-                loss_func = self._compute_loss_ce_sampled
+            loss_func = self._compute_loss_ce if self._loss_sample_count is None else self._compute_loss_ce_sampled
         else:
-            raise ValueError(f"Not supported loss type: {self._loss_type}")
+            msg = f"Not supported loss type: {self._loss_type}"
+            raise ValueError(msg)
 
         loss = loss_func(
             batch.features,
@@ -225,8 +213,10 @@ class SasRec(L.LightningModule):
         # [B x L x V]
         logits = self._model.forward(feature_tensors, padding_mask)
 
-        # Take only logits which correspond to non-padded tokens
-        # M = non_zero_count(target_padding_mask)
+        """
+        Take only logits which correspond to non-padded tokens
+        M = non_zero_count(target_padding_mask)
+        """
         logits = logits[target_padding_mask]  # [M x V]
         labels = positive_labels[target_padding_mask]  # [M]
 
@@ -318,7 +308,6 @@ class SasRec(L.LightningModule):
         loss = self._loss(logits, labels_flat)
         return loss
 
-    # pylint: disable=too-many-locals
     def _get_sampled_logits(
         self,
         feature_tensors: TensorMap,
@@ -354,7 +343,8 @@ class SasRec(L.LightningModule):
             else:
                 multinomial_sample_distribution = torch.softmax(positive_logits, dim=-1)
         else:
-            raise NotImplementedError(f"Unknown negative sampling strategy: {self._negative_sampling_strategy}")
+            msg = f"Unknown negative sampling strategy: {self._negative_sampling_strategy}"
+            raise NotImplementedError(msg)
         n_negative_samples = min(n_negative_samples, vocab_size)
 
         if self._negatives_sharing:
@@ -405,7 +395,8 @@ class SasRec(L.LightningModule):
         if self._loss_type == "CE":
             return torch.nn.CrossEntropyLoss()
 
-        raise NotImplementedError("Not supported loss_type")
+        msg = "Not supported loss_type"
+        raise NotImplementedError(msg)
 
     def get_all_embeddings(self) -> Dict[str, torch.nn.Embedding]:
         """
@@ -415,17 +406,18 @@ class SasRec(L.LightningModule):
 
     def set_item_embeddings_by_size(self, new_vocab_size: int):
         """
-        Set item embeddings initialized with xavier_normal_ by new size of vocabulary
-        to item embedder.
+        Keep the current item embeddings and expand vocabulary with new embeddings
+        initialized with xavier_normal_ for new items.
 
-        :param new_vocab_size: Size of vocabulary with new items.
+        :param new_vocab_size: Size of vocabulary with new items included.
             Must be greater then already fitted.
         """
         old_vocab_size = self._model.item_embedder.item_emb.weight.data.shape[0] - 1
         hidden_size = self._model.hidden_size
 
         if new_vocab_size <= old_vocab_size:
-            raise ValueError("New vocabulary size must be greater then already fitted")
+            msg = "New vocabulary size must be greater then already fitted"
+            raise ValueError(msg)
 
         new_embedding = torch.nn.Embedding(new_vocab_size + 1, hidden_size, padding_idx=new_vocab_size)
         torch.nn.init.xavier_normal_(new_embedding.weight)
@@ -443,16 +435,19 @@ class SasRec(L.LightningModule):
             shape (n, h), where n - number of all items, h - model hidden size.
         """
         if all_item_embeddings.dim() != 2:
-            raise ValueError("Input tensor must have (number of all items, model hidden size) shape")
+            msg = "Input tensor must have (number of all items, model hidden size) shape"
+            raise ValueError(msg)
 
         old_vocab_size = self._model.item_embedder.item_emb.weight.data.shape[0] - 1
         new_vocab_size = all_item_embeddings.shape[0]
         hidden_size = self._model.hidden_size
 
         if new_vocab_size < old_vocab_size:
-            raise ValueError("New vocabulary size can't be less then already fitted")
+            msg = "New vocabulary size can't be less then already fitted"
+            raise ValueError(msg)
         if all_item_embeddings.shape[1] != hidden_size:
-            raise ValueError("Input tensor second dimension doesn't match model hidden size")
+            msg = "Input tensor second dimension doesn't match model hidden size"
+            raise ValueError(msg)
 
         new_embedding = torch.nn.Embedding(new_vocab_size + 1, hidden_size, padding_idx=new_vocab_size)
         new_embedding.weight.data[:-1, :] = all_item_embeddings
@@ -467,14 +462,16 @@ class SasRec(L.LightningModule):
             n - number of only new items, h - model hidden size.
         """
         if item_embeddings.dim() != 2:
-            raise ValueError("Input tensor must have (number of new items, model hidden size) shape")
+            msg = "Input tensor must have (number of new items, model hidden size) shape"
+            raise ValueError(msg)
 
         old_vocab_size = self._model.item_embedder.item_emb.weight.data.shape[0] - 1
         new_vocab_size = item_embeddings.shape[0] + old_vocab_size
         hidden_size = self._model.hidden_size
 
         if item_embeddings.shape[1] != hidden_size:
-            raise ValueError("Input tensor second dimension doesn't match model hidden size")
+            msg = "Input tensor second dimension doesn't match model hidden size"
+            raise ValueError(msg)
 
         new_embedding = torch.nn.Embedding(new_vocab_size + 1, hidden_size, padding_idx=new_vocab_size)
         new_embedding.weight.data[:old_vocab_size, :] = self._model.item_embedder.item_emb.weight.data[:-1, :]
@@ -489,3 +486,11 @@ class SasRec(L.LightningModule):
         self._model.item_count = new_vocab_size
         self._model.padding_idx = new_vocab_size
         self._model.masking.padding_idx = new_vocab_size
+        self._model.candidates_to_score = torch.tensor(
+            list(range(new_embedding.weight.data.shape[0] - 1)),
+            device=self._model.candidates_to_score.device,
+            dtype=torch.long,
+        )
+        self._schema.item_id_features[self._schema.item_id_feature_name]._set_cardinality(
+            new_embedding.weight.data.shape[0] - 1
+        )

replay/models/nn/sequential/sasrec/model.py

@@ -1,18 +1,17 @@
 import abc
-from typing import Any, Optional, Tuple, Union, cast, Dict
+import contextlib
+from typing import Any, Dict, Optional, Tuple, Union, cast
 
 import torch
 
 from replay.data.nn import TensorMap, TensorSchema
 
 
-# pylint: disable=too-many-instance-attributes
 class SasRecModel(torch.nn.Module):
     """
     SasRec model
     """
 
-    # pylint: disable=too-many-arguments
     def __init__(
         self,
         schema: TensorSchema,
@@ -189,13 +188,10 @@ class SasRecModel(torch.nn.Module):
 
     def _init(self) -> None:
         for _, param in self.named_parameters():
-            try:
+            with contextlib.suppress(ValueError):
                 torch.nn.init.xavier_normal_(param.data)
-            except ValueError:
-                pass
 
 
-# pylint: disable=too-few-public-methods
 class SasRecMasks:
     """
     SasRec Masks
@@ -316,7 +312,6 @@ class SasRecEmbeddings(torch.nn.Module, BaseSasRecEmbeddings):
     Link: https://arxiv.org/pdf/1808.09781.pdf
     """
 
-    # pylint: disable=too-many-arguments
     def __init__(
         self,
         schema: TensorSchema,
@@ -406,11 +401,7 @@ class SasRecLayers(torch.nn.Module):
         """
         super().__init__()
         self.attention_layers = self._layers_stacker(
-            num_blocks,
-            torch.nn.MultiheadAttention,
-            hidden_size,
-            num_heads,
-            dropout
+            num_blocks, torch.nn.MultiheadAttention, hidden_size, num_heads, dropout
         )
         self.attention_layernorms = self._layers_stacker(num_blocks, torch.nn.LayerNorm, hidden_size, eps=1e-8)
         self.forward_layers = self._layers_stacker(num_blocks, SasRecPointWiseFeedForward, hidden_size, dropout)
@@ -513,7 +504,6 @@ class SasRecPositionalEmbedding(torch.nn.Module):
     Positional embedding.
     """
 
-    # pylint: disable=invalid-name
     def __init__(self, max_len: int, d_model: int) -> None:
         """
         :param max_len: Max sequence length.
@@ -542,7 +532,6 @@ class TiSasRecEmbeddings(torch.nn.Module, BaseSasRecEmbeddings):
     Link: https://cseweb.ucsd.edu/~jmcauley/pdfs/wsdm20b.pdf
     """
 
-    # pylint: disable=too-many-arguments
     def __init__(
         self,
         schema: TensorSchema,
@@ -678,7 +667,6 @@ class TiSasRecLayers(torch.nn.Module):
         self.attention_layernorms = self._layers_stacker(num_blocks, torch.nn.LayerNorm, hidden_size, eps=1e-8)
         self.forward_layernorms = self._layers_stacker(num_blocks, torch.nn.LayerNorm, hidden_size, eps=1e-8)
 
-    # pylint: disable=too-many-arguments
     def forward(
         self,
         seqs: torch.Tensor,
@@ -738,7 +726,6 @@ class TiSasRecAttention(torch.nn.Module):
         self.head_size = hidden_size // head_num
         self.dropout_rate = dropout_rate
 
-    # pylint: disable=too-many-arguments, invalid-name, too-many-locals
     def forward(
         self,
         queries: torch.LongTensor,

replay/models/pop_rec.py

@@ -1,8 +1,8 @@
-
 from replay.data.dataset import Dataset
-from .base_rec import NonPersonalizedRecommender
 from replay.utils import PYSPARK_AVAILABLE
 
+from .base_rec import NonPersonalizedRecommender
+
 if PYSPARK_AVAILABLE:
     from pyspark.sql import functions as sf
 
@@ -23,7 +23,11 @@ class PopRec(NonPersonalizedRecommender):
     >>> import pandas as pd
     >>> from replay.data.dataset import Dataset, FeatureSchema, FeatureInfo, FeatureHint, FeatureType
     >>> from replay.utils.spark_utils import convert2spark
-    >>> data_frame = pd.DataFrame({"user_id": [1, 1, 2, 2, 3, 4], "item_id": [1, 2, 2, 3, 3, 3], "rating": [0.5, 1, 0.1, 0.8, 0.7, 1]})
+    >>> data_frame = pd.DataFrame(
+    ...    {"user_id": [1, 1, 2, 2, 3, 4],
+    ...     "item_id": [1, 2, 2, 3, 3, 3],
+    ...     "rating": [0.5, 1, 0.1, 0.8, 0.7, 1]}
+    ... )
     >>> data_frame
         user_id   item_id     rating
     0         1         1        0.5
@@ -104,9 +108,7 @@ class PopRec(NonPersonalizedRecommender):
             `Cold_weight` value should be in interval (0, 1].
         """
         self.use_rating = use_rating
-        super().__init__(
-            add_cold_items=add_cold_items, cold_weight=cold_weight
-        )
+        super().__init__(add_cold_items=add_cold_items, cold_weight=cold_weight)
 
     @property
     def _init_args(self):
@@ -120,7 +122,6 @@ class PopRec(NonPersonalizedRecommender):
         self,
         dataset: Dataset,
     ) -> None:
-
         agg_func = sf.countDistinct(self.query_column).alias(self.rating_column)
         if self.use_rating:
             agg_func = sf.sum(self.rating_column).alias(self.rating_column)
@@ -128,9 +129,7 @@ class PopRec(NonPersonalizedRecommender):
         self.item_popularity = (
             dataset.interactions.groupBy(self.item_column)
             .agg(agg_func)
-            .withColumn(
-                self.rating_column, sf.col(self.rating_column) / sf.lit(self.queries_count)
-            )
+            .withColumn(self.rating_column, sf.col(self.rating_column) / sf.lit(self.queries_count))
         )
 
         self.item_popularity.cache().count()

replay/models/query_pop_rec.py

@@ -1,8 +1,8 @@
-
 from replay.data import Dataset
-from .base_rec import Recommender
 from replay.utils import PYSPARK_AVAILABLE, SparkDataFrame
 
+from .base_rec import Recommender
+
 if PYSPARK_AVAILABLE:
     from pyspark.sql import functions as sf
 
@@ -76,7 +76,6 @@ class QueryPopRec(Recommender):
         self,
         dataset: Dataset,
     ) -> None:
-
         query_rating_sum = (
             dataset.interactions.groupBy(self.query_column)
             .agg(sf.sum(self.rating_column).alias("query_rel_sum"))
@@ -94,9 +93,7 @@ class QueryPopRec(Recommender):
             .select(
                 self.query_column,
                 self.item_column,
-                (sf.col("query_item_rel_sum") / sf.col("query_rel_sum")).alias(
-                    self.rating_column
-                ),
+                (sf.col("query_item_rel_sum") / sf.col("query_rel_sum")).alias(self.rating_column),
             )
         )
         self.query_item_popularity.cache().count()
@@ -105,20 +102,15 @@ class QueryPopRec(Recommender):
         if hasattr(self, "query_item_popularity"):
             self.query_item_popularity.unpersist()
 
-    # pylint: disable=too-many-arguments
     def _predict(
         self,
-        dataset: Dataset,
-        k: int,
+        dataset: Dataset,  # noqa: ARG002
+        k: int,  # noqa: ARG002
         queries: SparkDataFrame,
         items: SparkDataFrame,
         filter_seen_items: bool = True,
     ) -> SparkDataFrame:
         if filter_seen_items:
-            self.logger.warning(
-                "QueryPopRec can't predict new items, recommendations will not be filtered"
-            )
+            self.logger.warning("QueryPopRec can't predict new items, recommendations will not be filtered")
 
-        return self.query_item_popularity.join(queries, on=self.query_column).join(
-            items, on=self.item_column
-        )
+        return self.query_item_popularity.join(queries, on=self.query_column).join(items, on=self.item_column)