replay-rec 0.20.0__py3-none-any.whl → 0.20.1__py3-none-any.whl

replay/models/base_rec.py

@@ -13,8 +13,9 @@ Base abstract classes:
 
 import warnings
 from abc import ABC, abstractmethod
+from collections.abc import Iterable
 from os.path import join
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+from typing import Any, Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -55,14 +56,14 @@ class IsSavable(ABC):
 
     @property
     @abstractmethod
-    def _init_args(self) -> Dict:
+    def _init_args(self) -> dict:
         """
         Dictionary of the model attributes passed during model initialization.
         Used for model saving and loading
         """
 
     @property
-    def _dataframes(self) -> Dict:
+    def _dataframes(self) -> dict:
         """
         Dictionary of the model dataframes required for inference.
         Used for model saving and loading
@@ -508,7 +509,7 @@ class BaseRecommender(IsSavable, IsOptimizible, RecommenderCommons, ABC):
             or None if `file_path` is provided
         """
         if dataset is not None:
-            interactions, query_features, item_features, pairs = [
+            interactions, query_features, item_features, pairs = (
                 convert2spark(df)
                 for df in [
                     dataset.interactions,
@@ -516,7 +517,7 @@ class BaseRecommender(IsSavable, IsOptimizible, RecommenderCommons, ABC):
                     dataset.item_features,
                     pairs,
                 ]
-            ]
+            )
             if set(pairs.columns) != {self.item_column, self.query_column}:
                 msg = "pairs must be a dataframe with columns strictly [user_idx, item_idx]"
                 raise ValueError(msg)
@@ -590,7 +591,7 @@ class BaseRecommender(IsSavable, IsOptimizible, RecommenderCommons, ABC):
 
     def _get_features_wrap(
         self, ids: SparkDataFrame, features: Optional[SparkDataFrame]
-    ) -> Optional[Tuple[SparkDataFrame, int]]:
+    ) -> Optional[tuple[SparkDataFrame, int]]:
         if self.query_column not in ids.columns and self.item_column not in ids.columns:
             msg = f"{self.query_column} or {self.item_column} missing"
             raise ValueError(msg)
@@ -599,7 +600,7 @@ class BaseRecommender(IsSavable, IsOptimizible, RecommenderCommons, ABC):
 
     def _get_features(
         self, ids: SparkDataFrame, features: Optional[SparkDataFrame]  # noqa: ARG002
-    ) -> Tuple[Optional[SparkDataFrame], Optional[int]]:
+    ) -> tuple[Optional[SparkDataFrame], Optional[int]]:
         """
         Get embeddings from model
 
@@ -679,7 +680,7 @@ class ItemVectorModel(BaseRecommender):
     """Parent for models generating items' vector representations"""
 
     can_predict_item_to_item: bool = True
