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

replay/metrics/base_metric.py

@@ -1,6 +1,7 @@
 import warnings
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Mapping, Optional, Union
+from collections.abc import Mapping
+from typing import Any, Optional, Union
 
 import numpy as np
 import polars as pl
@@ -14,7 +15,7 @@ if PYSPARK_AVAILABLE:
     from pyspark.sql.types import ArrayType, DoubleType, StructType
 
 
-MetricsDataFrameLike = Union[DataFrameLike, Dict]
+MetricsDataFrameLike = Union[DataFrameLike, dict]
 MetricsMeanReturnType = Mapping[str, float]
 MetricsPerUserReturnType = Mapping[str, Mapping[Any, float]]
 MetricsReturnType = Union[MetricsMeanReturnType, MetricsPerUserReturnType]
@@ -29,7 +30,7 @@ class Metric(ABC):
 
     def __init__(
         self,
-        topk: Union[List[int], int],
+        topk: Union[list[int], int],
         query_column: str = "query_id",
         item_column: str = "item_id",
         rating_column: str = "rating",
@@ -89,7 +90,7 @@ class Metric(ABC):
         if duplicates_count:
             self._duplicate_warn()
 
-    def _check_duplicates_dict(self, recommendations: Dict) -> None:
+    def _check_duplicates_dict(self, recommendations: dict) -> None:
         for items in recommendations.values():
             items_set = set(items)
             if len(items) != len(items_set):
@@ -143,7 +144,7 @@ class Metric(ABC):
             ground_truth=ground_truth,
         )
 
-    def _convert_pandas_to_dict_with_score(self, data: PandasDataFrame) -> Dict:
+    def _convert_pandas_to_dict_with_score(self, data: PandasDataFrame) -> dict:
         return (
             data.sort_values(by=[self.rating_column, self.item_column], ascending=False, kind="stable")
             .groupby(self.query_column)[self.item_column]
@@ -151,7 +152,7 @@ class Metric(ABC):
             .to_dict()
         )
 
-    def _convert_dict_to_dict_with_score(self, data: Dict) -> Dict:
+    def _convert_dict_to_dict_with_score(self, data: dict) -> dict:
         converted_data = {}
         for user, items in data.items():
             is_sorted = True
@@ -164,10 +165,10 @@ class Metric(ABC):
             converted_data[user] = [item for item, _ in items]
         return converted_data
 
-    def _convert_pandas_to_dict_without_score(self, data: PandasDataFrame) -> Dict:
+    def _convert_pandas_to_dict_without_score(self, data: PandasDataFrame) -> dict:
         return data.groupby(self.query_column)[self.item_column].apply(list).to_dict()
 
-    def _dict_call(self, users: List, **kwargs: Dict) -> MetricsReturnType:
+    def _dict_call(self, users: list, **kwargs: dict) -> MetricsReturnType:
         """
         Calculating metrics in dict format.
         kwargs can contain different dicts (for example, ground_truth or train), it depends on the metric.
@@ -287,7 +288,7 @@ class Metric(ABC):
         )
         return self._rearrange_columns(enriched_recommendations)
 
-    def _aggregate_results_per_user(self, distribution_per_user: Dict[Any, List[float]]) -> MetricsPerUserReturnType:
+    def _aggregate_results_per_user(self, distribution_per_user: dict[Any, list[float]]) -> MetricsPerUserReturnType:
         res: MetricsPerUserReturnType = {}
         for index, val in enumerate(self.topk):
             metric_name = f"{self.__name__}@{val}"
@@ -374,7 +375,7 @@ class Metric(ABC):
 
     @staticmethod
     @abstractmethod
-    def _get_metric_value_by_user(ks: List[int], *args: List) -> List[float]:  # pragma: no cover
+    def _get_metric_value_by_user(ks: list[int], *args: list) -> list[float]:  # pragma: no cover
         """
         Metric calculation for one user.
 

