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

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

@@ -10,7 +10,6 @@ from replay.data.nn import (
     TorchSequentialDataset,
     TorchSequentialValidationDataset,
 )
-from replay.utils import deprecation_warning
 
 
 class SasRecTrainingBatch(NamedTuple):
@@ -31,17 +30,13 @@ class SasRecTrainingDataset(TorchDataset):
     Dataset that generates samples to train SasRec-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,
         sequence_shift: int = 1,
         sliding_window_step: Optional[None] = None,
-        padding_value: int = 0,
+        padding_value: Optional[int] = None,
         label_feature_name: Optional[str] = None,
     ) -> None:
         """
@@ -127,15 +122,11 @@ class SasRecPredictionDataset(TorchDataset):
     Dataset that generates samples to infer SasRec-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.
@@ -179,17 +170,13 @@ class SasRecValidationDataset(TorchDataset):
     Dataset that generates samples to infer and validate SasRec-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,
     ):
         """

replay/models/nn/sequential/sasrec/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
@@ -341,7 +341,7 @@ class SasRec(lightning.LightningModule):
         positive_labels: torch.LongTensor,
         padding_mask: torch.BoolTensor,
         target_padding_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
         positive_labels = cast(
@@ -428,7 +428,7 @@ class SasRec(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/sasrec/model.py

@@ -1,6 +1,6 @@
 import abc
 import contextlib
-from typing import Any, Dict, Optional, Tuple, Union, cast
+from typing import Any, Optional, Union, cast
 
 import torch
 
@@ -212,7 +212,7 @@ class SasRecMasks:
         self,
         feature_tensor: TensorMap,
         padding_mask: torch.BoolTensor,
-    ) -> Tuple[torch.BoolTensor, torch.BoolTensor, TensorMap]:
+    ) -> tuple[torch.BoolTensor, torch.BoolTensor, TensorMap]:
         """
         :param feature_tensor: Batch of features.
         :param padding_mask: Padding mask where 0 - <PAD>, 1 otherwise.
@@ -260,7 +260,7 @@ class BaseSasRecEmbeddings(abc.ABC):
         """
 
     @abc.abstractmethod
-    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 a layer as a dict.
         """
@@ -366,7 +366,7 @@ class SasRecEmbeddings(torch.nn.Module, BaseSasRecEmbeddings):
         # Last one is reserved for padding, so we remove it
         return self.item_emb.weight[:-1, :]
 
-    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.
         """
@@ -579,7 +579,7 @@ class TiSasRecEmbeddings(torch.nn.Module, BaseSasRecEmbeddings):
         self,
         feature_tensor: TensorMap,
         padding_mask: torch.BoolTensor,
-    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]]:
+    ) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]]:
         """
         :param feature_tensor: Batch of features.
         :param padding_mask: Padding mask where 0 - <PAD>, 1 otherwise.
@@ -628,7 +628,7 @@ class TiSasRecEmbeddings(torch.nn.Module, BaseSasRecEmbeddings):
         # Last one is reserved for padding, so we remove it
         return self.item_emb.weight[:-1, :]
 
-    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.
         """
@@ -674,7 +674,7 @@ class TiSasRecLayers(torch.nn.Module):
         seqs: torch.Tensor,
         attention_mask: torch.BoolTensor,
         padding_mask: torch.BoolTensor,
-        ti_embeddings: Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
+        ti_embeddings: tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
         device: torch.device,
     ) -> torch.Tensor:
         """
@@ -734,7 +734,7 @@ class TiSasRecAttention(torch.nn.Module):
         keys: torch.LongTensor,
         time_mask: torch.LongTensor,
         attn_mask: torch.LongTensor,
-        ti_embeddings: Tuple[torch.LongTensor, torch.LongTensor, torch.LongTensor, torch.LongTensor],
+        ti_embeddings: tuple[torch.LongTensor, torch.LongTensor, torch.LongTensor, torch.LongTensor],
         device: torch.device,
     ) -> torch.Tensor:
         """

