replay-rec 0.19.0rc0__py3-none-any.whl → 0.20.0rc0__py3-none-any.whl

replay/__init__.py

@@ -1,3 +1,7 @@
-""" RecSys library """
+"""RecSys library"""
 
-__version__ = "0.19.0.preview"
+# NOTE: This ensures distutils monkey-patching is performed before any
+# functionality removed in Python 3.12 is used in downstream packages (like lightfm)
+import setuptools as _
+
+__version__ = "0.20.0.preview"

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)
 
@@ -610,9 +611,9 @@ class Dataset:
             if self.is_pandas:
                 try:
                     data[column] = data[column].astype(int)
-                except Exception:
+                except Exception as exc:
                     msg = f"IDs in {source.name}.{column} are not encoded. They are not int."
-                    raise ValueError(msg)
+                    raise ValueError(msg) from exc
 
         if self.is_pandas:
             is_int = np.issubdtype(dict(data.dtypes)[column], int)
@@ -775,10 +776,10 @@ def check_dataframes_types_equal(dataframe: DataFrameLike, other: DataFrameLike)
 
     :returns: True if dataframes have same type.
     """
-    if isinstance(dataframe, PandasDataFrame) and isinstance(other, PandasDataFrame):
-        return True
-    if isinstance(dataframe, SparkDataFrame) and isinstance(other, SparkDataFrame):
-        return True
-    if isinstance(dataframe, PolarsDataFrame) and isinstance(other, PolarsDataFrame):
-        return True
-    return False
+    return any(
+        [
+            isinstance(dataframe, PandasDataFrame) and isinstance(other, PandasDataFrame),
+            isinstance(dataframe, SparkDataFrame) and isinstance(other, SparkDataFrame),
+            isinstance(dataframe, PolarsDataFrame) and isinstance(other, PolarsDataFrame),
+        ]
+    )

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/__init__.py

@@ -14,17 +14,17 @@ if TORCH_AVAILABLE:
     )
 
     __all__ = [
+        "DEFAULT_GROUND_TRUTH_PADDING_VALUE",
+        "DEFAULT_TRAIN_PADDING_VALUE",
         "MutableTensorMap",
+        "PandasSequentialDataset",
+        "PolarsSequentialDataset",
+        "SequenceTokenizer",
+        "SequentialDataset",
         "TensorFeatureInfo",
         "TensorFeatureSource",
         "TensorMap",
         "TensorSchema",
-        "SequenceTokenizer",
-        "PandasSequentialDataset",
-        "PolarsSequentialDataset",
-        "SequentialDataset",
-        "DEFAULT_GROUND_TRUTH_PADDING_VALUE",
-        "DEFAULT_TRAIN_PADDING_VALUE",
         "TorchSequentialBatch",
         "TorchSequentialDataset",
         "TorchSequentialValidationBatch",

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 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,11 +15,11 @@ 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.model_handler import deprecation_warning
+from replay.utils import deprecation_warning
 
-from .schema import TensorFeatureInfo, TensorFeatureSource, TensorSchema
-from .sequential_dataset import PandasSequentialDataset, PolarsSequentialDataset, SequentialDataset
-from .utils import ensure_pandas, groupby_sequences
+if TYPE_CHECKING:
+    from .schema import TensorFeatureInfo, TensorFeatureSource, TensorSchema
+    from .sequential_dataset import SequentialDataset
 
 SequenceDataFrameLike = Union[PandasDataFrame, PolarsDataFrame]
 _T = TypeVar("_T")
@@ -34,7 +35,7 @@ class SequenceTokenizer:
 
     def __init__(
         self,
-        tensor_schema: TensorSchema,
+        tensor_schema: "TensorSchema",
         handle_unknown_rule: HandleUnknownStrategies = "error",
         default_value_rule: Optional[Union[int, str]] = None,
         allow_collect_to_master: bool = False,
@@ -77,7 +78,7 @@ class SequenceTokenizer:
         self,
         dataset: Dataset,
         tensor_features_to_keep: Optional[Sequence[str]] = None,
-    ) -> SequentialDataset:
+    ) -> "SequentialDataset":
         """
         :param dataset: input dataset to transform
         :param tensor_features_to_keep: specified feature names to transform
