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

replay/__init__.py

@@ -4,4 +4,4 @@
 # functionality removed in Python 3.12 is used in downstream packages (like lightfm)
 import setuptools as _
 
-__version__ = "0.20.0"
+__version__ = "0.20.1"

replay/data/dataset.py

@@ -5,8 +5,9 @@
 from __future__ import annotations
 
 import json
+from collections.abc import Iterable, Sequence
 from pathlib import Path
-from typing import Callable, Dict, Iterable, List, Optional, Sequence, Union
+from typing import Callable, Optional, Union
 
 import numpy as np
 from pandas import read_parquet as pd_read_parquet
@@ -315,7 +316,7 @@ class Dataset:
         :returns: Loaded Dataset.
         """
         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:
             dataset_dict = json.loads(file.read())
 
         if dataframe_type not in ["pandas", "spark", "polars", None]:
@@ -436,14 +437,14 @@ class Dataset:
         )
 
     def _get_feature_source_map(self):
-        self._feature_source_map: Dict[FeatureSource, DataFrameLike] = {
+        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] = {
+        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,
         }
@@ -499,10 +500,10 @@ class Dataset:
             )
         return FeatureSchema(features_list=features_list + filled_features)
 
-    def _fill_unlabeled_features_sources(self, feature_schema: FeatureSchema) -> List[FeatureInfo]:
+    def _fill_unlabeled_features_sources(self, feature_schema: FeatureSchema) -> list[FeatureInfo]:
         features_list = list(feature_schema.all_features)
 
-        source_mapping: Dict[str, FeatureSource] = {}
+        source_mapping: dict[str, FeatureSource] = {}
         for source in FeatureSource:
             dataframe = self._feature_source_map[source]
             if dataframe is not None:
@@ -524,7 +525,7 @@ class Dataset:
         self._set_cardinality(features_list=features_list)
         return features_list
 
-    def _get_unlabeled_columns(self, source: FeatureSource, feature_schema: FeatureSchema) -> List[FeatureInfo]:
+    def _get_unlabeled_columns(self, source: FeatureSource, feature_schema: FeatureSchema) -> list[FeatureInfo]:
         set_source_dataframe_columns = set(self._feature_source_map[source].columns)
         set_labeled_dataframe_columns = set(feature_schema.columns)
         unlabeled_columns = set_source_dataframe_columns - set_labeled_dataframe_columns
@@ -534,13 +535,13 @@ class Dataset:
         ]
         return unlabeled_features_list
 
-    def _fill_unlabeled_features(self, source: FeatureSource, feature_schema: FeatureSchema) -> List[FeatureInfo]:
+    def _fill_unlabeled_features(self, source: FeatureSource, feature_schema: FeatureSchema) -> list[FeatureInfo]:
         unlabeled_columns = self._get_unlabeled_columns(source=source, feature_schema=feature_schema)
         self._set_features_source(feature_list=unlabeled_columns, source=source)
         self._set_cardinality(features_list=unlabeled_columns)
         return unlabeled_columns
 
-    def _set_features_source(self, feature_list: List[FeatureInfo], source: FeatureSource) -> None:
+    def _set_features_source(self, feature_list: list[FeatureInfo], source: FeatureSource) -> None:
         for feature in feature_list:
             feature._set_feature_source(source)
 

replay/data/dataset_utils/dataset_label_encoder.py

@@ -6,7 +6,8 @@ Contains classes for encoding categorical data
 """
 
 import warnings
-from typing import Dict, Iterable, Iterator, Optional, Sequence, Set, Union
+from collections.abc import Iterable, Iterator, Sequence
+from typing import Optional, Union
 
 from replay.data import Dataset, FeatureHint, FeatureSchema, FeatureSource, FeatureType
 from replay.preprocessing import LabelEncoder, LabelEncodingRule, SequenceEncodingRule
@@ -45,9 +46,9 @@ class DatasetLabelEncoder:
         """
         self._handle_unknown_rule = handle_unknown_rule
         self._default_value_rule = default_value_rule
-        self._encoding_rules: Dict[str, LabelEncodingRule] = {}
+        self._encoding_rules: dict[str, LabelEncodingRule] = {}
 
-        self._features_columns: Dict[Union[FeatureHint, FeatureSource], Sequence[str]] = {}
+        self._features_columns: dict[Union[FeatureHint, FeatureSource], Sequence[str]] = {}
 
     def fit(self, dataset: Dataset) -> "DatasetLabelEncoder":
         """