replay/models/slim.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 import numpy as np
 import pandas as pd
@@ -21,7 +21,7 @@ class SLIM(NeighbourRec):
     """`SLIM: Sparse Linear Methods for Top-N Recommender Systems
     <http://glaros.dtc.umn.edu/gkhome/fetch/papers/SLIM2011icdm.pdf>`_"""
 
-    def _get_ann_infer_params(self) -> Dict[str, Any]:
+    def _get_ann_infer_params(self) -> dict[str, Any]:
         return {
             "features_col": None,
         }

replay/models/ucb.py

@@ -1,5 +1,5 @@
 import math
-from typing import Any, Dict, List, Optional
+from typing import Any, Optional
 
 from replay.data.dataset import Dataset
 from replay.metrics import NDCG, Metric
@@ -103,7 +103,7 @@ class UCB(NonPersonalizedRecommender):
         self,
         train_dataset: Dataset,  # noqa: ARG002
         test_dataset: Dataset,  # noqa: ARG002
-        param_borders: Optional[Dict[str, List[Any]]] = None,  # noqa: ARG002
+        param_borders: Optional[dict[str, list[Any]]] = None,  # noqa: ARG002
         criterion: Metric = NDCG,  # noqa: ARG002
         k: int = 10,  # noqa: ARG002
         budget: int = 10,  # noqa: ARG002

replay/models/word2vec.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Optional
+from typing import Any, Optional
 
 from replay.data import Dataset
 from replay.utils import PYSPARK_AVAILABLE, SparkDataFrame
@@ -24,7 +24,7 @@ class Word2VecRec(ANNMixin, Recommender, ItemVectorModel):
     Trains word2vec model where items are treated as words and queries as sentences.
     """
 
-    def _get_ann_infer_params(self) -> Dict[str, Any]:
+    def _get_ann_infer_params(self) -> dict[str, Any]:
         self.index_builder.index_params.dim = self.rank
         return {
             "features_col": "query_vector",
@@ -36,7 +36,7 @@ class Word2VecRec(ANNMixin, Recommender, ItemVectorModel):
         query_vectors = query_vectors.select(self.query_column, vector_to_array("query_vector").alias("query_vector"))
         return query_vectors
 
-    def _configure_index_builder(self, interactions: SparkDataFrame) -> Dict[str, Any]:
+    def _configure_index_builder(self, interactions: SparkDataFrame) -> dict[str, Any]:
         item_vectors = self._get_item_vectors()
         item_vectors = item_vectors.select(self.item_column, vector_to_array("item_vector").alias("item_vector"))
 

replay/preprocessing/discretizer.py

@@ -2,8 +2,9 @@ import abc
 import json
 import os
 import warnings
+from collections.abc import Sequence
 from pathlib import Path
-from typing import Dict, List, Literal, Sequence
+from typing import Literal
 
 import numpy as np
 import polars as pl
@@ -114,7 +115,7 @@ class GreedyDiscretizingRule(BaseDiscretizingRule):
         max_bin: int,
         total_cnt: int,
         min_data_in_bin: int,
-    ) -> List[float]:
+    ) -> list[float]:
         """
         Computes bound for bins.
 
@@ -363,7 +364,7 @@ Set 'keep' or 'skip' for processing NaN."
     @classmethod
     def load(cls, path: str) -> "GreedyDiscretizingRule":
         base_path = Path(path).with_suffix(".replay").resolve()
-        with open(base_path / "init_args.json", "r") as file:
+        with open(base_path / "init_args.json") as file:
             discretizer_rule_dict = json.loads(file.read())
 
         discretizer_rule = cls(**discretizer_rule_dict["init_args"])
@@ -590,7 +591,7 @@ Set 'keep' or 'skip' for processing NaN."
     @classmethod
     def load(cls, path: str) -> "QuantileDiscretizingRule":
         base_path = Path(path).with_suffix(".replay").resolve()
