replay-rec 0.17.0rc0__tar.gz → 0.17.1__tar.gz

{replay_rec-0.17.0rc0 → replay_rec-0.17.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: replay-rec
-Version: 0.17.0rc0
+Version: 0.17.1
 Summary: RecSys Library
 Home-page: https://sb-ai-lab.github.io/RePlay/
 License: Apache-2.0
@@ -20,25 +20,17 @@ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
 Provides-Extra: all
 Provides-Extra: spark
 Provides-Extra: torch
-Requires-Dist: d3rlpy (>=2.0.4,<3.0.0)
-Requires-Dist: gym (>=0.26.0,<0.27.0)
 Requires-Dist: hnswlib (==0.7.0)
-Requires-Dist: implicit (>=0.7.0,<0.8.0)
-Requires-Dist: lightautoml (>=0.3.1,<0.4.0)
-Requires-Dist: lightfm (==1.17)
 Requires-Dist: lightning (>=2.0.2,<3.0.0) ; extra == "torch" or extra == "all"
-Requires-Dist: llvmlite (>=0.32.1)
 Requires-Dist: nmslib (==2.1.1)
-Requires-Dist: numba (>=0.50)
 Requires-Dist: numpy (>=1.20.0)
 Requires-Dist: optuna (>=3.2.0,<3.3.0)
-Requires-Dist: pandas (>=1.3.5,<2.0.0)
+Requires-Dist: pandas (>=1.3.5,<=2.2.2)
 Requires-Dist: polars (>=0.20.7,<0.21.0)
 Requires-Dist: psutil (>=5.9.5,<5.10.0)
 Requires-Dist: pyarrow (>=12.0.1)
-Requires-Dist: pyspark (>=3.0,<3.3) ; extra == "spark" or extra == "all"
+Requires-Dist: pyspark (>=3.0,<3.5) ; extra == "spark" or extra == "all"
 Requires-Dist: pytorch-ranger (>=0.1.1,<0.2.0) ; extra == "torch" or extra == "all"
-Requires-Dist: sb-obp (>=0.5.7,<0.6.0)
 Requires-Dist: scikit-learn (>=1.0.2,<2.0.0)
 Requires-Dist: scipy (>=1.8.1,<1.9.0)
 Requires-Dist: torch (>=1.8,<2.0) ; extra == "torch" or extra == "all"

{replay_rec-0.17.0rc0 → replay_rec-0.17.1}/pyproject.toml

@@ -39,33 +39,26 @@ classifiers = [
 ]
 exclude = [
     "replay/conftest.py",
+    "replay/experimental",
 ]
-version = "0.17.0.preview"
+version = "0.17.1"
 
 [tool.poetry.dependencies]
 python = ">=3.8.1, <3.11"
 numpy = ">=1.20.0"
-pandas = "^1.3.5"
+pandas = ">=1.3.5,<=2.2.2"
 polars = "~0.20.7"
 optuna = "~3.2.0"
 scipy = "~1.8.1"
 psutil = "~5.9.5"
-pyspark = {version = ">=3.0,<3.3", optional = true}
+pyspark = {version = ">=3.0,<3.5", optional = true}
 scikit-learn = "^1.0.2"
 pyarrow = ">=12.0.1"
+torch = {version = "^1.8", optional = true}
+lightning = {version = "^2.0.2", optional = true}
+pytorch-ranger = {version = "^0.1.1", optional = true}
 nmslib = "2.1.1"
 hnswlib = "0.7.0"
-torch = "^1.8"
-lightning = "^2.0.2"
-pytorch-ranger = "^0.1.1"
-lightfm = "1.17"
-lightautoml = "~0.3.1"
-numba = ">=0.50"
-llvmlite = ">=0.32.1"
-sb-obp = "^0.5.7"
-d3rlpy = "^2.0.4"
-implicit = "~0.7.0"
-gym = "^0.26.0"
 
 [tool.poetry.extras]
 spark = ["pyspark"]
@@ -92,7 +85,7 @@ data-science-types = "0.2.23"
 
 [tool.poetry-dynamic-versioning]
 enable = false
-format-jinja = """0.17.0{{ env['PACKAGE_SUFFIX'] }}"""
+format-jinja = """0.17.1{{ env['PACKAGE_SUFFIX'] }}"""
 vcs = "git"
 
 [tool.ruff]
@@ -123,6 +116,9 @@ max-complexity = 13
 "tests/*" = ["ARG", "E402", "INP", "ISC", "N", "S", "SIM", "F811"]
 "tests/experimental/*" = ["F401", "F811"]
 "replay/experimental/models/extensions/spark_custom_models/als_extension.py" = ["ARG002", "N802", "N803", "N815"]
+"replay/data/nn/sequence_tokenizer.py" = ["ARG003"]
+"replay/splitters/base_splitter.py" = ["ARG003"]
+"replay/data/nn/sequential_dataset.py" = ["ARG003"]
 
 [tool.tomlsort]
 ignore_case = true

replay_rec-0.17.1/replay/__init__.py

@@ -0,0 +1,2 @@
+""" RecSys library """
+__version__ = "0.17.1"

{replay_rec-0.17.0rc0 → replay_rec-0.17.1}/replay/data/dataset.py

@@ -3,11 +3,22 @@
 """
 from __future__ import annotations
 
-from typing import Callable, Dict, Iterable, List, Optional, Sequence
+import json
+from pathlib import Path
+from typing import Callable, Dict, Iterable, List, Optional, Sequence, Union
 
 import numpy as np
-
-from replay.utils import PYSPARK_AVAILABLE, DataFrameLike, PandasDataFrame, PolarsDataFrame, SparkDataFrame
+from pandas import read_parquet as pd_read_parquet
+from polars import read_parquet as pl_read_parquet
+
+from replay.utils import (
+    PYSPARK_AVAILABLE,
+    DataFrameLike,
+    PandasDataFrame,
+    PolarsDataFrame,
+    SparkDataFrame,
+)
+from replay.utils.session_handler import get_spark_session
 
 from .schema import FeatureHint, FeatureInfo, FeatureSchema, FeatureSource, FeatureType
 
@@ -47,9 +58,7 @@ class Dataset:
         self._query_features = query_features
         self._item_features = item_features
 
-        self.is_pandas = isinstance(interactions, PandasDataFrame)
-        self.is_spark = isinstance(interactions, SparkDataFrame)
-        self.is_polars = isinstance(interactions, PolarsDataFrame)
+        self._assign_df_type()
 
         self._categorical_encoded = categorical_encoded
 
@@ -74,16 +83,8 @@ class Dataset:
             msg = "Interactions and query features should have the same type."
             raise TypeError(msg)
 
-        self._feature_source_map: Dict[FeatureSource, DataFrameLike] = {
-            FeatureSource.INTERACTIONS: self.interactions,
-            FeatureSource.QUERY_FEATURES: self.query_features,
-            FeatureSource.ITEM_FEATURES: self.item_features,
-        }
-
-        self._ids_feature_map: Dict[FeatureHint, DataFrameLike] = {
-            FeatureHint.QUERY_ID: self.query_features if self.query_features is not None else self.interactions,
-            FeatureHint.ITEM_ID: self.item_features if self.item_features is not None else self.interactions,
-        }
+        self._get_feature_source_map()
+        self._get_ids_source_map()
 
         self._feature_schema = self._fill_feature_schema(feature_schema)
 
@@ -92,7 +93,6 @@ class Dataset:
                 self._check_ids_consistency(hint=FeatureHint.QUERY_ID)
             if self.item_features is not None:
                 self._check_ids_consistency(hint=FeatureHint.ITEM_ID)
-
             if self._categorical_encoded:
                 self._check_encoded()
 
@@ -189,6 +189,157 @@ class Dataset:
         """
         return self._feature_schema
 
+    def _get_df_type(self) -> str:
+        """
+        :returns: Stored dataframe type.
+        """
+        if self.is_spark:
+            return "spark"
+        if self.is_pandas:
+            return "pandas"
+        if self.is_polars:
+            return "polars"
+        msg = "No known dataframe types are provided"
+        raise ValueError(msg)
+
+    def _to_parquet(self, df: DataFrameLike, path: Path) -> None:
+        """
+        Save the content of the dataframe in parquet format to the provided path.
+
+        :param df: Dataframe to save.
+        :param path: Path to save the dataframe to.
+        """
+        if self.is_spark:
+            path = str(path)
+            df = df.withColumn("idx", sf.monotonically_increasing_id())
+            df.write.mode("overwrite").parquet(path)
+        elif self.is_pandas:
+            df.to_parquet(path)
+        elif self.is_polars:
+            df.write_parquet(path)
+        else:
+            msg = """
+            _to_parquet() can only be used to save polars|pandas|spark dataframes;
+            No known dataframe types are provided
+            """
+            raise TypeError(msg)
+
+    @staticmethod
+    def _read_parquet(path: Path, mode: str) -> Union[SparkDataFrame, PandasDataFrame, PolarsDataFrame]:
+        """
+        Read the parquet file as dataframe.
+
+        :param path: The parquet file path.
+        :param mode: Dataframe type. Can be spark|pandas|polars.
+        :returns: The dataframe read from the file.
+        """
+        if mode == "spark":
+            path = str(path)
+            spark_session = get_spark_session()
+            df = spark_session.read.parquet(path)
+            if "idx" in df.columns:
+                df = df.orderBy("idx").drop("idx")
+            return df
+        if mode == "pandas":
+            df = pd_read_parquet(path)
+            if "idx" in df.columns:
+                df = df.set_index("idx").reset_index(drop=True)
+            return df
+        if mode == "polars":
+            df = pl_read_parquet(path, use_pyarrow=True)
+            if "idx" in df.columns:
+                df = df.sort("idx").drop("idx")
+            return df
+        msg = f"_read_parquet() can only be used to read polars|pandas|spark dataframes, not {mode}"
+        raise TypeError(msg)
+
+    def save(self, path: str) -> None:
+        """
+        Save the Dataset to the provided path.
+
+        :param path: Path to save the Dataset to.
+        """
+        dataset_dict = {}
+        dataset_dict["_class_name"] = self.__class__.__name__
+
+        interactions_type = self._get_df_type()
+        dataset_dict["init_args"] = {
+            "feature_schema": [],
+            "interactions": interactions_type,
+            "item_features": (interactions_type if self.item_features is not None else None),
+            "query_features": (interactions_type if self.query_features is not None else None),
+            "check_consistency": False,
+            "categorical_encoded": self._categorical_encoded,
+        }
+
+        for feature in self.feature_schema.all_features:
+            dataset_dict["init_args"]["feature_schema"].append(
+                {
+                    "column": feature.column,
+                    "feature_type": feature.feature_type.name,
+                    "feature_hint": (feature.feature_hint.name if feature.feature_hint else None),
+                }
+            )
+
+        base_path = Path(path).with_suffix(".replay").resolve()
+        base_path.mkdir(parents=True, exist_ok=True)
+
+        with open(base_path / "init_args.json", "w+") as file:
+            json.dump(dataset_dict, file)
+
+        df_data = {
+            "interactions": self.interactions,
+            "item_features": self.item_features,
+            "query_features": self.query_features,
+        }
+
+        for df_name, df in df_data.items():
+            if df is not None:
+                df_path = base_path / f"{df_name}.parquet"
+                self._to_parquet(df, df_path)
+
+    @classmethod
+    def load(
+        cls,
+        path: str,
+        dataframe_type: Optional[str] = None,
+    ) -> Dataset:
+        """
+        Load the Dataset from the provided path.
+
+        :param path: The file path
+        :dataframe_type: Dataframe type to use to store internal data.
+            Can be spark|pandas|polars|None.
+            If not provided automatically sets to the one used when the Dataset was saved.
+        :returns: Loaded Dataset.
+        """
+        base_path = Path(path).with_suffix(".replay").resolve()
+        with open(base_path / "init_args.json", "r") as file:
+            dataset_dict = json.loads(file.read())
+
+        if dataframe_type not in ["pandas", "spark", "polars", None]:
+            msg = f"Argument dataframe_type can be spark|pandas|polars|None, not {dataframe_type}"
+            raise ValueError(msg)
+
+        feature_schema_data = dataset_dict["init_args"]["feature_schema"]
+        features_list = []
+        for feature_data in feature_schema_data:
+            f_type = feature_data["feature_type"]
+            f_hint = feature_data["feature_hint"]
+            feature_data["feature_type"] = FeatureType[f_type] if f_type else None
+            feature_data["feature_hint"] = FeatureHint[f_hint] if f_hint else None
+            features_list.append(FeatureInfo(**feature_data))
+        dataset_dict["init_args"]["feature_schema"] = FeatureSchema(features_list)
+
+        for df_name in ["interactions", "query_features", "item_features"]:
+            df_type = dataset_dict["init_args"][df_name]
+            if df_type:
+                df_type = dataframe_type or df_type
+                load_path = base_path / f"{df_name}.parquet"
+                dataset_dict["init_args"][df_name] = cls._read_parquet(load_path, df_type)
+        dataset = cls(**dataset_dict["init_args"])
+        return dataset
+
     if PYSPARK_AVAILABLE:
 
         def persist(self, storage_level: StorageLevel = StorageLevel(True, True, False, True, 1)) -> None:
@@ -283,6 +434,24 @@ class Dataset:
             categorical_encoded=self._categorical_encoded,
         )
 