@@ -161,7 +162,7 @@ class DatasetLabelEncoder:
         """
         self._check_if_initialized()
 
-        columns_set: Set[str]
+        columns_set: set[str]
         columns_set = {columns} if isinstance(columns, str) else {*columns}
 
         def get_encoding_rules() -> Iterator[LabelEncodingRule]:

replay/data/nn/schema.py

@@ -1,17 +1,8 @@
+from collections import OrderedDict
+from collections.abc import ItemsView, Iterable, Iterator, KeysView, Mapping, Sequence, ValuesView
 from typing import (
-    Dict,
-    ItemsView,
-    Iterable,
-    Iterator,
-    KeysView,
-    List,
-    Mapping,
     Optional,
-    OrderedDict,
-    Sequence,
-    Set,
     Union,
-    ValuesView,
 )
 
 import torch
@@ -20,7 +11,7 @@ from replay.data import FeatureHint, FeatureSource, FeatureType
 
 # Alias
 TensorMap = Mapping[str, torch.Tensor]
-MutableTensorMap = Dict[str, torch.Tensor]
+MutableTensorMap = dict[str, torch.Tensor]
 
 
 class TensorFeatureSource:
@@ -79,7 +70,7 @@ class TensorFeatureInfo:
         feature_type: FeatureType,
         is_seq: bool = False,
         feature_hint: Optional[FeatureHint] = None,
-        feature_sources: Optional[List[TensorFeatureSource]] = None,
+        feature_sources: Optional[list[TensorFeatureSource]] = None,
         cardinality: Optional[int] = None,
         padding_value: int = 0,
         embedding_dim: Optional[int] = None,
@@ -154,13 +145,13 @@ class TensorFeatureInfo:
         self._feature_hint = hint
 
     @property
-    def feature_sources(self) -> Optional[List[TensorFeatureSource]]:
+    def feature_sources(self) -> Optional[list[TensorFeatureSource]]:
         """
         :returns: List of sources feature came from.
         """
         return self._feature_sources
 
-    def _set_feature_sources(self, sources: List[TensorFeatureSource]) -> None:
+    def _set_feature_sources(self, sources: list[TensorFeatureSource]) -> None:
         self._feature_sources = sources
 
     @property
@@ -276,7 +267,7 @@ class TensorSchema(Mapping[str, TensorFeatureInfo]):
 
         :returns: New tensor schema of given features.
         """
-        features: Set[TensorFeatureInfo] = set()
+        features: set[TensorFeatureInfo] = set()
         for feature_name in features_to_keep:
             features.add(self._tensor_schema[feature_name])
         return TensorSchema(list(features))
@@ -432,7 +423,7 @@ class TensorSchema(Mapping[str, TensorFeatureInfo]):
             return None
         return rating_features.item().name
 
-    def _get_object_args(self) -> Dict:
+    def _get_object_args(self) -> dict:
         """
         Returns list of features represented as dictionaries.
         """
@@ -456,7 +447,7 @@ class TensorSchema(Mapping[str, TensorFeatureInfo]):
         return features
 
     @classmethod