-        with open(base_path / "init_args.json", "r") as file:
+        with open(base_path / "init_args.json") as file:
             discretizer_rule_dict = json.loads(file.read())
 
         discretizer_rule = cls(**discretizer_rule_dict["init_args"])
@@ -655,7 +656,7 @@ class Discretizer:
         """
         return self.fit(df).transform(df)
 
-    def set_handle_invalid(self, handle_invalid_rules: Dict[str, HandleInvalidStrategies]) -> None:
+    def set_handle_invalid(self, handle_invalid_rules: dict[str, HandleInvalidStrategies]) -> None:
         """
         Modify handle_invalid strategy on already fitted Discretizer.
 
@@ -704,13 +705,13 @@ class Discretizer:
     @classmethod
     def load(cls, path: str) -> "Discretizer":
         base_path = Path(path).with_suffix(".replay").resolve()
-        with open(base_path / "init_args.json", "r") as file:
+        with open(base_path / "init_args.json") as file:
             discretizer_dict = json.loads(file.read())
         rules = []
         for root, dirs, files in os.walk(str(base_path) + "/rules/"):
             for d in dirs:
                 if d.split(".")[0] in discretizer_dict["rule_names"]:
-                    with open(root + d + "/init_args.json", "r") as file:
+                    with open(root + d + "/init_args.json") as file:
                         discretizer_rule_dict = json.loads(file.read())
                     rules.append(globals()[discretizer_rule_dict["_class_name"]].load(root + d))
 

replay/preprocessing/filters.py

@@ -4,7 +4,7 @@ Select or remove data by some criteria
 
 from abc import ABC, abstractmethod
 from datetime import datetime, timedelta
-from typing import Callable, Literal, Optional, Tuple, Union
+from typing import Callable, Literal, Optional, Union
 from uuid import uuid4
 
 import numpy as np
@@ -182,7 +182,7 @@ class InteractionEntriesFilter(_BaseFilter):
         non_agg_column: str,
         min_inter: Optional[int] = None,
         max_inter: Optional[int] = None,
-    ) -> Tuple[PandasDataFrame, int, int]:
+    ) -> tuple[PandasDataFrame, int, int]:
         filtered_interactions = interactions.copy(deep=True)
 
         filtered_interactions["count"] = filtered_interactions.groupby(agg_column, sort=False)[
@@ -207,7 +207,7 @@ class InteractionEntriesFilter(_BaseFilter):
         non_agg_column: str,
         min_inter: Optional[int] = None,
         max_inter: Optional[int] = None,
-    ) -> Tuple[SparkDataFrame, int, int]:
+    ) -> tuple[SparkDataFrame, int, int]:
         filtered_interactions = interactions.withColumn(
             "count", sf.count(non_agg_column).over(Window.partitionBy(agg_column))
         )
@@ -233,7 +233,7 @@ class InteractionEntriesFilter(_BaseFilter):
         non_agg_column: str,
         min_inter: Optional[int] = None,
         max_inter: Optional[int] = None,
-    ) -> Tuple[PolarsDataFrame, int, int]:
+    ) -> tuple[PolarsDataFrame, int, int]:
         filtered_interactions = interactions.with_columns(
             pl.col(non_agg_column).count().over(pl.col(agg_column)).alias("count")
         )

replay/preprocessing/history_based_fp.py

@@ -9,7 +9,7 @@ Contains classes for users' and items' features generation based on interactions
 """
 
 from datetime import datetime
-from typing import Dict, List, Optional
+from typing import Optional
 
 from replay.utils import PYSPARK_AVAILABLE, SparkDataFrame
 
@@ -64,7 +64,7 @@ class LogStatFeaturesProcessor(EmptyFeatureProcessor):
     user_log_features: Optional[SparkDataFrame] = None
     item_log_features: Optional[SparkDataFrame] = None
 
-    def _create_log_aggregates(self, agg_col: str = "user_idx") -> List:
+    def _create_log_aggregates(self, agg_col: str = "user_idx") -> list:
         """
         Create features based on relevance type
         (binary or not) and whether timestamp is present.