+    def _get_feature_source_map(self):
+        self._feature_source_map: Dict[FeatureSource, DataFrameLike] = {
+            FeatureSource.INTERACTIONS: self.interactions,
+            FeatureSource.QUERY_FEATURES: self.query_features,
+            FeatureSource.ITEM_FEATURES: self.item_features,
+        }
+
+    def _get_ids_source_map(self):
+        self._ids_feature_map: Dict[FeatureHint, DataFrameLike] = {
+            FeatureHint.QUERY_ID: self.query_features if self.query_features is not None else self.interactions,
+            FeatureHint.ITEM_ID: self.item_features if self.item_features is not None else self.interactions,
+        }
+
+    def _assign_df_type(self):
+        self.is_pandas = isinstance(self.interactions, PandasDataFrame)
+        self.is_spark = isinstance(self.interactions, SparkDataFrame)
+        self.is_polars = isinstance(self.interactions, PolarsDataFrame)
+
     def _get_cardinality(self, feature: FeatureInfo) -> Callable:
         def callback(column: str) -> int:
             if feature.feature_hint in [FeatureHint.ITEM_ID, FeatureHint.QUERY_ID]:
@@ -381,7 +550,11 @@ class Dataset:
             is_consistent = (
                 self.interactions.select(ids_column)
                 .distinct()
-                .join(features_df.select(ids_column).distinct(), on=[ids_column], how="leftanti")
+                .join(
+                    features_df.select(ids_column).distinct(),
+                    on=[ids_column],
+                    how="leftanti",
+                )
                 .count()
             ) == 0
         else:
@@ -389,7 +562,11 @@ class Dataset:
                 len(
                     self.interactions.select(ids_column)
                     .unique()
-                    .join(features_df.select(ids_column).unique(), on=ids_column, how="anti")
+                    .join(
+                        features_df.select(ids_column).unique(),
+                        on=ids_column,
+                        how="anti",
+                    )
                 )
                 == 0
             )
@@ -399,7 +576,11 @@ class Dataset:
             raise ValueError(msg)
 
     def _check_column_encoded(
-        self, data: DataFrameLike, column: str, source: FeatureSource, cardinality: Optional[int]
+        self,
+        data: DataFrameLike,
+        column: str,
+        source: FeatureSource,
+        cardinality: Optional[int],
     ) -> None:
         """
         Checks that IDs are encoded:
@@ -482,6 +663,51 @@ class Dataset:
                     feature.cardinality,
                 )
 
+    def to_pandas(self) -> None:
+        """
+        Convert internally stored dataframes to pandas.DataFrame.
+        """
+        from replay.utils.common import convert2pandas
+
+        self._interactions = convert2pandas(self._interactions)
+        if self._query_features is not None:
+            self._query_features = convert2pandas(self._query_features)
+        if self._item_features is not None:
+            self._item_features = convert2pandas(self.item_features)
+        self._get_feature_source_map()
+        self._get_ids_source_map()
+        self._assign_df_type()
+
+    def to_spark(self):
+        """
+        Convert internally stored dataframes to pyspark.sql.DataFrame.
+        """
+        from replay.utils.common import convert2spark
+
+        self._interactions = convert2spark(self._interactions)
+        if self._query_features is not None:
+            self._query_features = convert2spark(self._query_features)
+        if self._item_features is not None:
+            self._item_features = convert2spark(self._item_features)
+        self._get_feature_source_map()
+        self._get_ids_source_map()
+        self._assign_df_type()
+
+    def to_polars(self):
+        """
+        Convert internally stored dataframes to polars.DataFrame.
+        """
+        from replay.utils.common import convert2polars
+
+        self._interactions = convert2polars(self._interactions)
+        if self._query_features is not None:
+            self._query_features = convert2polars(self._query_features)
+        if self._item_features is not None:
+            self._item_features = convert2polars(self._item_features)
+        self._get_feature_source_map()
+        self._get_ids_source_map()
+        self._assign_df_type()
+
 
 def nunique(data: DataFrameLike, column: str) -> int:
     """