-    def _create_object_by_args(cls, args: Dict) -> "TensorSchema":
+    def _create_object_by_args(cls, args: dict) -> "TensorSchema":
         features_list = []
         for feature_data in args:
             feature_data["feature_sources"] = (

replay/data/nn/sequence_tokenizer.py

@@ -2,8 +2,9 @@ import abc
 import json
 import pickle
 import warnings
+from collections.abc import Sequence
 from pathlib import Path
-from typing import TYPE_CHECKING, Dict, Generic, List, Optional, Sequence, Set, Tuple, TypeVar, Union
+from typing import TYPE_CHECKING, Generic, Optional, TypeVar, Union
 
 import numpy as np
 import polars as pl
@@ -14,7 +15,6 @@ from replay.data import Dataset, FeatureHint, FeatureSchema, FeatureSource, Feat
 from replay.data.dataset_utils import DatasetLabelEncoder
 from replay.preprocessing import LabelEncoder, LabelEncodingRule
 from replay.preprocessing.label_encoder import HandleUnknownStrategies
-from replay.utils import deprecation_warning
 
 if TYPE_CHECKING:
     from .schema import TensorFeatureInfo, TensorFeatureSource, TensorSchema
@@ -187,7 +187,7 @@ class SequenceTokenizer:
     def _group_dataset(
         self,
         dataset: Dataset,
-    ) -> Tuple[SequenceDataFrameLike, Optional[SequenceDataFrameLike], Optional[SequenceDataFrameLike]]:
+    ) -> tuple[SequenceDataFrameLike, Optional[SequenceDataFrameLike], Optional[SequenceDataFrameLike]]:
         from replay.data.nn.utils import ensure_pandas, groupby_sequences
 
         grouped_interactions = groupby_sequences(
@@ -268,13 +268,13 @@ class SequenceTokenizer:
         tensor_schema: "TensorSchema",
         query_id_column: str,
         item_id_column: str,
-    ) -> Set[str]:
+    ) -> set[str]:
         # We need only features, which related to tensor schema, otherwise feature should
         # be ignored for efficiency reasons. The code below does feature filtering, and
         # keeps features used as a source in tensor schema.
 
         # Query and item IDs are always needed
-        features_subset: List[str] = [
+        features_subset: list[str] = [
             query_id_column,
             item_id_column,
         ]
@@ -303,7 +303,7 @@ class SequenceTokenizer:
                 msg = "All tensor features must have sources defined"
                 raise ValueError(msg)
 
-            source_tables: List[FeatureSource] = [s.source for s in feature_sources]
+            source_tables: list[FeatureSource] = [s.source for s in feature_sources]
 
             unexpected_tables = list(filter(lambda x: not isinstance(x, FeatureSource), source_tables))
             if len(unexpected_tables) > 0:
@@ -327,7 +327,7 @@ class SequenceTokenizer:
         tensor_features_to_keep: Optional[Sequence[str]] = None,
     ) -> None:
         # Check if all source columns specified in tensor schema exist in provided data frames
-        sources_for_tensors: List["TensorFeatureSource"] = []
+        sources_for_tensors: list["TensorFeatureSource"] = []
         for tensor_feature_name, tensor_feature in tensor_schema.items():
             if tensor_features_to_keep is not None and tensor_feature_name not in tensor_features_to_keep:
                 continue
@@ -405,7 +405,6 @@ class SequenceTokenizer:
             tensor_feature._set_cardinality(dataset_feature.cardinality)
 
     @classmethod
-    @deprecation_warning("with `use_pickle` equals to `True` will be deprecated in future versions")
     def load(cls, path: str, use_pickle: bool = False, **kwargs) -> "SequenceTokenizer":
         """
         Load tokenizer object from the given path.
@@ -421,7 +420,7 @@ class SequenceTokenizer:
 
         if not use_pickle:
             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:
                 tokenizer_dict = json.loads(file.read())
 
             # load tensor_schema, tensor_features
@@ -449,12 +448,16 @@ class SequenceTokenizer:
             tokenizer._encoder._features_columns = encoder_features_columns
             tokenizer._encoder._encoding_rules = tokenizer_dict["encoder"]["encoding_rules"]
         else:
+            warnings.warn(
+                "with `use_pickle` equals to `True` will be deprecated in future versions",
+                DeprecationWarning,
+                stacklevel=2,
+            )
             with open(path, "rb") as file:
                 tokenizer = pickle.load(file)
 
         return tokenizer
 
-    @deprecation_warning("with `use_pickle` equals to `True` will be deprecated in future versions")
     def save(self, path: str, use_pickle: bool = False) -> None:
         """
         Save the tokenizer to the given path.
@@ -495,6 +498,11 @@ class SequenceTokenizer:
             with open(base_path / "init_args.json", "w+") as file:
                 json.dump(tokenizer_dict, file)
         else:
+            warnings.warn(
+                "with `use_pickle` equals to `True` will be deprecated in future versions",
+                DeprecationWarning,
+                stacklevel=2,
+            )
             with open(path, "wb") as file:
                 pickle.dump(self, file)
 
@@ -625,7 +633,7 @@ class _PandasSequenceProcessor(_BaseSequenceProcessor[PandasDataFrame]):
         """
         :returns: processed Pandas DataFrame with all features from tensor schema.
         """