@@ -289,12 +289,12 @@ class ConditionalPopularityProcessor(EmptyFeatureProcessor):
     If user features are provided, item features will be generated and vice versa.
     """
 
-    conditional_pop_dict: Optional[Dict[str, SparkDataFrame]]
+    conditional_pop_dict: Optional[dict[str, SparkDataFrame]]
     entity_name: str
 
     def __init__(
         self,
-        cat_features_list: List,
+        cat_features_list: list,
     ):
         """
         :param cat_features_list: List of columns with categorical features to use
@@ -397,8 +397,8 @@ class HistoryBasedFeaturesProcessor:
         self,
         use_log_features: bool = True,
         use_conditional_popularity: bool = True,
-        user_cat_features_list: Optional[List] = None,
-        item_cat_features_list: Optional[List] = None,
+        user_cat_features_list: Optional[list] = None,
+        item_cat_features_list: Optional[list] = None,
     ):
         """
         :param use_log_features: if add statistical log-based features

replay/preprocessing/label_encoder.py

@@ -10,8 +10,9 @@ import abc
 import json
 import os
 import warnings
+from collections.abc import Mapping, Sequence
 from pathlib import Path
-from typing import Dict, List, Literal, Mapping, Optional, Sequence, Union
+from typing import Literal, Optional, Union
 
 import polars as pl
 
@@ -162,7 +163,7 @@ class LabelEncodingRule(BaseLabelEncodingRule):
     def _make_inverse_mapping(self) -> Mapping:
         return {val: key for key, val in self.get_mapping().items()}
 
-    def _make_inverse_mapping_list(self) -> List:
+    def _make_inverse_mapping_list(self) -> list:
         inverse_mapping_list = [0 for _ in range(len(self.get_mapping()))]
         for k, value in self.get_mapping().items():
             inverse_mapping_list[value] = k
@@ -543,7 +544,7 @@ Convert type to string, integer, or float."
     @classmethod
     def load(cls, path: str) -> "LabelEncodingRule":
         base_path = Path(path).with_suffix(".replay").resolve()
-        with open(base_path / "init_args.json", "r") as file:
+        with open(base_path / "init_args.json") as file:
             encoder_rule_dict = json.loads(file.read())
 
         string_column_type = encoder_rule_dict["fitted_args"]["column_type"]
@@ -901,7 +902,7 @@ class LabelEncoder:
         """
         return self.fit(df).transform(df)
 
-    def set_handle_unknowns(self, handle_unknown_rules: Dict[str, HandleUnknownStrategies]) -> None:
+    def set_handle_unknowns(self, handle_unknown_rules: dict[str, HandleUnknownStrategies]) -> None:
         """
         Modify handle unknown strategy on already fitted encoder.
 
@@ -923,7 +924,7 @@ class LabelEncoder:
             rule = list(filter(lambda x: x.column == column, self.rules))
             rule[0].set_handle_unknown(handle_unknown)
 
-    def set_default_values(self, default_value_rules: Dict[str, Optional[Union[int, str]]]) -> None:
+    def set_default_values(self, default_value_rules: dict[str, Optional[Union[int, str]]]) -> None:
         """
         Modify handle unknown strategy on already fitted encoder.
         Default value that will fill the unknown labels
@@ -974,13 +975,13 @@ class LabelEncoder:
     @classmethod
     def load(cls, path: str) -> "LabelEncoder":
         base_path = Path(path).with_suffix(".replay").resolve()
-        with open(base_path / "init_args.json", "r") as file:
+        with open(base_path / "init_args.json") as file:
             encoder_dict = json.loads(file.read())
         rules = []
         for root, dirs, files in os.walk(str(base_path) + "/rules/"):
             for d in dirs:
                 if d.split(".")[0] in encoder_dict["rule_names"]:
-                    with open(root + d + "/init_args.json", "r") as file:
+                    with open(root + d + "/init_args.json") as file:
                         encoder_rule_dict = json.loads(file.read())
                     rules.append(globals()[encoder_rule_dict["_class_name"]].load(root + d))
 

replay/scenarios/fallback.py

@@ -1,4 +1,5 @@
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+from collections.abc import Iterable
+from typing import Any, Optional, Union
 
 from replay.data import Dataset
 from replay.metrics import NDCG, Metric
@@ -125,12 +126,12 @@ class Fallback(BaseRecommender):
         self,
         train_dataset: Dataset,
         test_dataset: Dataset,
-        param_borders: Optional[Dict[str, Dict[str, List[Any]]]] = None,
+        param_borders: Optional[dict[str, dict[str, list[Any]]]] = None,
         criterion: Metric = NDCG,
         k: int = 10,
         budget: int = 10,
         new_study: bool = True,
-    ) -> Tuple[Dict[str, Any]]:
+    ) -> tuple[dict[str, Any]]:
         """
         Searches best parameters with optuna.
 

replay/splitters/base_splitter.py

@@ -1,7 +1,7 @@
 import json
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Optional, Tuple
+from typing import Optional
 
 import polars as pl
 
@@ -20,7 +20,7 @@ if PYSPARK_AVAILABLE:
     )
 
 
-SplitterReturnType = Tuple[DataFrameLike, DataFrameLike]
+SplitterReturnType = tuple[DataFrameLike, DataFrameLike]
 
 
 class Splitter(ABC):
@@ -90,7 +90,7 @@ class Splitter(ABC):
         Method for loading splitter from `.replay` directory.
         """
         base_path = Path(path).with_suffix(".replay").resolve()
-        with open(base_path / "init_args.json", "r") as file:
+        with open(base_path / "init_args.json") as file:
             splitter_dict = json.loads(file.read())
         splitter = cls(**splitter_dict["init_args"])
 

replay/splitters/cold_user_random_splitter.py

@@ -1,4 +1,4 @@
-from typing import Optional, Tuple
+from typing import Optional
 
 import polars as pl
 
@@ -62,7 +62,7 @@ class ColdUserRandomSplitter(Splitter):
 
     def _core_split_pandas(
         self, interactions: PandasDataFrame, threshold: float
-    ) -> Tuple[PandasDataFrame, PandasDataFrame]:
+    ) -> tuple[PandasDataFrame, PandasDataFrame]:
         users = PandasDataFrame(interactions[self.query_column].unique(), columns=[self.query_column])
         train_users = users.sample(frac=(1 - threshold), random_state=self.seed)
         train_users["is_test"] = False