@@ -89,7 +90,7 @@ class SequenceTokenizer:
     def fit_transform(
         self,
         dataset: Dataset,
-    ) -> SequentialDataset:
+    ) -> "SequentialDataset":
         """
         :param dataset: input dataset to transform
         :returns: SequentialDataset
@@ -97,7 +98,7 @@ class SequenceTokenizer:
         return self.fit(dataset)._transform_unchecked(dataset)
 
     @property
-    def tensor_schema(self) -> TensorSchema:
+    def tensor_schema(self) -> "TensorSchema":
         """
         :returns: tensor schema
         """
@@ -149,7 +150,9 @@ class SequenceTokenizer:
         self,
         dataset: Dataset,
         tensor_features_to_keep: Optional[Sequence[str]] = None,
-    ) -> SequentialDataset:
+    ) -> "SequentialDataset":
+        from replay.data.nn.sequential_dataset import PandasSequentialDataset, PolarsSequentialDataset
+
         schema = self._tensor_schema
         if tensor_features_to_keep is not None:
             schema = schema.subset(tensor_features_to_keep)
@@ -185,7 +188,9 @@ 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(
             events=dataset.interactions,
             groupby_col=dataset.feature_schema.query_id_column,
@@ -218,7 +223,7 @@ class SequenceTokenizer:
 
     def _make_sequence_features(
         self,
-        schema: TensorSchema,
+        schema: "TensorSchema",
         feature_schema: FeatureSchema,
         grouped_interactions: SequenceDataFrameLike,
         query_features: Optional[SequenceDataFrameLike],
@@ -242,7 +247,7 @@ class SequenceTokenizer:
     def _match_features_with_tensor_schema(
         cls,
         dataset: Dataset,
-        tensor_schema: TensorSchema,
+        tensor_schema: "TensorSchema",
     ) -> Dataset:
         feature_subset_filter = cls._get_features_filter_from_schema(
             tensor_schema,
@@ -261,16 +266,16 @@ class SequenceTokenizer:
     @classmethod
     def _get_features_filter_from_schema(
         cls,
-        tensor_schema: TensorSchema,
+        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,
         ]
@@ -291,7 +296,7 @@ class SequenceTokenizer:
         return set(features_subset)
 
     @classmethod
-    def _check_tensor_schema(cls, tensor_schema: TensorSchema) -> None:
+    def _check_tensor_schema(cls, tensor_schema: "TensorSchema") -> None:
         # Check consistency of sequential features
         for tensor_feature in tensor_schema.all_features:
             feature_sources = tensor_feature.feature_sources
@@ -299,7 +304,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:
@@ -319,11 +324,11 @@ class SequenceTokenizer:
     def _check_if_tensor_schema_matches_data(  # noqa: C901
         cls,
         dataset: Dataset,
-        tensor_schema: TensorSchema,
+        tensor_schema: "TensorSchema",
         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
@@ -413,9 +418,11 @@ class SequenceTokenizer:
 
         :returns: Loaded tokenizer object.
         """
+        from replay.data.nn import TensorSchema
+
         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