replay/metrics/categorical_diversity.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
@@ -62,7 +62,7 @@ class CategoricalDiversity(Metric):
 
     def __init__(
         self,
-        topk: Union[List, int],
+        topk: Union[list, int],
         query_column: str = "query_id",
         category_column: str = "category_id",
         rating_column: str = "rating",
@@ -195,7 +195,7 @@ class CategoricalDiversity(Metric):
             return self._polars_compute_per_user(recs)
         return self._polars_compute_agg(recs)
 
-    def _convert_pandas_to_dict_with_score(self, data: PandasDataFrame) -> Dict:
+    def _convert_pandas_to_dict_with_score(self, data: PandasDataFrame) -> dict:
         return (
             data.sort_values(by=self.rating_column, ascending=False)
             .groupby(self.query_column)[self.category_column]
@@ -203,7 +203,7 @@ class CategoricalDiversity(Metric):
             .to_dict()
         )
 
-    def _precalculate_unique_cats(self, recommendations: Dict) -> Dict:
+    def _precalculate_unique_cats(self, recommendations: dict) -> dict:
         """
         Precalculate unique categories for each prefix for each user.
         """
@@ -217,14 +217,14 @@ class CategoricalDiversity(Metric):
             answer[user] = unique_len
         return answer
 
-    def _dict_compute_per_user(self, precalculated_answer: Dict) -> MetricsPerUserReturnType:
+    def _dict_compute_per_user(self, precalculated_answer: dict) -> MetricsPerUserReturnType:
         distribution_per_user = defaultdict(list)
         for k in self.topk:
             for user, unique_cats in precalculated_answer.items():
                 distribution_per_user[user].append(unique_cats[min(len(unique_cats), k) - 1] / k)
         return self._aggregate_results_per_user(distribution_per_user)
 
-    def _dict_compute_mean(self, precalculated_answer: Dict) -> MetricsMeanReturnType:
+    def _dict_compute_mean(self, precalculated_answer: dict) -> MetricsMeanReturnType:
         distribution_list = []
         for unique_cats in precalculated_answer.values():
             metrics_per_user = []
@@ -238,7 +238,7 @@ class CategoricalDiversity(Metric):
         metrics = [self._mode.cpu(distribution[:, k]) for k in range(distribution.shape[1])]
         return self._aggregate_results(metrics)
 
-    def _dict_call(self, precalculated_answer: Dict) -> MetricsReturnType:
+    def _dict_call(self, precalculated_answer: dict) -> MetricsReturnType:
         """
         Calculating metrics in dict format.
         """
@@ -247,5 +247,5 @@ class CategoricalDiversity(Metric):
         return self._dict_compute_mean(precalculated_answer)
 
     @staticmethod
-    def _get_metric_value_by_user(ks: List[int], *args: List) -> List[float]:  # pragma: no cover
+    def _get_metric_value_by_user(ks: list[int], *args: list) -> list[float]:  # pragma: no cover
         pass

replay/metrics/coverage.py

@@ -1,6 +1,6 @@
 import functools
 import operator
-from typing import Dict, List, Union
+from typing import Union
 
 import polars as pl
 
@@ -60,7 +60,7 @@ class Coverage(Metric):
 
     def __init__(
         self,
-        topk: Union[List, int],
+        topk: Union[list, int],
         query_column: str = "query_id",
         item_column: str = "item_id",
         rating_column: str = "rating",
@@ -173,7 +173,7 @@ class Coverage(Metric):
         recs = self._get_enriched_recommendations(recommendations)
         return self._polars_compute(recs, train)
 
-    def _dict_call(self, recommendations: Dict, train: Dict) -> MetricsReturnType:
+    def _dict_call(self, recommendations: dict, train: dict) -> MetricsReturnType:
         """
         Calculating metrics in dict format.
         """
@@ -229,5 +229,5 @@ class Coverage(Metric):
         return self._dict_call(recommendations, train)
 
     @staticmethod
-    def _get_metric_value_by_user(ks, *args) -> List[float]:  # pragma: no cover
+    def _get_metric_value_by_user(ks, *args) -> list[float]:  # pragma: no cover
         pass

replay/metrics/experiment.py

@@ -1,4 +1,4 @@
-from typing import Dict, List, Optional, Union
+from typing import Optional, Union
 
 import pandas as pd
 
@@ -102,10 +102,10 @@ class Experiment:
 
     def __init__(
         self,
-        metrics: List[Metric],
+        metrics: list[Metric],
         ground_truth: MetricsDataFrameLike,
         train: Optional[MetricsDataFrameLike] = None,
-        base_recommendations: Optional[Union[MetricsDataFrameLike, Dict[str, MetricsDataFrameLike]]] = None,
+        base_recommendations: Optional[Union[MetricsDataFrameLike, dict[str, MetricsDataFrameLike]]] = None,
         query_column: str = "query_id",
         item_column: str = "item_id",
         rating_column: str = "rating",

replay/metrics/hitrate.py

@@ -1,5 +1,3 @@
-from typing import List
-
 from .base_metric import Metric
 
 
@@ -62,7 +60,7 @@ class HitRate(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/map.py

@@ -1,5 +1,3 @@
-from typing import List
-
 from .base_metric import Metric
 
 
@@ -63,7 +61,7 @@ class MAP(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]
         res = []

replay/metrics/mrr.py

@@ -1,5 +1,3 @@
-from typing import List
-
 from .base_metric import Metric
 
 
@@ -55,7 +53,7 @@ class MRR(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/ndcg.py

@@ -1,5 +1,4 @@
 import math
-from typing import List
 
 from .base_metric import Metric
 
@@ -80,7 +79,7 @@ class NDCG(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 pred or not ground_truth:
             return [0.0 for _ in ks]
         set_gt = set(ground_truth)

replay/metrics/novelty.py

@@ -1,11 +1,11 @@
-from typing import TYPE_CHECKING, List, Type
+from typing import TYPE_CHECKING
 
 from replay.utils import PandasDataFrame, PolarsDataFrame, SparkDataFrame
 
 from .base_metric import Metric, MetricsDataFrameLike, MetricsReturnType
 
 if TYPE_CHECKING:  # pragma: no cover
-    __class__: Type
+    __class__: type
 
 
 class Novelty(Metric):
@@ -139,7 +139,7 @@ class Novelty(Metric):
         return self._polars_compute(recs)
 
     @staticmethod
-    def _get_metric_value_by_user(ks: List[int], pred: List, train: List) -> List[float]:
+    def _get_metric_value_by_user(ks: list[int], pred: list, train: list) -> list[float]:
         if not train or not pred:
             return [1.0 for _ in ks]
         set_train = set(train)

replay/metrics/offline_metrics.py

@@ -1,5 +1,5 @@
 import warnings
-from typing import Dict, List, Optional, Tuple, Union
+from typing import Optional, Union
 
 from replay.utils import PandasDataFrame, PolarsDataFrame, SparkDataFrame
 
@@ -132,7 +132,7 @@ class OfflineMetrics:
     <BLANKLINE>
     """
 
-    _metrics_call_requirement_map: Dict[str, List[str]] = {
+    _metrics_call_requirement_map: dict[str, list[str]] = {
         "HitRate": ["ground_truth"],
         "MAP": ["ground_truth"],
         "NDCG": ["ground_truth"],
@@ -147,7 +147,7 @@ class OfflineMetrics:
 
     def __init__(
         self,
-        metrics: List[Metric],
+        metrics: list[Metric],
         query_column: str = "query_id",
         item_column: str = "item_id",
         rating_column: str = "rating",
@@ -174,9 +174,9 @@ class OfflineMetrics:
         :param allow_caching: (bool): The flag for using caching to optimize calculations.
             Default: ``True``.
         """
-        self.unexpectedness_metric: List[Metric] = []
-        self.diversity_metric: List[Metric] = []
-        self.main_metrics: List[Metric] = []
+        self.unexpectedness_metric: list[Metric] = []
+        self.diversity_metric: list[Metric] = []
+        self.main_metrics: list[Metric] = []
         self._allow_caching = allow_caching
 
         for metric in metrics:
@@ -198,7 +198,7 @@ class OfflineMetrics:
         recommendations: Union[SparkDataFrame, PolarsDataFrame],
         ground_truth: Union[SparkDataFrame, PolarsDataFrame],
         train: Optional[Union[SparkDataFrame, PolarsDataFrame]],
-    ) -> Tuple[Dict[str, Union[SparkDataFrame, PolarsDataFrame]], Optional[Union[SparkDataFrame, PolarsDataFrame]]]:
+    ) -> tuple[dict[str, Union[SparkDataFrame, PolarsDataFrame]], Optional[Union[SparkDataFrame, PolarsDataFrame]]]:
         if len(self.main_metrics) == 0:
             return {}, train
         result_dict = {}
@@ -257,21 +257,21 @@ class OfflineMetrics:
 
         return result_dict, train
 
-    def _cache_dataframes(self, dataframes: Dict[str, SparkDataFrame]) -> None:
+    def _cache_dataframes(self, dataframes: dict[str, SparkDataFrame]) -> None:
         for data in dataframes.values():
             data.cache()
 
-    def _unpersist_dataframes(self, dataframes: Dict[str, SparkDataFrame]) -> None:
+    def _unpersist_dataframes(self, dataframes: dict[str, SparkDataFrame]) -> None:
         for data in dataframes.values():
             data.unpersist()
 
     def _calculate_metrics(
         self,
-        enriched_recs_dict: Dict[str, Union[SparkDataFrame, PolarsDataFrame]],
+        enriched_recs_dict: dict[str, Union[SparkDataFrame, PolarsDataFrame]],
         train: Optional[Union[SparkDataFrame, PolarsDataFrame]] = None,
         is_spark: bool = True,
     ) -> MetricsReturnType:
-        result: Dict = {}
+        result: dict = {}
         for metric in self.metrics:
             metric_args = {}
             if metric.__class__.__name__ == "Coverage" and train is not None:
@@ -295,7 +295,7 @@ class OfflineMetrics:
         recommendations: MetricsDataFrameLike,
         ground_truth: MetricsDataFrameLike,
         train: Optional[MetricsDataFrameLike],
-        base_recommendations: Optional[Union[MetricsDataFrameLike, Dict[str, MetricsDataFrameLike]]],
+        base_recommendations: Optional[Union[MetricsDataFrameLike, dict[str, MetricsDataFrameLike]]],
     ) -> None:
         types = set()
         types.add(type(recommendations))
@@ -379,8 +379,8 @@ class OfflineMetrics:
         recommendations: MetricsDataFrameLike,
         ground_truth: MetricsDataFrameLike,
         train: Optional[MetricsDataFrameLike] = None,
-        base_recommendations: Optional[Union[MetricsDataFrameLike, Dict[str, MetricsDataFrameLike]]] = None,
-    ) -> Dict[str, float]:
+        base_recommendations: Optional[Union[MetricsDataFrameLike, dict[str, MetricsDataFrameLike]]] = None,
+    ) -> dict[str, float]:
         """
         Compute metrics.
 
@@ -450,12 +450,12 @@ class OfflineMetrics:
             if is_spark and self._allow_caching:
                 self._unpersist_dataframes(enriched_recs_dict)
         else:  # Calculating metrics in dict format
-            current_map: Dict[str, Union[PandasDataFrame, Dict]] = {
+            current_map: dict[str, Union[PandasDataFrame, dict]] = {
                 "ground_truth": ground_truth,
                 "train": train,
             }
             for metric in self.metrics:
-                args_to_call: Dict[str, Union[PandasDataFrame, Dict]] = {"recommendations": recommendations}
+                args_to_call: dict[str, Union[PandasDataFrame, dict]] = {"recommendations": recommendations}
                 for data_name in self._metrics_call_requirement_map[str(metric.__class__.__name__)]:
                     args_to_call[data_name] = current_map[data_name]
                 result.update(metric(**args_to_call))

replay/metrics/precision.py

@@ -1,5 +1,3 @@
-from typing import List
-
 from .base_metric import Metric
 
 
@@ -61,7 +59,7 @@ class Precision(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/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, {