@@ -78,7 +78,7 @@ class ColdUserRandomSplitter(Splitter):
 
     def _core_split_spark(
         self, interactions: SparkDataFrame, threshold: float
-    ) -> Tuple[SparkDataFrame, SparkDataFrame]:
+    ) -> tuple[SparkDataFrame, SparkDataFrame]:
         users = interactions.select(self.query_column).distinct()
         train_users, _ = users.randomSplit(
             [1 - threshold, threshold],
@@ -97,7 +97,7 @@ class ColdUserRandomSplitter(Splitter):
 
     def _core_split_polars(
         self, interactions: PolarsDataFrame, threshold: float
-    ) -> Tuple[PolarsDataFrame, PolarsDataFrame]:
+    ) -> tuple[PolarsDataFrame, PolarsDataFrame]:
         train_users = (
             interactions.select(self.query_column)
             .unique(maintain_order=True)

replay/splitters/k_folds.py

@@ -1,4 +1,4 @@
-from typing import Literal, Optional, Tuple
+from typing import Literal, Optional
 
 import polars as pl
 
@@ -83,7 +83,7 @@ class KFolds(Splitter):
         """
         return self._core_split(interactions)
 
-    def _query_split_spark(self, interactions: SparkDataFrame) -> Tuple[SparkDataFrame, SparkDataFrame]:
+    def _query_split_spark(self, interactions: SparkDataFrame) -> tuple[SparkDataFrame, SparkDataFrame]:
         dataframe = interactions.withColumn("_rand", sf.rand(self.seed))
         dataframe = dataframe.withColumn(
             "fold",
@@ -100,7 +100,7 @@ class KFolds(Splitter):
             test = self._drop_cold_items_and_users(train, test)
             yield train, test
 
-    def _query_split_pandas(self, interactions: PandasDataFrame) -> Tuple[PandasDataFrame, PandasDataFrame]:
+    def _query_split_pandas(self, interactions: PandasDataFrame) -> tuple[PandasDataFrame, PandasDataFrame]:
         dataframe = interactions.sample(frac=1, random_state=self.seed).sort_values(self.query_column)
         dataframe["fold"] = (dataframe.groupby(self.query_column, sort=False).cumcount() + 1) % self.n_folds
         for i in range(self.n_folds):
@@ -115,7 +115,7 @@ class KFolds(Splitter):
             test = self._drop_cold_items_and_users(train, test)
             yield train, test
 
-    def _query_split_polars(self, interactions: PolarsDataFrame) -> Tuple[PolarsDataFrame, PolarsDataFrame]:
+    def _query_split_polars(self, interactions: PolarsDataFrame) -> tuple[PolarsDataFrame, PolarsDataFrame]:
         dataframe = interactions.sample(fraction=1, shuffle=True, seed=self.seed).sort(self.query_column)
         dataframe = dataframe.with_columns(
             (pl.cum_count(self.query_column).over(self.query_column) % self.n_folds).alias("fold")

replay/splitters/last_n_splitter.py

@@ -1,4 +1,4 @@
-from typing import List, Literal, Optional, Tuple
+from typing import Literal, Optional
 
 import numpy as np
 import pandas as pd
@@ -240,7 +240,7 @@ class LastNSplitter(Splitter):
 
         return interactions
 
-    def _partial_split_interactions(self, interactions: DataFrameLike, n: int) -> Tuple[DataFrameLike, DataFrameLike]:
+    def _partial_split_interactions(self, interactions: DataFrameLike, n: int) -> tuple[DataFrameLike, DataFrameLike]:
         res = self._add_time_partition(interactions)
         if isinstance(interactions, SparkDataFrame):
             return self._partial_split_interactions_spark(res, n)
@@ -250,7 +250,7 @@ class LastNSplitter(Splitter):
 
     def _partial_split_interactions_pandas(
         self, interactions: PandasDataFrame, n: int
-    ) -> Tuple[PandasDataFrame, PandasDataFrame]:
+    ) -> tuple[PandasDataFrame, PandasDataFrame]:
         interactions["count"] = interactions.groupby(self.divide_column, sort=False)[self.divide_column].transform(len)
         interactions["is_test"] = interactions["row_num"] > (interactions["count"] - float(n))
         if self.session_id_column:
@@ -263,7 +263,7 @@ class LastNSplitter(Splitter):
 
     def _partial_split_interactions_spark(
         self, interactions: SparkDataFrame, n: int
-    ) -> Tuple[SparkDataFrame, SparkDataFrame]:
+    ) -> tuple[SparkDataFrame, SparkDataFrame]:
         interactions = interactions.withColumn(
             "count",
             sf.count(self.timestamp_column).over(Window.partitionBy(self.divide_column)),
@@ -281,7 +281,7 @@ class LastNSplitter(Splitter):
 
     def _partial_split_interactions_polars(
         self, interactions: PolarsDataFrame, n: int
-    ) -> Tuple[PolarsDataFrame, PolarsDataFrame]:
+    ) -> tuple[PolarsDataFrame, PolarsDataFrame]:
         interactions = interactions.with_columns(
             pl.col(self.timestamp_column).count().over(self.divide_column).alias("count")
         )
@@ -296,7 +296,7 @@ class LastNSplitter(Splitter):
 
     def _partial_split_timedelta(
         self, interactions: DataFrameLike, timedelta: int
-    ) -> Tuple[DataFrameLike, DataFrameLike]:
+    ) -> tuple[DataFrameLike, DataFrameLike]:
         if isinstance(interactions, SparkDataFrame):
             return self._partial_split_timedelta_spark(interactions, timedelta)
         if isinstance(interactions, PandasDataFrame):
@@ -305,7 +305,7 @@ class LastNSplitter(Splitter):
 
     def _partial_split_timedelta_pandas(
         self, interactions: PandasDataFrame, timedelta: int
-    ) -> Tuple[PandasDataFrame, PandasDataFrame]:
+    ) -> tuple[PandasDataFrame, PandasDataFrame]:
         res = interactions.copy(deep=True)
         res["diff_timestamp"] = (
             res.groupby(self.divide_column)[self.timestamp_column].transform(max) - res[self.timestamp_column]
@@ -321,7 +321,7 @@ class LastNSplitter(Splitter):
 
     def _partial_split_timedelta_spark(
         self, interactions: SparkDataFrame, timedelta: int
-    ) -> Tuple[SparkDataFrame, SparkDataFrame]:
+    ) -> tuple[SparkDataFrame, SparkDataFrame]:
         inter_with_max_time = interactions.withColumn(
             "max_timestamp",
             sf.max(self.timestamp_column).over(Window.partitionBy(self.divide_column)),
@@ -343,7 +343,7 @@ class LastNSplitter(Splitter):
 
     def _partial_split_timedelta_polars(
         self, interactions: PolarsDataFrame, timedelta: int
-    ) -> Tuple[PolarsDataFrame, PolarsDataFrame]:
+    ) -> tuple[PolarsDataFrame, PolarsDataFrame]:
         res = interactions.with_columns(
             (pl.col(self.timestamp_column).max().over(self.divide_column) - pl.col(self.timestamp_column)).alias(
                 "diff_timestamp"
@@ -358,7 +358,7 @@ class LastNSplitter(Splitter):
 
         return train, test
 
-    def _core_split(self, interactions: DataFrameLike) -> List[DataFrameLike]:
+    def _core_split(self, interactions: DataFrameLike) -> list[DataFrameLike]:
         if self.strategy == "timedelta":
             interactions = self._to_unix_timestamp(interactions)
         train, test = getattr(self, "_partial_split_" + self.strategy)(interactions, self.N)

replay/splitters/new_users_splitter.py

@@ -1,4 +1,4 @@
-from typing import Optional, Tuple
+from typing import Optional
 
 import polars as pl
 
@@ -100,7 +100,7 @@ class NewUsersSplitter(Splitter):
 
     def _core_split_pandas(
         self, interactions: PandasDataFrame, threshold: float
-    ) -> Tuple[PandasDataFrame, PandasDataFrame]:
+    ) -> tuple[PandasDataFrame, PandasDataFrame]:
         start_date_by_user = (
             interactions.groupby(self.query_column).agg(_start_dt_by_user=(self.timestamp_column, "min")).reset_index()
         )
@@ -134,7 +134,7 @@ class NewUsersSplitter(Splitter):
 
     def _core_split_spark(
         self, interactions: SparkDataFrame, threshold: float
-    ) -> Tuple[SparkDataFrame, SparkDataFrame]:
+    ) -> tuple[SparkDataFrame, SparkDataFrame]:
         start_date_by_user = interactions.groupby(self.query_column).agg(
             sf.min(self.timestamp_column).alias("_start_dt_by_user")
         )
@@ -171,7 +171,7 @@ class NewUsersSplitter(Splitter):
 
     def _core_split_polars(
         self, interactions: PolarsDataFrame, threshold: float
-    ) -> Tuple[PolarsDataFrame, PolarsDataFrame]:
+    ) -> tuple[PolarsDataFrame, PolarsDataFrame]:
         start_date_by_user = interactions.group_by(self.query_column).agg(
             pl.col(self.timestamp_column).min().alias("_start_dt_by_user")
         )