-        all_features: Dict[str, Union[np.ndarray, List[np.ndarray]]] = {}
+        all_features: dict[str, Union[np.ndarray, list[np.ndarray]]] = {}
         all_features[self._query_id_column] = self._grouped_interactions[self._query_id_column].values
 
         for tensor_feature_name in self._tensor_schema:
@@ -635,7 +643,7 @@ class _PandasSequenceProcessor(_BaseSequenceProcessor[PandasDataFrame]):
 
     def _process_num_interaction_feature(
         self, tensor_feature: "TensorFeatureInfo"
-    ) -> Union[List[np.ndarray], List[List]]:
+    ) -> Union[list[np.ndarray], list[list]]:
         """
         Process numerical interaction feature.
 
@@ -656,7 +664,7 @@ class _PandasSequenceProcessor(_BaseSequenceProcessor[PandasDataFrame]):
                 values.append(np.array(sequence))
         return values
 
-    def _process_num_item_feature(self, tensor_feature: "TensorFeatureInfo") -> Union[List[np.ndarray], List[List]]:
+    def _process_num_item_feature(self, tensor_feature: "TensorFeatureInfo") -> Union[list[np.ndarray], list[list]]:
         """
         Process numerical feature from item features dataset.
 
@@ -682,7 +690,7 @@ class _PandasSequenceProcessor(_BaseSequenceProcessor[PandasDataFrame]):
 
         return values
 
-    def _process_num_query_feature(self, tensor_feature: "TensorFeatureInfo") -> List[np.ndarray]:
+    def _process_num_query_feature(self, tensor_feature: "TensorFeatureInfo") -> list[np.ndarray]:
         """
         Process numerical feature from query features dataset.
 
@@ -694,7 +702,7 @@ class _PandasSequenceProcessor(_BaseSequenceProcessor[PandasDataFrame]):
 
     def _process_cat_interaction_feature(
         self, tensor_feature: "TensorFeatureInfo"
-    ) -> Union[List[np.ndarray], List[List]]:
+    ) -> Union[list[np.ndarray], list[list]]:
         """
         Process categorical interaction feature.
 
@@ -715,7 +723,7 @@ class _PandasSequenceProcessor(_BaseSequenceProcessor[PandasDataFrame]):
                 values.append(np.array(sequence))
         return values
 
-    def _process_cat_query_feature(self, tensor_feature: "TensorFeatureInfo") -> List[np.ndarray]:
+    def _process_cat_query_feature(self, tensor_feature: "TensorFeatureInfo") -> list[np.ndarray]:
         """
         Process categorical feature from query features dataset.
 
@@ -744,7 +752,7 @@ class _PandasSequenceProcessor(_BaseSequenceProcessor[PandasDataFrame]):
                 ]
         return [np.array([query_feature[i]]).reshape(-1) for i in range(len(self._grouped_interactions))]
 
-    def _process_cat_item_feature(self, tensor_feature: "TensorFeatureInfo") -> Union[List[np.ndarray], List[List]]:
+    def _process_cat_item_feature(self, tensor_feature: "TensorFeatureInfo") -> Union[list[np.ndarray], list[list]]:
         """
         Process categorical feature from item features dataset.
 
@@ -760,7 +768,7 @@ class _PandasSequenceProcessor(_BaseSequenceProcessor[PandasDataFrame]):
         assert source is not None
 
         item_feature = self._item_features[source.column]
-        values: List[np.ndarray] = []
+        values: list[np.ndarray] = []
 
         for item_id_sequence in self._grouped_interactions[self._item_id_column]:
             feature_sequence = item_feature.loc[item_id_sequence].values

replay/data/nn/sequential_dataset.py

@@ -1,7 +1,7 @@
 import abc
 import json
 from pathlib import Path
-from typing import TYPE_CHECKING, Tuple, Union
+from typing import TYPE_CHECKING, Union
 
 import numpy as np
 import pandas as pd