-    item_to_item_metrics: List[str] = [
+    item_to_item_metrics: list[str] = [
         "euclidean_distance_sim",
         "cosine_similarity",
         "dot_product",
@@ -899,7 +900,7 @@ class HybridRecommender(BaseRecommender, ABC):
 
     def get_features(
         self, ids: SparkDataFrame, features: Optional[SparkDataFrame]
-    ) -> Optional[Tuple[SparkDataFrame, int]]:
+    ) -> Optional[tuple[SparkDataFrame, int]]:
         """
         Returns query or item feature vectors as a Column with type ArrayType
         If a model does not have a vector for some ids they are not present in the final result.
@@ -1026,7 +1027,7 @@ class Recommender(BaseRecommender, ABC):
             recs_file_path=recs_file_path,
         )
 
-    def get_features(self, ids: SparkDataFrame) -> Optional[Tuple[SparkDataFrame, int]]:
+    def get_features(self, ids: SparkDataFrame) -> Optional[tuple[SparkDataFrame, int]]:
         """
         Returns query or item feature vectors as a Column with type ArrayType
 

replay/models/cat_pop_rec.py

@@ -1,5 +1,6 @@
+from collections.abc import Iterable
 from os.path import join
-from typing import Iterable, Optional, Union
+from typing import Optional, Union
 
 from replay.data import Dataset
 from replay.utils import PYSPARK_AVAILABLE, SparkDataFrame

replay/models/extensions/ann/ann_mixin.py

@@ -2,7 +2,8 @@ import importlib
 import logging
 import sys
 from abc import abstractmethod
-from typing import Any, Iterable, Optional, Union
+from collections.abc import Iterable
+from typing import Any, Optional, Union
 
 from replay.data import Dataset
 from replay.models.common import RecommenderCommons

replay/models/extensions/ann/index_builders/executor_hnswlib_index_builder.py

@@ -1,5 +1,6 @@
 import logging
-from typing import Iterator, Optional
+from collections.abc import Iterator
+from typing import Optional
 
 import numpy as np
 

replay/models/extensions/ann/index_builders/executor_nmslib_index_builder.py

@@ -1,5 +1,6 @@
 import logging
-from typing import Iterator, Optional
+from collections.abc import Iterator
+from typing import Optional
 
 import pandas as pd
 

replay/models/lin_ucb.py

@@ -1,5 +1,6 @@
 import warnings
-from typing import List, Tuple, Union
+from os.path import join
+from typing import Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -8,7 +9,11 @@ from tqdm import tqdm
 
 from replay.data.dataset import Dataset
 from replay.utils import SparkDataFrame
-from replay.utils.spark_utils import convert2spark
+from replay.utils.spark_utils import (
+    convert2spark,
+    load_pickled_from_parquet,
+    save_picklable_to_parquet,
+)
 
 from .base_rec import HybridRecommender
 
@@ -70,7 +75,7 @@ class HybridArm:
         # right-hand side of the regression
         self.b = np.zeros(d, dtype=float)
 
-    def feature_update(self, usr_features, usr_itm_features, relevances) -> Tuple[np.ndarray, np.ndarray]:
+    def feature_update(self, usr_features, usr_itm_features, relevances) -> tuple[np.ndarray, np.ndarray]:
         """
         Function to update featurs or each Lin-UCB hand in the current model.
 
@@ -175,8 +180,9 @@ class LinUCB(HybridRecommender):
         "alpha": {"type": "uniform", "args": [0.001, 10.0]},
     }
     _study = None  # field required for proper optuna's optimization
-    linucb_arms: List[Union[DisjointArm, HybridArm]]  # initialize only when working within fit method
+    linucb_arms: list[Union[DisjointArm, HybridArm]]  # initialize only when working within fit method
     rel_matrix: np.array  # matrix with relevance scores from predict method
+    _num_items: int  # number of items/arms
 
     def __init__(
         self,
@@ -195,7 +201,7 @@ class LinUCB(HybridRecommender):
 
     @property
     def _init_args(self):
-        return {"is_hybrid": self.is_hybrid}
+        return {"is_hybrid": self.is_hybrid, "eps": self.eps, "alpha": self.alpha}
 
     def _verify_features(self, dataset: Dataset):
         if dataset.query_features is None:
@@ -230,6 +236,7 @@ class LinUCB(HybridRecommender):
         self._num_items = item_features.shape[0]
         self._user_dim_size = user_features.shape[1] - 1
         self._item_dim_size = item_features.shape[1] - 1
+        self._user_idxs_list = set(user_features[feature_schema.query_id_column].values)
 
         # now initialize an arm object for each potential arm instance
         if self.is_hybrid:
@@ -248,11 +255,14 @@ class LinUCB(HybridRecommender):
             ]
 
             for i in tqdm(range(self._num_items)):
-                B = log.loc[log[feature_schema.item_id_column] == i]  # noqa: N806
-                idxs_list = B[feature_schema.query_id_column].values
-                rel_list = B[feature_schema.interactions_rating_column].values
+                B = log.loc[  # noqa: N806
+                    (log[feature_schema.item_id_column] == i)
+                    & (log[feature_schema.query_id_column].isin(self._user_idxs_list))
+                ]
                 if not B.empty:
                     # if we have at least one user interacting with the hand i
+                    idxs_list = B[feature_schema.query_id_column].values
+                    rel_list = B[feature_schema.interactions_rating_column].values
                     cur_usrs = scs.csr_matrix(
                         user_features.query(f"{feature_schema.query_id_column} in @idxs_list")
                         .drop(columns=[feature_schema.query_id_column])
@@ -284,11 +294,14 @@ class LinUCB(HybridRecommender):
             ]
 
             for i in range(self._num_items):
-                B = log.loc[log[feature_schema.item_id_column] == i]  # noqa: N806
-                idxs_list = B[feature_schema.query_id_column].values  # noqa: F841
-                rel_list = B[feature_schema.interactions_rating_column].values
+                B = log.loc[  # noqa: N806
+                    (log[feature_schema.item_id_column] == i)
+                    & (log[feature_schema.query_id_column].isin(self._user_idxs_list))
+                ]
                 if not B.empty:
                     # if we have at least one user interacting with the hand i
+                    idxs_list = B[feature_schema.query_id_column].values  # noqa: F841
+                    rel_list = B[feature_schema.interactions_rating_column].values
                     cur_usrs = user_features.query(f"{feature_schema.query_id_column} in @idxs_list").drop(
                         columns=[feature_schema.query_id_column]
                     )
@@ -318,8 +331,10 @@ class LinUCB(HybridRecommender):
         user_features = dataset.query_features
         item_features = dataset.item_features
         big_k = min(oversample * k, item_features.shape[0])
+        self._user_idxs_list = set(user_features[feature_schema.query_id_column].values)
 
         users = users.toPandas()
+        users = users[users[feature_schema.query_id_column].isin(self._user_idxs_list)]
         num_user_pred = users.shape[0]
         rel_matrix = np.zeros((num_user_pred, self._num_items), dtype=float)
 
@@ -404,3 +419,34 @@ class LinUCB(HybridRecommender):
         warnings.warn(warn_msg)
         dataset.to_spark()
         return convert2spark(res_df)
+
+    def _save_model(self, path: str, additional_params: Optional[dict] = None):
+        super()._save_model(path, additional_params)
+
+        save_picklable_to_parquet(self.linucb_arms, join(path, "linucb_arms.dump"))
+
+        if self.is_hybrid:
+            linucb_hybrid_shared_params = {
+                "A_0": self.A_0,
+                "A_0_inv": self.A_0_inv,
+                "b_0": self.b_0,
+                "beta": self.beta,
+            }
+            save_picklable_to_parquet(
+                linucb_hybrid_shared_params,
+                join(path, "linucb_hybrid_shared_params.dump"),
+            )
+
+    def _load_model(self, path: str):
+        super()._load_model(path)
+
+        loaded_linucb_arms = load_pickled_from_parquet(join(path, "linucb_arms.dump"))
+        self.linucb_arms = loaded_linucb_arms
+        self._num_items = len(loaded_linucb_arms)
+
+        if self.is_hybrid:
+            loaded_linucb_hybrid_shared_params = load_pickled_from_parquet(
+                join(path, "linucb_hybrid_shared_params.dump")
+            )
+            for param, value in loaded_linucb_hybrid_shared_params.items():
+                setattr(self, param, value)

replay/models/nn/optimizer_utils/optimizer_factory.py

@@ -1,5 +1,5 @@
 import abc
-from typing import Iterator, Tuple
+from collections.abc import Iterator
 
 import torch
 
@@ -47,7 +47,7 @@ class FatOptimizerFactory(OptimizerFactory):
         learning_rate: float = 0.001,
         weight_decay: float = 0.0,
         sgd_momentum: float = 0.0,
-        betas: Tuple[float, float] = (0.9, 0.98),
+        betas: tuple[float, float] = (0.9, 0.98),
     ) -> None:
         super().__init__()
         self.optimizer = optimizer

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

@@ -1,5 +1,5 @@
 import abc
-from typing import NamedTuple, Optional, Tuple, cast
+from typing import NamedTuple, Optional, cast
 
 import torch
 from torch.utils.data import Dataset as TorchDataset
@@ -12,7 +12,6 @@ from replay.data.nn import (
     TorchSequentialDataset,
     TorchSequentialValidationDataset,
 )
-from replay.utils import deprecation_warning
 
 
 class Bert4RecTrainingBatch(NamedTuple):
@@ -89,10 +88,6 @@ class Bert4RecTrainingDataset(TorchDataset):
     Dataset that generates samples to train BERT-like model
     """
 
-    @deprecation_warning(
-        "`padding_value` parameter will be removed in future versions. "
-        "Instead, you should specify `padding_value` for each column in TensorSchema"
-    )
     def __init__(
         self,
         sequential: SequentialDataset,
@@ -101,7 +96,7 @@ class Bert4RecTrainingDataset(TorchDataset):
         sliding_window_step: Optional[int] = None,
         label_feature_name: Optional[str] = None,
         custom_masker: Optional[Bert4RecMasker] = None,
-        padding_value: int = 0,
+        padding_value: Optional[int] = None,
     ) -> None:
         """
         :param sequential: Sequential dataset with training data.
@@ -181,15 +176,11 @@ class Bert4RecPredictionDataset(TorchDataset):
     Dataset that generates samples to infer BERT-like model
     """
 
-    @deprecation_warning(
-        "`padding_value` parameter will be removed in future versions. "
-        "Instead, you should specify `padding_value` for each column in TensorSchema"
-    )
     def __init__(
         self,
         sequential: SequentialDataset,
         max_sequence_length: int,
-        padding_value: int = 0,
+        padding_value: Optional[int] = None,
     ) -> None:
         """
         :param sequential: Sequential dataset with data to make predictions at.
@@ -239,17 +230,13 @@ class Bert4RecValidationDataset(TorchDataset):
     Dataset that generates samples to infer and validate BERT-like model
     """
 
-    @deprecation_warning(
-        "`padding_value` parameter will be removed in future versions. "
-        "Instead, you should specify `padding_value` for each column in TensorSchema"
-    )
     def __init__(
         self,
         sequential: SequentialDataset,
         ground_truth: SequentialDataset,
         train: SequentialDataset,
         max_sequence_length: int,
-        padding_value: int = 0,
+        padding_value: Optional[int] = None,
         label_feature_name: Optional[str] = None,
     ):
         """
@@ -295,7 +282,7 @@ def _shift_features(
     schema: TensorSchema,
     features: TensorMap,
     padding_mask: torch.BoolTensor,
-) -> Tuple[TensorMap, torch.BoolTensor, torch.BoolTensor]:
+) -> tuple[TensorMap, torch.BoolTensor, torch.BoolTensor]:
     shifted_features: MutableTensorMap = {}
     for feature_name, feature in schema.items():
         if feature.is_seq:

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

@@ -1,5 +1,5 @@
 import math
-from typing import Any, Dict, Literal, Optional, Tuple, Union, cast
+from typing import Any, Literal, Optional, Union, cast
 
 import lightning
 import torch
@@ -338,7 +338,7 @@ class Bert4Rec(lightning.LightningModule):
         positive_labels: torch.LongTensor,
         padding_mask: torch.BoolTensor,
         tokens_mask: torch.BoolTensor,
-    ) -> Tuple[torch.Tensor, torch.Tensor, torch.LongTensor, torch.LongTensor, int]:
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.LongTensor, torch.LongTensor, int]:
         assert self._loss_sample_count is not None
         n_negative_samples = self._loss_sample_count
 
@@ -440,7 +440,7 @@ class Bert4Rec(lightning.LightningModule):
         msg = "Not supported loss_type"
         raise NotImplementedError(msg)
 
-    def get_all_embeddings(self) -> Dict[str, torch.nn.Embedding]:
+    def get_all_embeddings(self) -> dict[str, torch.nn.Embedding]:
         """
         :returns: copy of all embeddings as a dictionary.
         """

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

@@ -1,6 +1,6 @@
 import contextlib
 from abc import ABC, abstractmethod
-from typing import Dict, Optional, Union
+from typing import Optional, Union
 
 import torch
 import torch.nn as nn
@@ -303,7 +303,7 @@ class BertEmbedding(torch.nn.Module):
         """
         return self.cat_embeddings[self.schema.item_id_feature_name].weight
 
-    def get_all_embeddings(self) -> Dict[str, torch.Tensor]:
+    def get_all_embeddings(self) -> dict[str, torch.Tensor]:
         """
         :returns: copy of all embeddings presented in this layer as a dict.
         """

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

@@ -1,5 +1,5 @@
 import abc
-from typing import Generic, List, Optional, Protocol, Tuple, TypeVar, cast
+from typing import Generic, Optional, Protocol, TypeVar, cast
 
 import lightning
 import torch
@@ -38,7 +38,7 @@ class BasePredictionCallback(lightning.Callback, Generic[_T]):
         query_column: str,
         item_column: str,
         rating_column: str = "rating",
-        postprocessors: Optional[List[BasePostProcessor]] = None,
+        postprocessors: Optional[list[BasePostProcessor]] = None,
     ) -> None:
         """
         :param top_k: Takes the highest k scores in the ranking.
@@ -52,10 +52,10 @@ class BasePredictionCallback(lightning.Callback, Generic[_T]):
         self.item_column = item_column
         self.rating_column = rating_column
         self._top_k = top_k
-        self._postprocessors: List[BasePostProcessor] = postprocessors or []
-        self._query_batches: List[torch.Tensor] = []
-        self._item_batches: List[torch.Tensor] = []
-        self._item_scores: List[torch.Tensor] = []
+        self._postprocessors: list[BasePostProcessor] = postprocessors or []
+        self._query_batches: list[torch.Tensor] = []
+        self._item_batches: list[torch.Tensor] = []
+        self._item_scores: list[torch.Tensor] = []
 
     def on_predict_epoch_start(
         self, trainer: lightning.Trainer, pl_module: lightning.LightningModule  # noqa: ARG002
@@ -97,7 +97,7 @@ class BasePredictionCallback(lightning.Callback, Generic[_T]):
 
     def _compute_pipeline(
         self, query_ids: torch.LongTensor, scores: torch.Tensor
-    ) -> Tuple[torch.LongTensor, torch.Tensor]:
+    ) -> tuple[torch.LongTensor, torch.Tensor]:
         for postprocessor in self._postprocessors:
             query_ids, scores = postprocessor.on_prediction(query_ids, scores)
         return query_ids, scores
@@ -166,7 +166,7 @@ class SparkPredictionCallback(BasePredictionCallback[SparkDataFrame]):
         item_column: str,
         rating_column: str,
         spark_session: SparkSession,
-        postprocessors: Optional[List[BasePostProcessor]] = None,
+        postprocessors: Optional[list[BasePostProcessor]] = None,
     ) -> None:
         """
         :param top_k: Takes the highest k scores in the ranking.
