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

replay/metrics/recall.py

@@ -1,5 +1,3 @@
-from typing import List
-
 from .base_metric import Metric
 
 
@@ -65,7 +63,7 @@ class Recall(Metric):
     """
 
     @staticmethod
-    def _get_metric_value_by_user(ks: List[int], ground_truth: List, pred: List) -> List[float]:
+    def _get_metric_value_by_user(ks: list[int], ground_truth: list, pred: list) -> list[float]:
         if not ground_truth or not pred:
             return [0.0 for _ in ks]
         set_gt = set(ground_truth)

replay/metrics/rocauc.py

@@ -1,5 +1,3 @@
-from typing import List
-
 from .base_metric import Metric
 
 
@@ -74,7 +72,7 @@ class RocAuc(Metric):
     """
 
     @staticmethod
-    def _get_metric_value_by_user(ks: List[int], ground_truth: List, pred: List) -> List[float]:
+    def _get_metric_value_by_user(ks: list[int], ground_truth: list, pred: list) -> list[float]:
         if not ground_truth or not pred:
             return [0.0 for _ in ks]
         set_gt = set(ground_truth)

replay/metrics/surprisal.py

@@ -1,5 +1,5 @@
 from collections import defaultdict
-from typing import Dict, List, Union
+from typing import Union
 
 import numpy as np
 import polars as pl
@@ -82,7 +82,7 @@ class Surprisal(Metric):
     <BLANKLINE>
     """
 
-    def _get_weights(self, train: Dict) -> Dict:
+    def _get_weights(self, train: dict) -> dict:
         n_users = len(train.keys())
         items_counter = defaultdict(set)
         for user, items in train.items():
@@ -93,7 +93,7 @@ class Surprisal(Metric):
             weights[item] = np.log2(n_users / len(users)) / np.log2(n_users)
         return weights
 
-    def _get_recommendation_weights(self, recommendations: Dict, train: Dict) -> Dict:
+    def _get_recommendation_weights(self, recommendations: dict, train: dict) -> dict:
         weights = self._get_weights(train)
         recs_with_weights = {}
         for user, items in recommendations.items():
@@ -183,7 +183,7 @@ class Surprisal(Metric):
         )
 
     @staticmethod
-    def _get_metric_value_by_user(ks: List[int], pred_item_ids: List, pred_weights: List) -> List[float]:
+    def _get_metric_value_by_user(ks: list[int], pred_item_ids: list, pred_weights: list) -> list[float]:
         if not pred_item_ids:
             return [0.0 for _ in ks]
         res = []

replay/metrics/torch_metrics_builder.py

@@ -1,6 +1,7 @@
 import abc
+from collections.abc import Mapping
 from dataclasses import dataclass
-from typing import Any, Dict, List, Literal, Mapping, Optional, Set
+from typing import Any, Literal, Optional
 
 import numpy as np
 
@@ -19,13 +20,13 @@ MetricName = Literal[
     "coverage",
 ]
 
-DEFAULT_METRICS: List[MetricName] = [
+DEFAULT_METRICS: list[MetricName] = [
     "map",
     "ndcg",
     "recall",
 ]
 
-DEFAULT_KS: List[int] = [1, 5, 10, 20]
+DEFAULT_KS: list[int] = [1, 5, 10, 20]
 
 
 @dataclass
@@ -34,7 +35,7 @@ class _MetricRequirements:
     Stores description of metrics which need to be computed
     """
 
-    top_k: List[int]
+    top_k: list[int]
     need_recall: bool
     need_precision: bool
     need_ndcg: bool
@@ -68,14 +69,14 @@ class _MetricRequirements:
         self._metric_names = metrics
 
     @property
-    def metric_names(self) -> List[str]:
+    def metric_names(self) -> list[str]:
         """
         Getting metric names
         """
         return self._metric_names
 
     @classmethod
-    def from_metrics(cls, metrics: Set[str], top_k: List[int]) -> "_MetricRequirements":
+    def from_metrics(cls, metrics: set[str], top_k: list[int]) -> "_MetricRequirements":
         """
         Creating a class based on a given list of metrics and K values
         """
@@ -96,7 +97,7 @@ class _CoverageHelper:
     Computes coverage metric over multiple batches
     """
 
-    def __init__(self, top_k: List[int], item_count: Optional[int]) -> None:
+    def __init__(self, top_k: list[int], item_count: Optional[int]) -> None:
         """
         :param top_k: (list): Consider the highest k scores in the ranking.
         :param item_count: (optional, int): the total number of items in the dataset.