@@ -90,7 +90,7 @@ class SequentialDataset(abc.ABC):
     @staticmethod
     def keep_common_query_ids(
         lhs: "SequentialDataset", rhs: "SequentialDataset"
-    ) -> Tuple["SequentialDataset", "SequentialDataset"]:
+    ) -> tuple["SequentialDataset", "SequentialDataset"]:
         """
         Returns `SequentialDatasets` that contain query ids from both datasets.
 
@@ -110,17 +110,27 @@ class SequentialDataset(abc.ABC):
 
         sequential_dict = {}
         sequential_dict["_class_name"] = self.__class__.__name__
-        self._sequences.reset_index().to_json(base_path / "sequences.json")
+
+        df = SequentialDataset._convert_array_to_list(self._sequences)
+        df.reset_index().to_parquet(base_path / "sequences.parquet")
         sequential_dict["init_args"] = {
             "tensor_schema": self._tensor_schema._get_object_args(),
             "query_id_column": self._query_id_column,
             "item_id_column": self._item_id_column,
-            "sequences_path": "sequences.json",
+            "sequences_path": "sequences.parquet",
         }
 
         with open(base_path / "init_args.json", "w+") as file:
             json.dump(sequential_dict, file)
 
+    @staticmethod
+    def _convert_array_to_list(df):
+        return df.map(lambda x: x.tolist() if isinstance(x, np.ndarray) else x)
+
+    @staticmethod
+    def _convert_list_to_array(df):
+        return df.map(lambda x: np.array(x) if isinstance(x, list) else x)
+
 
 class PandasSequentialDataset(SequentialDataset):
     """
@@ -149,7 +159,7 @@ class PandasSequentialDataset(SequentialDataset):
         if sequences.index.name != query_id_column:
             sequences = sequences.set_index(query_id_column)
 
-        self._sequences = sequences
+        self._sequences = SequentialDataset._convert_list_to_array(sequences)
 
     def __len__(self) -> int:
         return len(self._sequences)
@@ -203,10 +213,11 @@ class PandasSequentialDataset(SequentialDataset):
         from replay.data.nn import TensorSchema
 
         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:
             sequential_dict = json.loads(file.read())
 