@@ -213,7 +213,7 @@ class SparkPredictionCallback(BasePredictionCallback[SparkDataFrame]):
         return prediction
 
 
-class TorchPredictionCallback(BasePredictionCallback[Tuple[torch.LongTensor, torch.LongTensor, torch.Tensor]]):
+class TorchPredictionCallback(BasePredictionCallback[tuple[torch.LongTensor, torch.LongTensor, torch.Tensor]]):
     """
     Callback for predition stage with tuple of tensors
     """
@@ -221,7 +221,7 @@ class TorchPredictionCallback(BasePredictionCallback[Tuple[torch.LongTensor, tor
     def __init__(
         self,
         top_k: int,
-        postprocessors: Optional[List[BasePostProcessor]] = None,
+        postprocessors: Optional[list[BasePostProcessor]] = None,
     ) -> None:
         """
         :param top_k: Takes the highest k scores in the ranking.
@@ -240,7 +240,7 @@ class TorchPredictionCallback(BasePredictionCallback[Tuple[torch.LongTensor, tor
         query_ids: torch.Tensor,
         item_ids: torch.Tensor,
         item_scores: torch.Tensor,
-    ) -> Tuple[torch.LongTensor, torch.LongTensor, torch.Tensor]:
+    ) -> tuple[torch.LongTensor, torch.LongTensor, torch.Tensor]:
         return (
             cast(torch.LongTensor, query_ids.flatten().cpu().long()),
             cast(torch.LongTensor, item_ids.cpu().long()),
@@ -254,7 +254,7 @@ class QueryEmbeddingsPredictionCallback(lightning.Callback):
     """
 
     def __init__(self):
-        self._embeddings_per_batch: List[torch.Tensor] = []
+        self._embeddings_per_batch: list[torch.Tensor] = []
 
     def on_predict_epoch_start(
         self, trainer: lightning.Trainer, pl_module: lightning.LightningModule  # noqa: ARG002

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

@@ -1,4 +1,4 @@
-from typing import Any, List, Literal, Optional, Protocol, Tuple
+from typing import Any, Literal, Optional, Protocol
 
 import lightning
 import torch
@@ -38,9 +38,9 @@ class ValidationMetricsCallback(lightning.Callback):
 
     def __init__(
         self,
-        metrics: Optional[List[CallbackMetricName]] = None,
-        ks: Optional[List[int]] = None,
-        postprocessors: Optional[List[BasePostProcessor]] = None,
+        metrics: Optional[list[CallbackMetricName]] = None,
+        ks: Optional[list[int]] = None,
+        postprocessors: Optional[list[BasePostProcessor]] = None,
         item_count: Optional[int] = None,
     ):
         """
@@ -52,11 +52,11 @@ class ValidationMetricsCallback(lightning.Callback):
         self._metrics = metrics
         self._ks = ks
         self._item_count = item_count
-        self._metrics_builders: List[TorchMetricsBuilder] = []
-        self._dataloaders_size: List[int] = []
-        self._postprocessors: List[BasePostProcessor] = postprocessors or []
+        self._metrics_builders: list[TorchMetricsBuilder] = []
+        self._dataloaders_size: list[int] = []
+        self._postprocessors: list[BasePostProcessor] = postprocessors or []
 
-    def _get_dataloaders_size(self, dataloaders: Optional[Any]) -> List[int]:
+    def _get_dataloaders_size(self, dataloaders: Optional[Any]) -> list[int]:
         if isinstance(dataloaders, torch.utils.data.DataLoader):
             return [len(dataloaders)]
         return [len(dataloader) for dataloader in dataloaders]
@@ -85,7 +85,7 @@ class ValidationMetricsCallback(lightning.Callback):
 
     def _compute_pipeline(
         self, query_ids: torch.LongTensor, scores: torch.Tensor, ground_truth: torch.LongTensor
-    ) -> Tuple[torch.LongTensor, torch.Tensor, torch.LongTensor]:
+    ) -> tuple[torch.LongTensor, torch.Tensor, torch.LongTensor]:
         for postprocessor in self._postprocessors:
             query_ids, scores, ground_truth = postprocessor.on_validation(query_ids, scores, ground_truth)
         return query_ids, scores, ground_truth

replay/models/nn/sequential/compiled/base_compiled_model.py

@@ -1,7 +1,7 @@
 import pathlib
 import tempfile
 from abc import abstractmethod
-from typing import Any, Dict, List, Literal, Optional, Tuple, Union
+from typing import Any, Literal, Optional, Union
 
 import lightning
 import openvino as ov
@@ -68,7 +68,7 @@ class BaseCompiledModel:
         """
         self._batch_size: int
         self._max_seq_len: int
-        self._inputs_names: List[str]
+        self._inputs_names: list[str]
         self._output_name: str
 
         self._set_inner_params_from_openvino_model(compiled_model)
@@ -171,9 +171,9 @@ class BaseCompiledModel:
     @staticmethod
     def _run_model_compilation(
         lightning_model: lightning.LightningModule,
-        model_input_sample: Tuple[Union[torch.Tensor, Dict[str, torch.Tensor]]],
-        model_input_names: List[str],
-        model_dynamic_axes_in_input: Dict[str, Dict],
+        model_input_sample: tuple[Union[torch.Tensor, dict[str, torch.Tensor]]],
+        model_input_names: list[str],
+        model_dynamic_axes_in_input: dict[str, dict],
         batch_size: int,
         num_candidates_to_score: Union[int, None],
         num_threads: Optional[int] = None,

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

@@ -1,5 +1,4 @@
 import abc
-from typing import Tuple
 
 import torch
 
@@ -10,7 +9,7 @@ class BasePostProcessor(abc.ABC):  # pragma: no cover
     """
 
     @abc.abstractmethod
-    def on_prediction(self, query_ids: torch.LongTensor, scores: torch.Tensor) -> Tuple[torch.LongTensor, torch.Tensor]:
+    def on_prediction(self, query_ids: torch.LongTensor, scores: torch.Tensor) -> tuple[torch.LongTensor, torch.Tensor]:
         """
         Prediction step.
 
@@ -24,7 +23,7 @@ class BasePostProcessor(abc.ABC):  # pragma: no cover
     @abc.abstractmethod
     def on_validation(
         self, query_ids: torch.LongTensor, scores: torch.Tensor, ground_truth: torch.LongTensor
-    ) -> Tuple[torch.LongTensor, torch.Tensor, torch.LongTensor]:
+    ) -> tuple[torch.LongTensor, torch.Tensor, torch.LongTensor]:
         """
         Validation step.
 

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

@@ -1,4 +1,4 @@
-from typing import List, Optional, Set, Tuple, Union, cast
+from typing import Optional, Union, cast
 
 import numpy as np
 import pandas as pd
@@ -22,7 +22,7 @@ class RemoveSeenItems(BasePostProcessor):
 
     def on_validation(
         self, query_ids: torch.LongTensor, scores: torch.Tensor, ground_truth: torch.LongTensor
-    ) -> Tuple[torch.LongTensor, torch.Tensor, torch.LongTensor]:
+    ) -> tuple[torch.LongTensor, torch.Tensor, torch.LongTensor]:
         """
         Validation step.
 
@@ -36,7 +36,7 @@ class RemoveSeenItems(BasePostProcessor):
         modified_scores = self._compute_scores(query_ids, scores)
         return query_ids, modified_scores, ground_truth
 
-    def on_prediction(self, query_ids: torch.LongTensor, scores: torch.Tensor) -> Tuple[torch.LongTensor, torch.Tensor]:
+    def on_prediction(self, query_ids: torch.LongTensor, scores: torch.Tensor) -> tuple[torch.LongTensor, torch.Tensor]:
         """
         Prediction step.
 
@@ -51,7 +51,7 @@ class RemoveSeenItems(BasePostProcessor):
 
     def _compute_scores(self, query_ids: torch.LongTensor, scores: torch.Tensor) -> torch.Tensor:
         flat_seen_item_ids = self._get_flat_seen_item_ids(query_ids)
-        return self._fill_item_ids(scores, flat_seen_item_ids, -np.inf)
+        return self._fill_item_ids(scores.clone(), flat_seen_item_ids, -np.inf)
 
     def _fill_item_ids(
         self,
@@ -124,13 +124,13 @@ class SampleItems(BasePostProcessor):
         self.sample_count = sample_count
         users = grouped_validation_items[user_col].to_numpy()
         items = grouped_validation_items[item_col].to_numpy()
-        self.items_list: List[Set[int]] = [set() for _ in range(users.shape[0])]
+        self.items_list: list[set[int]] = [set() for _ in range(users.shape[0])]
         for i in range(users.shape[0]):
             self.items_list[users[i]] = set(items[i])
 
     def on_validation(
         self, query_ids: torch.LongTensor, scores: torch.Tensor, ground_truth: torch.LongTensor
-    ) -> Tuple[torch.LongTensor, torch.Tensor, torch.LongTensor]:
+    ) -> tuple[torch.LongTensor, torch.Tensor, torch.LongTensor]:
         """
         Validation step.
 
@@ -143,7 +143,7 @@ class SampleItems(BasePostProcessor):
         modified_score = self._compute_score(query_ids, scores, ground_truth)
         return query_ids, modified_score, ground_truth
 
-    def on_prediction(self, query_ids: torch.LongTensor, scores: torch.Tensor) -> Tuple[torch.LongTensor, torch.Tensor]:
+    def on_prediction(self, query_ids: torch.LongTensor, scores: torch.Tensor) -> tuple[torch.LongTensor, torch.Tensor]:
         """
         Prediction step.
 
@@ -160,8 +160,8 @@ class SampleItems(BasePostProcessor):
     ) -> torch.Tensor:
         batch_size = query_ids.shape[0]
         item_ids = ground_truth.cpu().numpy() if ground_truth is not None else None
-        candidate_ids: List[torch.Tensor] = []
-        candidate_labels: List[torch.Tensor] = []
+        candidate_ids: list[torch.Tensor] = []
+        candidate_labels: list[torch.Tensor] = []
         for user in range(batch_size):
             ground_truth_items = set(item_ids[user]) if ground_truth is not None else set()
             sample, label = self._generate_samples_for_user(ground_truth_items, self.items_list[user])
@@ -183,8 +183,8 @@ class SampleItems(BasePostProcessor):
         return new_scores.reshape_as(scores)
 
     def _generate_samples_for_user(
-        self, ground_truth_items: Set[int], input_items: Set[int]
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        self, ground_truth_items: set[int], input_items: set[int]
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         negative_sample_count = self.sample_count - len(ground_truth_items)
         assert negative_sample_count > 0