fixes #82

Author: Andrey-Matyashov

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

@@ -442,7 +442,7 @@ def _layers_stacker(self, num_blocks: int, layer_class: Any, *args, **kwargs) ->
 
 class SasRecNormalizer(torch.nn.Module):
     """
-    SasRec notmilization layers
+    SasRec normalization layers
 
     Link: https://arxiv.org/pdf/1808.09781.pdf
     """

replay/splitters/cold_user_random_splitter.py

@@ -38,12 +38,16 @@ def __init__(
         item_column: Optional[str] = "item_id",
     ):
         """
-        :param test_size: fraction of users to be in test
-        :param drop_cold_items: flag to drop cold items from test
-        :param drop_cold_users: flag to drop cold users from test
-        :param seed: random seed
-        :param query_column: query id column name
-        :param item_column: item id column name
+        :param test_size: The proportion of users to allocate to the test set.
+            Must be a float between 0.0 and 1.0.
+        :param drop_cold_items: Drop items from test DataFrame
+            which are not in train DataFrame, default: False.
+        :param seed: Seed for the random number generator to ensure
+            reproducibility of the split, default: None.
+        :param query_column: Name of query interaction column.
+            default: ``query_id``.
+        :param item_column: Name of item interaction column.
+            default: ``item_id``.
         """
         super().__init__(
             drop_cold_items=drop_cold_items,
@@ -81,7 +85,9 @@ def _core_split_spark(
             seed=self.seed,
         )
         interactions = interactions.join(
-            train_users.withColumn("is_test", sf.lit(False)), on=self.query_column, how="left"
+            train_users.withColumn("is_test", sf.lit(False)),
+            on=self.query_column,
+            how="left",
         ).na.fill({"is_test": True})
 
         train = interactions.filter(~sf.col("is_test")).drop("is_test")

replay/splitters/last_n_splitter.py

@@ -4,7 +4,13 @@
 import pandas as pd
 import polars as pl
 
-from replay.utils import PYSPARK_AVAILABLE, DataFrameLike, PandasDataFrame, PolarsDataFrame, SparkDataFrame
+from replay.utils import (
+    PYSPARK_AVAILABLE,
+    DataFrameLike,
+    PandasDataFrame,
+    PolarsDataFrame,
+    SparkDataFrame,
+)
 
 from .base_splitter import Splitter
 
@@ -118,14 +124,12 @@ def __init__(
         session_id_processing_strategy: str = "test",
     ):
         """
-        :param N: Array of interactions/timedelta to split.
+        :param N: Number of last interactions or size of the time window in seconds
         :param divide_column: Name of column for dividing
             in dataframe, default: ``query_id``.
-        :param time_column_format: Format of time_column,
-            needs for convert time_column into unix_timestamp type.
-            If strategy is set to 'interactions', then you can omit this parameter.
-            If time_column has already transformed into unix_timestamp type,
-            then you can omit this parameter.
+        :param time_column_format: Format of the timestamp column,
+            used for converting string dates to a numerical timestamp when strategy is 'timedelta'.
+            If the column is already a datetime object or a numerical timestamp, this parameter is ignored.
             default: ``yyyy-MM-dd HH:mm:ss``
         :param strategy: Defines the type of data splitting.
             Must be ``interactions`` or ``timedelta``.
@@ -223,7 +227,8 @@ def _to_unix_timestamp_spark(self, interactions: SparkDataFrame) -> SparkDataFra
         time_column_type = dict(interactions.dtypes)[self.timestamp_column]
         if time_column_type == "date":
             interactions = interactions.withColumn(
-                self.timestamp_column, sf.unix_timestamp(self.timestamp_column, self.time_column_format)
+                self.timestamp_column,
+                sf.unix_timestamp(self.timestamp_column, self.time_column_format),
             )
 
         return interactions
@@ -260,7 +265,8 @@ def _partial_split_interactions_spark(
         self, interactions: SparkDataFrame, n: int
     ) -> Tuple[SparkDataFrame, SparkDataFrame]:
         interactions = interactions.withColumn(
-            "count", sf.count(self.timestamp_column).over(Window.partitionBy(self.divide_column))
+            "count",
+            sf.count(self.timestamp_column).over(Window.partitionBy(self.divide_column)),
         )
         # float(n) - because DataFrame.filter is changing order
         # of sorted DataFrame to descending
@@ -317,7 +323,8 @@ def _partial_split_timedelta_spark(
         self, interactions: SparkDataFrame, timedelta: int
     ) -> Tuple[SparkDataFrame, SparkDataFrame]:
         inter_with_max_time = interactions.withColumn(
-            "max_timestamp", sf.max(self.timestamp_column).over(Window.partitionBy(self.divide_column))
+            "max_timestamp",
+            sf.max(self.timestamp_column).over(Window.partitionBy(self.divide_column)),
         )
         inter_with_diff = inter_with_max_time.withColumn(
             "diff_timestamp", sf.col("max_timestamp") - sf.col(self.timestamp_column)