-        sequences = pd.read_json(base_path / sequential_dict["init_args"]["sequences_path"])
+        sequences = pd.read_parquet(base_path / sequential_dict["init_args"]["sequences_path"])
+        sequences = cls._convert_array_to_list(sequences)
         dataset = cls(
             tensor_schema=TensorSchema._create_object_by_args(sequential_dict["init_args"]["tensor_schema"]),
             query_id_column=sequential_dict["init_args"]["query_id_column"],
@@ -258,18 +269,11 @@ class PolarsSequentialDataset(PandasSequentialDataset):
 
     def _convert_polars_to_pandas(self, df: PolarsDataFrame) -> PandasDataFrame:
         pandas_df = PandasDataFrame(df.to_dict(as_series=False))
-
-        for column in pandas_df.select_dtypes(include="object").columns:
-            if isinstance(pandas_df[column].iloc[0], list):
-                pandas_df[column] = pandas_df[column].apply(lambda x: np.array(x))
-
+        pandas_df = SequentialDataset._convert_list_to_array(pandas_df)
         return pandas_df
 
     def _convert_pandas_to_polars(self, df: PandasDataFrame) -> PolarsDataFrame:
-        for column in df.select_dtypes(include="object").columns:
-            if isinstance(df[column].iloc[0], np.ndarray):
-                df[column] = df[column].apply(lambda x: x.tolist())
-
+        df = SequentialDataset._convert_array_to_list(df)
         return pl.from_dict(df.to_dict("list"))
 
     @classmethod
@@ -287,10 +291,10 @@ class PolarsSequentialDataset(PandasSequentialDataset):
         from replay.data.nn import TensorSchema
 
         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:
             sequential_dict = json.loads(file.read())
 
-        sequences = pl.DataFrame(pd.read_json(base_path / sequential_dict["init_args"]["sequences_path"]))
+        sequences = pl.from_pandas(pd.read_parquet(base_path / sequential_dict["init_args"]["sequences_path"]))
         dataset = cls(
             tensor_schema=TensorSchema._create_object_by_args(sequential_dict["init_args"]["tensor_schema"]),
             query_id_column=sequential_dict["init_args"]["query_id_column"],

replay/data/nn/torch_sequential_dataset.py

@@ -1,11 +1,11 @@
-from typing import TYPE_CHECKING, Generator, NamedTuple, Optional, Sequence, Tuple, Union, cast
+import warnings
+from collections.abc import Generator, Sequence
+from typing import TYPE_CHECKING, NamedTuple, Optional, Union, cast
 
 import numpy as np
 import torch
 from torch.utils.data import Dataset as TorchDataset
 
-from replay.utils import deprecation_warning
-
 if TYPE_CHECKING:
     from .schema import TensorFeatureInfo, TensorMap, TensorSchema
     from .sequential_dataset import SequentialDataset
@@ -28,16 +28,12 @@ class TorchSequentialDataset(TorchDataset):
     Torch dataset for sequential recommender models
     """
 
-    @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,
         sliding_window_step: Optional[int] = None,
-        padding_value: int = 0,
+        padding_value: Optional[int] = None,
     ) -> None:
         """
         :param sequential: sequential dataset
@@ -52,6 +48,15 @@ class TorchSequentialDataset(TorchDataset):
         self._sequential = sequential
         self._max_sequence_length = max_sequence_length
         self._sliding_window_step = sliding_window_step
+        if padding_value is not None:
+            warnings.warn(
+                "`padding_value` parameter will be removed in future versions. "
+                "Instead, you should specify `padding_value` for each column in TensorSchema",
+                DeprecationWarning,
+                stacklevel=2,
+            )
+        else:
+            padding_value = 0
         self._padding_value = padding_value
         self._index2sequence_map = self._build_index2sequence_map()
 
@@ -110,7 +115,7 @@ class TorchSequentialDataset(TorchDataset):
             return sequence
 
         # form shape for padded_sequence. Now supported one and two-dimentions features
-        padded_sequence_shape: Union[Tuple[int, int], Tuple[int]]
+        padded_sequence_shape: Union[tuple[int, int], tuple[int]]
         if len(sequence.shape) == 1:
             padded_sequence_shape = (self._max_sequence_length,)
         elif len(sequence.shape) == 2:
@@ -134,10 +139,10 @@ class TorchSequentialDataset(TorchDataset):
             return torch.float32
         assert False, "Unknown tensor feature type"
 
-    def _build_index2sequence_map(self) -> Sequence[Tuple[int, int]]:
+    def _build_index2sequence_map(self) -> Sequence[tuple[int, int]]:
         return list(self._iter_with_window())
 
-    def _iter_with_window(self) -> Generator[Tuple[int, int], None, None]:
+    def _iter_with_window(self) -> Generator[tuple[int, int], None, None]:
         for i in range(len(self._sequential)):
             actual_seq_len = self._sequential.get_sequence_length(i)
             left_seq_len = actual_seq_len - self._max_sequence_length
@@ -176,17 +181,13 @@ class TorchSequentialValidationDataset(TorchDataset):
     Torch dataset for sequential recommender models that additionally stores ground truth
     """
 
-    @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,
         sliding_window_step: Optional[int] = None,
         label_feature_name: Optional[str] = None,
     ):

replay/data/nn/utils.py

@@ -1,4 +1,5 @@
-from typing import Iterable, Optional
+from collections.abc import Iterable
+from typing import Optional
 
 import polars as pl
 

replay/data/schema.py

@@ -1,18 +1,9 @@
+from collections.abc import ItemsView, Iterable, Iterator, KeysView, Mapping, Sequence, ValuesView
 from enum import Enum
 from typing import (
     Callable,
-    Dict,
-    ItemsView,
-    Iterable,
-    Iterator,
-    KeysView,
-    List,
-    Mapping,
     Optional,
-    Sequence,
-    Set,
     Union,
-    ValuesView,
 )
 
 
@@ -162,7 +153,7 @@ class FeatureSchema(Mapping[str, FeatureInfo]):
             in original schema to keep in subset.
         :returns: new feature schema of given features.
         """
-        features: Set[FeatureInfo] = set()
+        features: set[FeatureInfo] = set()
         for feature_column in features_to_keep:
             if feature_column in self._features_schema:
                 features.add(self._features_schema[feature_column])
@@ -438,7 +429,7 @@ class FeatureSchema(Mapping[str, FeatureInfo]):
         """
         unique_columns = set()
         duplicates = set()
-        item_query_names: Dict[FeatureHint, List[str]] = {
+        item_query_names: dict[FeatureHint, list[str]] = {
             FeatureHint.ITEM_ID: [],
             FeatureHint.QUERY_ID: [],
         }