@@ -110,7 +111,7 @@ class _CoverageHelper:
         Reload the metric counter
         """
         self._train_hist = torch.zeros(self.item_count)
-        self._pred_hist: Dict[int, torch.Tensor] = {k: torch.zeros(self.item_count) for k in self._top_k}
+        self._pred_hist: dict[int, torch.Tensor] = {k: torch.zeros(self.item_count) for k in self._top_k}
 
     def _ensure_hists_on_device(self, device: torch.device) -> None:
         self._train_hist = self._train_hist.to(device)
@@ -197,8 +198,8 @@ class TorchMetricsBuilder(_MetricBuilder):
 
     def __init__(
         self,
-        metrics: List[MetricName] = DEFAULT_METRICS,
-        top_k: Optional[List[int]] = DEFAULT_KS,
+        metrics: list[MetricName] = DEFAULT_METRICS,
+        top_k: Optional[list[int]] = DEFAULT_KS,
         item_count: Optional[int] = None,
     ) -> None:
         """
@@ -331,8 +332,8 @@ class TorchMetricsBuilder(_MetricBuilder):
 
     def _compute_metrics_sum(
         self, predictions: torch.LongTensor, ground_truth: torch.LongTensor, train: Optional[torch.LongTensor]
-    ) -> List[float]:
-        result: List[float] = []
+    ) -> list[float]:
+        result: list[float] = []
 
         # Getting a tensor of the same size as predictions
         # The tensor contains information about whether the item from the prediction is present in the test set

replay/metrics/unexpectedness.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Union
+from typing import Optional, Union
 
 from replay.utils import PandasDataFrame, PolarsDataFrame, SparkDataFrame
 
@@ -152,7 +152,7 @@ class Unexpectedness(Metric):
         )
 
     @staticmethod
-    def _get_metric_value_by_user(ks: List[int], base_recs: Optional[List], recs: Optional[List]) -> List[float]:
+    def _get_metric_value_by_user(ks: list[int], base_recs: Optional[list], recs: Optional[list]) -> list[float]:
         if not base_recs or not recs:
             return [0.0 for _ in ks]
         return [1.0 - len(set(recs[:k]) & set(base_recs[:k])) / k for k in ks]

replay/models/als.py

@@ -1,5 +1,5 @@
 from os.path import join
-from typing import Optional, Tuple
+from typing import Optional
 
 from replay.data import Dataset
 from replay.utils import PYSPARK_AVAILABLE, SparkDataFrame
@@ -148,7 +148,7 @@ class ALSWrap(Recommender, ItemVectorModel):
 
     def _get_features(
         self, ids: SparkDataFrame, features: Optional[SparkDataFrame]  # noqa: ARG002
-    ) -> Tuple[Optional[SparkDataFrame], Optional[int]]:
+    ) -> tuple[Optional[SparkDataFrame], Optional[int]]:
         entity = "user" if self.query_column in ids.columns else "item"
         entity_col = self.query_column if self.query_column in ids.columns else self.item_column
 

replay/models/association_rules.py

@@ -1,4 +1,5 @@
-from typing import Any, Dict, Iterable, List, Optional, Union
+from collections.abc import Iterable
+from typing import Any, Optional, Union
 
 import numpy as np
 
@@ -97,13 +98,13 @@ class AssociationRulesItemRec(NeighbourRec):
     In this case all items in sessions should have the same rating.
     """
 
-    def _get_ann_infer_params(self) -> Dict[str, Any]:
+    def _get_ann_infer_params(self) -> dict[str, Any]:
         return {
             "features_col": None,
         }
 
     can_predict_item_to_item = True
-    item_to_item_metrics: List[str] = ["lift", "confidence", "confidence_gain"]
+    item_to_item_metrics: list[str] = ["lift", "confidence", "confidence_gain"]
     similarity: SparkDataFrame
     can_change_metric = True
     _search_space = {

replay/models/base_neighbour_rec.py

@@ -4,7 +4,8 @@ Part of set of abstract classes (from base_rec.py)
 """
 
 from abc import ABC
-from typing import Any, Dict, Iterable, Optional, Union
+from collections.abc import Iterable
+from typing import Any, Optional, Union
 
 from replay.data.dataset import Dataset
 from replay.utils import PYSPARK_AVAILABLE, MissingImport, SparkDataFrame
@@ -187,7 +188,7 @@ class NeighbourRec(ANNMixin, Recommender, ABC):
             "similarity" if metric is None else metric,
         )
 
-    def _configure_index_builder(self, interactions: SparkDataFrame) -> Dict[str, Any]:
+    def _configure_index_builder(self, interactions: SparkDataFrame) -> dict[str, Any]:
         similarity_df = self.similarity.select("similarity", "item_idx_one", "item_idx_two")
         self.index_builder.index_params.items_count = interactions.select(sf.max(self.item_column)).first()[0] + 1
         return similarity_df, {

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,5 @@
 import warnings
-from typing import List, Tuple, Union
+from typing import Union
 
 import numpy as np
 import pandas as pd
@@ -70,7 +70,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,7 +175,7 @@ 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
 
     def __init__(

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
@@ -295,7 +295,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.