{replay_rec-0.17.0rc0 → replay_rec-0.17.1}/replay/data/nn/schema.py

@@ -408,6 +408,48 @@ class TensorSchema(Mapping[str, TensorFeatureInfo]):
             return None
         return rating_features.item().name
 
+    def _get_object_args(self) -> Dict:
+        """
+        Returns list of features represented as dictionaries.
+        """
+        features = [
+            {
+                "name": feature.name,
+                "feature_type": feature.feature_type.name,
+                "is_seq": feature.is_seq,
+                "feature_hint": feature.feature_hint.name if feature.feature_hint else None,
+                "feature_sources": [
+                    {"source": x.source.name, "column": x.column, "index": x.index} for x in feature.feature_sources
+                ]
+                if feature.feature_sources
+                else None,
+                "cardinality": feature.cardinality if feature.feature_type == FeatureType.CATEGORICAL else None,
+                "embedding_dim": feature.embedding_dim if feature.feature_type == FeatureType.CATEGORICAL else None,
+                "tensor_dim": feature.tensor_dim if feature.feature_type == FeatureType.NUMERICAL else None,
+            }
+            for feature in self.all_features
+        ]
+        return features
+
+    @classmethod
+    def _create_object_by_args(cls, args: Dict) -> "TensorSchema":
+        features_list = []
+        for feature_data in args:
+            feature_data["feature_sources"] = (
+                [
+                    TensorFeatureSource(source=FeatureSource[x["source"]], column=x["column"], index=x["index"])
+                    for x in feature_data["feature_sources"]
+                ]
+                if feature_data["feature_sources"]
+                else None
+            )
+            f_type = feature_data["feature_type"]
+            f_hint = feature_data["feature_hint"]
+            feature_data["feature_type"] = FeatureType[f_type] if f_type else None
+            feature_data["feature_hint"] = FeatureHint[f_hint] if f_hint else None
+            features_list.append(TensorFeatureInfo(**feature_data))
+        return TensorSchema(features_list)
+
     def filter(
         self,
         name: Optional[str] = None,

{replay_rec-0.17.0rc0 → replay_rec-0.17.1}/replay/data/nn/sequence_tokenizer.py

@@ -24,7 +24,10 @@ SequenceDataFrameLike = Union[PandasDataFrame, PolarsDataFrame]
 
 class SequenceTokenizer:
     """
-    Data tokenizer for transformers
+    Data tokenizer for transformers;
+    Encodes all categorical features (the ones marked as FeatureType.CATEGORICAL in
+    the FeatureSchema) and stores all data as items sequences (sorted by time if a
+    feature of type FeatureHint.TIMESTAMP is provided, unsorted otherwise).
     """
 
     def __init__(
@@ -278,17 +281,17 @@ class SequenceTokenizer:
         ]
 
         for tensor_feature in tensor_schema.values():
-            source = tensor_feature.feature_source
-            assert source is not None
+            for source in tensor_feature.feature_sources:
+                assert source is not None
 
-            # Some columns already added to encoder, skip them
-            if source.column in features_subset:
-                continue
+                # Some columns already added to encoder, skip them
+                if source.column in features_subset:
+                    continue
 
-            if isinstance(source.source, FeatureSource):
-                features_subset.append(source.column)
-            else:
-                assert False, "Unknown tensor feature source"
+                if isinstance(source.source, FeatureSource):
+                    features_subset.append(source.column)
+                else:
+                    assert False, "Unknown tensor feature source"
 
         return set(features_subset)
 
@@ -404,7 +407,7 @@ class SequenceTokenizer:
 
     @classmethod
     @deprecation_warning("with `use_pickle` equals to `True` will be deprecated in future versions")
-    def load(cls, path: str, use_pickle: bool = False) -> "SequenceTokenizer":
+    def load(cls, path: str, use_pickle: bool = False, **kwargs) -> "SequenceTokenizer":
         """
         Load tokenizer object from the given path.
 
@@ -422,18 +425,7 @@ class SequenceTokenizer:
 
             # load tensor_schema, tensor_features
             tensor_schema_data = tokenizer_dict["init_args"]["tensor_schema"]
-            features_list = []
-            for feature_data in tensor_schema_data:
-                feature_data["feature_sources"] = [
-                    TensorFeatureSource(source=FeatureSource[x["source"]], column=x["column"], index=x["index"])
-                    for x in feature_data["feature_sources"]
-                ]
-                f_type = feature_data["feature_type"]
-                f_hint = feature_data["feature_hint"]
-                feature_data["feature_type"] = FeatureType[f_type] if f_type else None
-                feature_data["feature_hint"] = FeatureHint[f_hint] if f_hint else None
-                features_list.append(TensorFeatureInfo(**feature_data))
-            tokenizer_dict["init_args"]["tensor_schema"] = TensorSchema(features_list)
+            tokenizer_dict["init_args"]["tensor_schema"] = TensorSchema._create_object_by_args(tensor_schema_data)
 
             # Load encoder columns and rules
             types = list(FeatureHint) + list(FeatureSource)
@@ -447,7 +439,7 @@ class SequenceTokenizer:
                 rule_data = rules_dict[rule]
                 if rule_data["mapping"] and rule_data["is_int"]:
                     rule_data["mapping"] = {int(key): value for key, value in rule_data["mapping"].items()}
-                    del rule_data["is_int"]
+                del rule_data["is_int"]
 
                 tokenizer_dict["encoder"]["encoding_rules"][rule] = LabelEncodingRule(**rule_data)
 
@@ -478,31 +470,9 @@ class SequenceTokenizer:
                 "allow_collect_to_master": self._allow_collect_to_master,
                 "handle_unknown_rule": self._encoder._handle_unknown_rule,
                 "default_value_rule": self._encoder._default_value_rule,
-                "tensor_schema": [],
+                "tensor_schema": self._tensor_schema._get_object_args(),
             }
 
-            # save tensor schema
-            for feature in list(self._tensor_schema.values()):
-                tokenizer_dict["init_args"]["tensor_schema"].append(
-                    {
-                        "name": feature.name,
-                        "feature_type": feature.feature_type.name,
-                        "is_seq": feature.is_seq,
-                        "feature_hint": feature.feature_hint.name if feature.feature_hint else None,
-                        "feature_sources": [
-                            {"source": x.source.name, "column": x.column, "index": x.index}
-                            for x in feature.feature_sources
-                        ]
-                        if feature.feature_sources
-                        else None,
-                        "cardinality": feature.cardinality if feature.feature_type == FeatureType.CATEGORICAL else None,
-                        "embedding_dim": feature.embedding_dim
-                        if feature.feature_type == FeatureType.CATEGORICAL
-                        else None,
-                        "tensor_dim": feature.tensor_dim if feature.feature_type == FeatureType.NUMERICAL else None,
-                    }
-                )
-
             # save DatasetLabelEncoder
             tokenizer_dict["encoder"] = {
                 "features_columns": {key.name: value for key, value in self._encoder._features_columns.items()},