@@ -500,7 +507,7 @@ class _BaseSequenceProcessor(Generic[_T]):
 
     def __init__(
         self,
-        tensor_schema: TensorSchema,
+        tensor_schema: "TensorSchema",
         query_id_column: str,
         item_id_column: str,
         grouped_interactions: _T,
@@ -535,7 +542,7 @@ class _BaseSequenceProcessor(Generic[_T]):
             return self._process_num_feature(tensor_feature)
         assert False, "Unknown tensor feature type"
 
-    def _process_num_feature(self, tensor_feature: TensorFeatureInfo) -> _T:
+    def _process_num_feature(self, tensor_feature: "TensorFeatureInfo") -> _T:
         """
         Process numerical tensor feature depends on it source.
         """
@@ -548,7 +555,7 @@ class _BaseSequenceProcessor(Generic[_T]):
             return self._process_num_item_feature(tensor_feature)
         assert False, "Unknown tensor feature source table"
 
-    def _process_cat_feature(self, tensor_feature: TensorFeatureInfo) -> _T:
+    def _process_cat_feature(self, tensor_feature: "TensorFeatureInfo") -> _T:
         """
         Process categorical tensor feature depends on it source.
         """
@@ -562,27 +569,27 @@ class _BaseSequenceProcessor(Generic[_T]):
         assert False, "Unknown tensor feature source table"
 
     @abc.abstractmethod
-    def _process_cat_interaction_feature(self, tensor_feature: TensorFeatureInfo) -> _T:  # pragma: no cover
+    def _process_cat_interaction_feature(self, tensor_feature: "TensorFeatureInfo") -> _T:  # pragma: no cover
         pass
 
     @abc.abstractmethod
-    def _process_cat_query_feature(self, tensor_feature: TensorFeatureInfo) -> _T:  # pragma: no cover
+    def _process_cat_query_feature(self, tensor_feature: "TensorFeatureInfo") -> _T:  # pragma: no cover
         pass
 
     @abc.abstractmethod
-    def _process_cat_item_feature(self, tensor_feature: TensorFeatureInfo) -> _T:  # pragma: no cover
+    def _process_cat_item_feature(self, tensor_feature: "TensorFeatureInfo") -> _T:  # pragma: no cover
         pass
 
     @abc.abstractmethod
-    def _process_num_interaction_feature(self, tensor_feature: TensorFeatureInfo) -> _T:  # pragma: no cover
+    def _process_num_interaction_feature(self, tensor_feature: "TensorFeatureInfo") -> _T:  # pragma: no cover
         pass
 
     @abc.abstractmethod
-    def _process_num_query_feature(self, tensor_feature: TensorFeatureInfo) -> _T:  # pragma: no cover
+    def _process_num_query_feature(self, tensor_feature: "TensorFeatureInfo") -> _T:  # pragma: no cover
         pass
 
     @abc.abstractmethod
-    def _process_num_item_feature(self, tensor_feature: TensorFeatureInfo) -> _T:  # pragma: no cover
+    def _process_num_item_feature(self, tensor_feature: "TensorFeatureInfo") -> _T:  # pragma: no cover
         pass
 
 
@@ -597,7 +604,7 @@ class _PandasSequenceProcessor(_BaseSequenceProcessor[PandasDataFrame]):
 
     def __init__(
         self,
-        tensor_schema: TensorSchema,
+        tensor_schema: "TensorSchema",
         query_id_column: str,
         item_id_column: str,
         grouped_interactions: PandasDataFrame,
@@ -619,7 +626,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:
@@ -628,8 +635,8 @@ class _PandasSequenceProcessor(_BaseSequenceProcessor[PandasDataFrame]):
         return PandasDataFrame(all_features)
 
     def _process_num_interaction_feature(
-        self, tensor_feature: TensorFeatureInfo
-    ) -> Union[List[np.ndarray], List[List]]:
+        self, tensor_feature: "TensorFeatureInfo"
+    ) -> Union[list[np.ndarray], list[list]]:
         """
         Process numerical interaction feature.
 
@@ -650,7 +657,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.
 
@@ -676,7 +683,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.
 
@@ -687,8 +694,8 @@ class _PandasSequenceProcessor(_BaseSequenceProcessor[PandasDataFrame]):
         return self._process_cat_query_feature(tensor_feature)
 
     def _process_cat_interaction_feature(
-        self, tensor_feature: TensorFeatureInfo
-    ) -> Union[List[np.ndarray], List[List]]:
+        self, tensor_feature: "TensorFeatureInfo"
+    ) -> Union[list[np.ndarray], list[list]]:
         """
         Process categorical interaction feature.
 
@@ -709,7 +716,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.
 
@@ -738,7 +745,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.
 
@@ -754,7 +761,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
@@ -784,7 +791,7 @@ class _PolarsSequenceProcessor(_BaseSequenceProcessor[PolarsDataFrame]):
             data = data.join(self._process_feature(tensor_feature_name), on=self._query_id_column, how="left")
         return data
 
-    def _process_num_interaction_feature(self, tensor_feature: TensorFeatureInfo) -> PolarsDataFrame:
+    def _process_num_interaction_feature(self, tensor_feature: "TensorFeatureInfo") -> PolarsDataFrame:
         """
         Process numerical interaction feature.
 
@@ -794,7 +801,7 @@ class _PolarsSequenceProcessor(_BaseSequenceProcessor[PolarsDataFrame]):
         """
         return self._process_cat_interaction_feature(tensor_feature)
 
-    def _process_num_query_feature(self, tensor_feature: TensorFeatureInfo) -> PolarsDataFrame:
+    def _process_num_query_feature(self, tensor_feature: "TensorFeatureInfo") -> PolarsDataFrame:
         """
         Process numerical feature from query features dataset.
 
@@ -805,7 +812,7 @@ class _PolarsSequenceProcessor(_BaseSequenceProcessor[PolarsDataFrame]):
         """
         return self._process_cat_query_feature(tensor_feature)
 
-    def _process_num_item_feature(self, tensor_feature: TensorFeatureInfo) -> PolarsDataFrame:
+    def _process_num_item_feature(self, tensor_feature: "TensorFeatureInfo") -> PolarsDataFrame:
         """
         Process numerical feature from item features dataset.
 
@@ -816,7 +823,7 @@ class _PolarsSequenceProcessor(_BaseSequenceProcessor[PolarsDataFrame]):
         """
         return self._process_cat_item_feature(tensor_feature)
 
-    def _process_cat_interaction_feature(self, tensor_feature: TensorFeatureInfo) -> PolarsDataFrame:
+    def _process_cat_interaction_feature(self, tensor_feature: "TensorFeatureInfo") -> PolarsDataFrame:
         """
         Process categorical interaction feature.
 
@@ -833,7 +840,7 @@ class _PolarsSequenceProcessor(_BaseSequenceProcessor[PolarsDataFrame]):
             {source.column: tensor_feature.name}
         )
 
-    def _process_cat_query_feature(self, tensor_feature: TensorFeatureInfo) -> PolarsDataFrame:
+    def _process_cat_query_feature(self, tensor_feature: "TensorFeatureInfo") -> PolarsDataFrame:
         """
         Process categorical feature from query features dataset.
 
@@ -877,7 +884,7 @@ class _PolarsSequenceProcessor(_BaseSequenceProcessor[PolarsDataFrame]):
             {source.column: tensor_feature.name}
         )
 
-    def _process_cat_item_feature(self, tensor_feature: TensorFeatureInfo) -> PolarsDataFrame:
+    def _process_cat_item_feature(self, tensor_feature: "TensorFeatureInfo") -> PolarsDataFrame:
         """
         Process categorical feature from item features dataset.