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

replay/splitters/ratio_splitter.py

@@ -1,16 +1,17 @@
-import polars as pl
 from typing import List, Optional, Tuple
 
-from .base_splitter import Splitter
+import polars as pl
+
 from replay.utils import PYSPARK_AVAILABLE, DataFrameLike, PandasDataFrame, PolarsDataFrame, SparkDataFrame
 
+from .base_splitter import Splitter
+
 if PYSPARK_AVAILABLE:
     import pyspark.sql.functions as sf
     from pyspark.sql import Window
     from pyspark.sql.types import IntegerType
 
 
-# pylint: disable=too-few-public-methods, too-many-instance-attributes
 class RatioSplitter(Splitter):
     """
     Split interactions into train and test by ratio. Split is made for each user separately.
@@ -82,6 +83,7 @@ class RatioSplitter(Splitter):
     14        3        2 2020-01-05
     <BLANKLINE>
     """
+
     _init_arg_names = [
         "test_size",
         "divide_column",
@@ -96,7 +98,6 @@ class RatioSplitter(Splitter):
         "session_id_processing_strategy",
     ]
 
-    # pylint: disable=too-many-arguments
     def __init__(
         self,
         test_size: float,
@@ -160,7 +161,8 @@ class RatioSplitter(Splitter):
         self.min_interactions_per_group = min_interactions_per_group
         self.split_by_fractions = split_by_fractions
         if test_size < 0 or test_size > 1:
-            raise ValueError("test_size must between 0 and 1")
+            msg = "test_size must between 0 and 1"
+            raise ValueError(msg)
         self.test_size = test_size
 
     def _add_time_partition(self, interactions: DataFrameLike) -> DataFrameLike:
@@ -171,7 +173,8 @@ class RatioSplitter(Splitter):
         if isinstance(interactions, PolarsDataFrame):
             return self._add_time_partition_to_polars(interactions)
 
-        raise NotImplementedError(f"{self} is not implemented for {type(interactions)}")
+        msg = f"{self} is not implemented for {type(interactions)}"
+        raise NotImplementedError(msg)
 
     def _add_time_partition_to_pandas(self, interactions: PandasDataFrame) -> PandasDataFrame:
         res = interactions.copy(deep=True)
@@ -189,14 +192,8 @@ class RatioSplitter(Splitter):
         return res
 
     def _add_time_partition_to_polars(self, interactions: PolarsDataFrame) -> PolarsDataFrame:
-        res = (
-            interactions
-            .sort(self.timestamp_column)
-            .with_columns(
-                pl.cum_count(self.divide_column)
-                .over(self.divide_column)
-                .alias("row_num")
-            )
+        res = interactions.sort(self.timestamp_column).with_columns(
+            pl.cum_count(self.divide_column).over(self.divide_column).alias("row_num")
         )
 
         return res
@@ -262,8 +259,7 @@ class RatioSplitter(Splitter):
         self, interactions: PolarsDataFrame, train_size: float
     ) -> Tuple[PolarsDataFrame, PolarsDataFrame]:
         interactions = interactions.with_columns(
-            pl.count(self.timestamp_column).over(pl.col(self.divide_column))
-            .alias("count")
+            pl.count(self.timestamp_column).over(pl.col(self.divide_column)).alias("count")
         )
         if self.min_interactions_per_group is not None:
             interactions = interactions.with_columns(
@@ -274,18 +270,14 @@ class RatioSplitter(Splitter):
             )
         else:
             interactions = interactions.with_columns(
-                (pl.col("row_num") / pl.col("count")).round(self._precision)
-                .alias("frac")
+                (pl.col("row_num") / pl.col("count")).round(self._precision).alias("frac")
             )
 
-        interactions = interactions.with_columns(
-            (pl.col("frac") > train_size)
-            .alias("is_test")
-        )
+        interactions = interactions.with_columns((pl.col("frac") > train_size).alias("is_test"))
         if self.session_id_column:
             interactions = self._recalculate_with_session_id_column(interactions)
 
-        train = interactions.filter(~pl.col("is_test")).drop("row_num", "count", "frac", "is_test")  # pylint: disable=invalid-unary-operand-type
+        train = interactions.filter(~pl.col("is_test")).drop("row_num", "count", "frac", "is_test")
         test = interactions.filter(pl.col("is_test")).drop("row_num", "count", "frac", "is_test")
 
         return train, test
@@ -316,7 +308,7 @@ class RatioSplitter(Splitter):
                 "train_size",
             ] = (
                 interactions["train_size"] - 1
-            )  # pylint: disable=C0325
+            )
 
         interactions["is_test"] = interactions["row_num"] > interactions["train_size"]
         if self.session_id_column:
@@ -327,9 +319,7 @@ class RatioSplitter(Splitter):
 
         return train, test
 
-    def _partial_split_spark(
-        self, interactions: SparkDataFrame, ratio: float
-    ) -> Tuple[SparkDataFrame, SparkDataFrame]:
+    def _partial_split_spark(self, interactions: SparkDataFrame, ratio: float) -> Tuple[SparkDataFrame, SparkDataFrame]:
         interactions = interactions.withColumn(
             "count", sf.count(self.timestamp_column).over(Window.partitionBy(self.divide_column))
         )
@@ -364,51 +354,37 @@ class RatioSplitter(Splitter):
         self, interactions: PolarsDataFrame, ratio: float
     ) -> Tuple[PolarsDataFrame, PolarsDataFrame]:
         interactions = interactions.with_columns(
-            pl.count(self.timestamp_column).over(self.divide_column)
-            .alias("count")
+            pl.count(self.timestamp_column).over(self.divide_column).alias("count")
         )
         if self.min_interactions_per_group is not None:
             interactions = interactions.with_columns(
-                pl.when(
-                    pl.col("count") >= self.min_interactions_per_group
-                )
-                .then(
-                    pl.col("count") - (pl.col("count") * ratio).cast(interactions.get_column("count").dtype)
-                )
+                pl.when(pl.col("count") >= self.min_interactions_per_group)
+                .then(pl.col("count") - (pl.col("count") * ratio).cast(interactions.get_column("count").dtype))
                 .otherwise(pl.col("count"))
                 .alias("train_size")
             )
         else:
-            interactions = (
-                interactions
-                .with_columns(
-                    (pl.col("count") - (pl.col("count") * ratio).cast(interactions.get_column("count").dtype))
-                    .alias("train_size")
-                )
-                .with_columns(
-                    pl.when(
-                        (pl.col("count") * ratio > 0) & (pl.col("count") * ratio < 1) & (pl.col("train_size") > 1)
-                    )
-                    .then(pl.col("train_size") - 1)
-                    .otherwise(pl.col("train_size"))
-                    .alias("train_size")
+            interactions = interactions.with_columns(
+                (pl.col("count") - (pl.col("count") * ratio).cast(interactions.get_column("count").dtype)).alias(
+                    "train_size"
                 )
+            ).with_columns(
+                pl.when((pl.col("count") * ratio > 0) & (pl.col("count") * ratio < 1) & (pl.col("train_size") > 1))
+                .then(pl.col("train_size") - 1)
+                .otherwise(pl.col("train_size"))
+                .alias("train_size")
             )
 
-        interactions = interactions.with_columns(
-            (pl.col("row_num") > pl.col("train_size"))
-            .alias("is_test")
-        )
+        interactions = interactions.with_columns((pl.col("row_num") > pl.col("train_size")).alias("is_test"))
 
         if self.session_id_column:
             interactions = self._recalculate_with_session_id_column(interactions)
 
-        train = interactions.filter(~pl.col("is_test")).drop("row_num", "count", "train_size", "is_test")  # pylint: disable=invalid-unary-operand-type
+        train = interactions.filter(~pl.col("is_test")).drop("row_num", "count", "train_size", "is_test")
         test = interactions.filter(pl.col("is_test")).drop("row_num", "count", "train_size", "is_test")
 
         return train, test
 
-    # pylint: disable=invalid-name
     def _core_split(self, interactions: DataFrameLike) -> List[DataFrameLike]:
         if self.split_by_fractions:
             return self._partial_split_fractions(interactions, self.test_size)

replay/splitters/time_splitter.py

@@ -3,21 +3,21 @@ from typing import List, Optional, Tuple, Union
 
 import polars as pl
 
-from .base_splitter import Splitter
 from replay.utils import (
     PYSPARK_AVAILABLE,
     DataFrameLike,
     PandasDataFrame,
-    SparkDataFrame,
     PolarsDataFrame,
+    SparkDataFrame,
 )
 
+from .base_splitter import Splitter
+
 if PYSPARK_AVAILABLE:
     import pyspark.sql.functions as sf
     from pyspark.sql import Window
 
 
-# pylint: disable=too-few-public-methods
 class TimeSplitter(Splitter):
     """
     Split interactions by time.
@@ -85,6 +85,7 @@ class TimeSplitter(Splitter):
     14        3        2 2020-01-05
     <BLANKLINE>
     """
+
     _init_arg_names = [
         "time_threshold",
         "drop_cold_users",
@@ -97,10 +98,9 @@ class TimeSplitter(Splitter):
         "time_column_format",
     ]
 
-    # pylint: disable=too-many-arguments
     def __init__(
         self,
-        time_threshold: Union[datetime, str, int, float],
+        time_threshold: Union[datetime, str, float],
         query_column: str = "query_id",
         drop_cold_users: bool = False,
         drop_cold_items: bool = False,
@@ -144,7 +144,8 @@ class TimeSplitter(Splitter):
         self._precision = 3
         self.time_column_format = time_column_format
         if isinstance(time_threshold, float) and (time_threshold < 0 or time_threshold > 1):
-            raise ValueError("time_threshold must be between 0 and 1")
+            msg = "time_threshold must be between 0 and 1"
+            raise ValueError(msg)
         self.time_threshold = time_threshold
 
     def _partial_split(
@@ -160,7 +161,8 @@ class TimeSplitter(Splitter):
         if isinstance(interactions, PolarsDataFrame):
             return self._partial_split_polars(interactions, threshold)
 
-        raise NotImplementedError(f"{self} is not implemented for {type(interactions)}")
+        msg = f"{self} is not implemented for {type(interactions)}"
+        raise NotImplementedError(msg)
 
     def _partial_split_pandas(
         self, interactions: PandasDataFrame, threshold: Union[datetime, str, int]
@@ -191,9 +193,7 @@ class TimeSplitter(Splitter):
             )
             test_start = int(dates.count() * (1 - threshold)) + 1
             test_start = (
-                dates.filter(sf.col("_row_number_by_ts") == test_start)
-                .select(self.timestamp_column)
-                .collect()[0][0]
+                dates.filter(sf.col("_row_number_by_ts") == test_start).select(self.timestamp_column).collect()[0][0]
             )
             res = interactions.withColumn("is_test", sf.col(self.timestamp_column) >= test_start)
         else:
@@ -212,20 +212,15 @@ class TimeSplitter(Splitter):
         if isinstance(threshold, float):
             test_start = int(len(interactions) * (1 - threshold)) + 1
 
-            res = (
-                interactions
-                .sort(self.timestamp_column)
-                .with_columns(
-                    (pl.col(self.timestamp_column).cum_count() >= test_start)
-                    .alias("is_test")
-                )
+            res = interactions.sort(self.timestamp_column).with_columns(
+                (pl.col(self.timestamp_column).cum_count() >= test_start).alias("is_test")
             )
         else:
             res = interactions.with_columns((pl.col(self.timestamp_column) >= threshold).alias("is_test"))
 
         if self.session_id_column:
             res = self._recalculate_with_session_id_column(res)
-        train = res.filter(~pl.col("is_test")).drop("is_test")  # pylint: disable=invalid-unary-operand-type
+        train = res.filter(~pl.col("is_test")).drop("is_test")
         test = res.filter("is_test").drop("is_test")
 
         return train, test

replay/splitters/two_stage_splitter.py

@@ -1,18 +1,19 @@
 """
 This splitter split data by two columns.
 """
-from typing import Optional, Union
+from typing import Optional, Tuple
+
 import polars as pl
 
-from .base_splitter import Splitter, SplitterReturnType
 from replay.utils import PYSPARK_AVAILABLE, DataFrameLike, PandasDataFrame, PolarsDataFrame, SparkDataFrame
 
+from .base_splitter import Splitter, SplitterReturnType
+
 if PYSPARK_AVAILABLE:
     import pyspark.sql.functions as sf
     from pyspark.sql import Window
 
 
-# pylint: disable=too-few-public-methods
 class TwoStageSplitter(Splitter):
     """
     Split data by two columns.
@@ -73,11 +74,10 @@ class TwoStageSplitter(Splitter):
         "timestamp_column",
     ]
 
-    # pylint: disable=too-many-arguments
     def __init__(
         self,
-        first_divide_size: Union[float, int],
-        second_divide_size: Union[float, int],
+        first_divide_size: float,
+        second_divide_size: float,
         first_divide_column: str = "query_id",
         second_divide_column: str = "item_id",
         shuffle=False,
@@ -147,17 +147,12 @@ class TwoStageSplitter(Splitter):
             else:
                 value_error = True
         if value_error:
-            raise ValueError(
-                f"""
-            Invalid value for user_test_size: {self.first_divide_size}
-            """
-            )
+            msg = f"Invalid value for user_test_size: {self.first_divide_size}"
+            raise ValueError(msg)
         if isinstance(interactions, SparkDataFrame):
             test_users = (
                 all_values.withColumn("_rand", sf.rand(self.seed))
-                .withColumn(
-                    "_row_num", sf.row_number().over(Window.orderBy("_rand"))
-                )
+                .withColumn("_row_num", sf.row_number().over(Window.orderBy("_rand")))
                 .filter(f"_row_num <= {test_user_count}")
                 .drop("_rand", "_row_num")
             )
@@ -168,11 +163,9 @@ class TwoStageSplitter(Splitter):
 
         return test_users
 
-    def _split_proportion_spark(self, interactions: SparkDataFrame) -> Union[SparkDataFrame, SparkDataFrame]:
+    def _split_proportion_spark(self, interactions: SparkDataFrame) -> Tuple[SparkDataFrame, SparkDataFrame]:
         counts = interactions.groupBy(self.first_divide_column).count()
-        test_users = self._get_test_values(interactions).withColumn(
-            "is_test", sf.lit(True)
-        )
+        test_users = self._get_test_values(interactions).withColumn("is_test", sf.lit(True))
         if self.shuffle:
             res = self._add_random_partition_spark(
                 interactions.join(test_users, how="left", on=self.first_divide_column)
@@ -202,10 +195,10 @@ class TwoStageSplitter(Splitter):
 
         return train, test
 
-    def _split_proportion_pandas(self, interactions: PandasDataFrame) -> Union[PandasDataFrame, PandasDataFrame]:
-        counts = interactions.groupby(self.first_divide_column).agg(
-            count=(self.first_divide_column, "count")
-        ).reset_index()
+    def _split_proportion_pandas(self, interactions: PandasDataFrame) -> Tuple[PandasDataFrame, PandasDataFrame]:
+        counts = (
+            interactions.groupby(self.first_divide_column).agg(count=(self.first_divide_column, "count")).reset_index()
+        )
         test_users = self._get_test_values(interactions)
         test_users["is_test"] = True
         if self.shuffle:
@@ -229,11 +222,9 @@ class TwoStageSplitter(Splitter):
 
         return train, test
 
-    def _split_proportion_polars(self, interactions: PolarsDataFrame) -> Union[PolarsDataFrame, PolarsDataFrame]:
+    def _split_proportion_polars(self, interactions: PolarsDataFrame) -> Tuple[PolarsDataFrame, PolarsDataFrame]:
         counts = interactions.group_by(self.first_divide_column).count()
-        test_users = self._get_test_values(interactions).with_columns(
-            pl.lit(True).alias("is_test")
-        )
+        test_users = self._get_test_values(interactions).with_columns(pl.lit(True).alias("is_test"))
         if self.shuffle:
             res = self._add_random_partition_polars(
                 interactions.join(test_users, how="left", on=self.first_divide_column)
@@ -245,18 +236,15 @@ class TwoStageSplitter(Splitter):
             )
 
         res = res.join(counts, on=self.first_divide_column, how="left")
-        res = res.with_columns(
-            (pl.col("_row_num") / pl.col("count"))
-            .alias("_frac")
-        )
+        res = res.with_columns((pl.col("_row_num") / pl.col("count")).alias("_frac"))
         res = res.fill_null(False)
 
-        train = res.filter(
-            (pl.col("_frac") > self.second_divide_size) | (~pl.col("is_test"))  # pylint: disable=invalid-unary-operand-type
-        ).drop("_rand", "_row_num", "count", "_frac", "is_test")
-        test = res.filter(
-            (pl.col("_frac") <= self.second_divide_size) & pl.col("is_test")
-        ).drop("_rand", "_row_num", "count", "_frac", "is_test")
+        train = res.filter((pl.col("_frac") > self.second_divide_size) | (~pl.col("is_test"))).drop(
+            "_rand", "_row_num", "count", "_frac", "is_test"
+        )
+        test = res.filter((pl.col("_frac") <= self.second_divide_size) & pl.col("is_test")).drop(
+            "_rand", "_row_num", "count", "_frac", "is_test"
+        )
 
         return train, test
 
@@ -274,12 +262,11 @@ class TwoStageSplitter(Splitter):
         if isinstance(interactions, PolarsDataFrame):
             return self._split_proportion_polars(interactions)
 
-        raise NotImplementedError(f"{self} is not implemented for {type(interactions)}")
+        msg = f"{self} is not implemented for {type(interactions)}"
+        raise NotImplementedError(msg)
 
     def _split_quantity_spark(self, interactions: SparkDataFrame) -> SparkDataFrame:
-        test_users = self._get_test_values(interactions).withColumn(
-            "is_test", sf.lit(True)
-        )
+        test_users = self._get_test_values(interactions).withColumn("is_test", sf.lit(True))
         if self.shuffle:
             res = self._add_random_partition_spark(
                 interactions.join(test_users, how="left", on=self.first_divide_column)
@@ -328,9 +315,7 @@ class TwoStageSplitter(Splitter):
         return train, test
 
     def _split_quantity_polars(self, interactions: PolarsDataFrame) -> PolarsDataFrame:
-        test_users = self._get_test_values(interactions).with_columns(
-            pl.lit(True).alias("is_test")
-        )
+        test_users = self._get_test_values(interactions).with_columns(pl.lit(True).alias("is_test"))
         if self.shuffle:
             res = self._add_random_partition_polars(
                 interactions.join(test_users, how="left", on=self.first_divide_column)
@@ -342,12 +327,12 @@ class TwoStageSplitter(Splitter):
             )
 
         res = res.fill_null(False)
-        train = res.filter(
-            (pl.col("_row_num") > self.second_divide_size) | (~pl.col("is_test"))  # pylint: disable=invalid-unary-operand-type
-        ).drop("_row_num", "is_test")
-        test = res.filter(
-            (pl.col("_row_num") <= self.second_divide_size) & pl.col("is_test")
-        ).drop("_row_num", "is_test")
+        train = res.filter((pl.col("_row_num") > self.second_divide_size) | (~pl.col("is_test"))).drop(
+            "_row_num", "is_test"
+        )
+        test = res.filter((pl.col("_row_num") <= self.second_divide_size) & pl.col("is_test")).drop(
+            "_row_num", "is_test"
+        )
 
         return train, test
 
@@ -365,7 +350,8 @@ class TwoStageSplitter(Splitter):
         if isinstance(interactions, PolarsDataFrame):
             return self._split_quantity_polars(interactions)
 
-        raise NotImplementedError(f"{self} is not implemented for {type(interactions)}")
+        msg = f"{self} is not implemented for {type(interactions)}"
+        raise NotImplementedError(msg)
 
     def _core_split(self, interactions: DataFrameLike) -> SplitterReturnType:
         if 0 <= self.second_divide_size < 1.0:
@@ -373,11 +359,8 @@ class TwoStageSplitter(Splitter):
         elif self.second_divide_size >= 1 and isinstance(self.second_divide_size, int):
             train, test = self._split_quantity(interactions)
         else:
-            raise ValueError(
-                "`test_size` value must be [0, 1) or "
-                "a positive integer; "
-                f"test_size={self.second_divide_size}"
-            )
+            msg = f"`test_size` value must be [0, 1) or a positive integer; test_size={self.second_divide_size}"
+            raise ValueError(msg)
 
         return train, test
 
@@ -391,9 +374,7 @@ class TwoStageSplitter(Splitter):
         dataframe = dataframe.withColumn("_rand", sf.rand(self.seed))
         dataframe = dataframe.withColumn(
             "_row_num",
-            sf.row_number().over(
-                Window.partitionBy(self.first_divide_column).orderBy("_rand")
-            ),
+            sf.row_number().over(Window.partitionBy(self.first_divide_column).orderBy("_rand")),
         )
         return dataframe
 
@@ -404,14 +385,8 @@ class TwoStageSplitter(Splitter):
         return res
 
     def _add_random_partition_polars(self, dataframe: PolarsDataFrame) -> PolarsDataFrame:
-        res = (
-            dataframe
-            .sample(fraction=1, shuffle=True, seed=self.seed)
-            .with_columns(
-                pl.cum_count(self.first_divide_column)
-                .over(self.first_divide_column)
-                .alias("_row_num")
-            )
+        res = dataframe.sample(fraction=1, shuffle=True, seed=self.seed).with_columns(
+            pl.cum_count(self.first_divide_column).over(self.first_divide_column).alias("_row_num")
         )
         return res
 
@@ -431,11 +406,7 @@ class TwoStageSplitter(Splitter):
         """
         res = dataframe.withColumn(
             "_row_num",
-            sf.row_number().over(
-                Window.partitionBy(query_column).orderBy(
-                    sf.col(date_column).desc()
-                )
-            ),
+            sf.row_number().over(Window.partitionBy(query_column).orderBy(sf.col(date_column).desc())),
         )
         return res
 
@@ -456,13 +427,7 @@ class TwoStageSplitter(Splitter):
         query_column: str = "query_id",
         date_column: str = "timestamp",
     ) -> PolarsDataFrame:
-        res = (
-            dataframe
-            .sort(date_column, descending=True)
-            .with_columns(
-                pl.cum_count(query_column)
-                .over(query_column)
-                .alias("_row_num")
-            )
+        res = dataframe.sort(date_column, descending=True).with_columns(
+            pl.cum_count(query_column).over(query_column).alias("_row_num")
         )
         return res

replay/utils/__init__.py

@@ -7,6 +7,6 @@ from .types import (
     MissingImportType,
     NumType,
     PandasDataFrame,
-    SparkDataFrame,
     PolarsDataFrame,
+    SparkDataFrame,
 )

replay/utils/common.py

@@ -0,0 +1,65 @@
+import json
+from pathlib import Path
+from typing import Union
+
+from replay.splitters import (
+    ColdUserRandomSplitter,
+    KFolds,
+    LastNSplitter,
+    NewUsersSplitter,
+    RandomSplitter,
+    RatioSplitter,
+    TimeSplitter,
+    TwoStageSplitter,
+)
+from replay.utils import TORCH_AVAILABLE
+
+SavableObject = Union[
+    ColdUserRandomSplitter,
+    KFolds,
+    LastNSplitter,
+    NewUsersSplitter,
+    RandomSplitter,
+    RatioSplitter,
+    TimeSplitter,
+    TwoStageSplitter,
+]
+
+if TORCH_AVAILABLE:
+    from replay.data.nn import SequenceTokenizer
+
+    SavableObject = Union[
+        ColdUserRandomSplitter,
+        KFolds,
+        LastNSplitter,
+        NewUsersSplitter,
+        RandomSplitter,
+        RatioSplitter,
+        TimeSplitter,
+        TwoStageSplitter,
+        SequenceTokenizer,
+    ]
+
+
+def save_to_replay(obj: SavableObject, path: Union[str, Path]) -> None:
+    """
+    General function to save RePlay models, splitters and tokenizer.
+
+    :param path: Path to save the object.
+    """
+    obj.save(path)
+
+
+def load_from_replay(path: Union[str, Path]) -> SavableObject:
+    """
+    General function to load RePlay models, splitters and tokenizer.
+
+    :param path: Path to save the object.
+    """
+    path = Path(path).with_suffix(".replay").resolve()
+    with open(path / "init_args.json", "r") as file:
+        class_name = json.loads(file.read())["_class_name"]
+    obj_type = globals()[class_name]
+    obj = obj_type.load(path)
+
+    return obj