replay-rec 0.20.3__py3-none-any.whl → 0.21.0__py3-none-any.whl

replay/data/nn/parquet/partitioned_iterable_dataset.py

@@ -0,0 +1,56 @@
+from collections.abc import Iterable, Iterator
+from typing import Optional
+
+import torch
+import torch.utils.data as data
+
+from replay.data.nn.parquet import DEFAULT_REPLICAS_INFO
+
+from .impl.named_columns import NamedColumns
+from .info.replicas import ReplicasInfoProtocol
+from .iterable_dataset import IterableDataset
+
+Batch = dict[str, torch.Tensor]
+
+
+class PartitionedIterableDataset(data.IterableDataset):
+    """
+    A dataset that implements iteration over partitioned data.
+
+    This implementation allows large amounts of data to be processed in batch-wise mode,
+    which is especially useful when used in distributed training.
+    """
+
+    def __init__(
+        self,
+        iterable: Iterable[NamedColumns],
+        batch_size: int,
+        generator: Optional[torch.Generator] = None,
+        replicas_info: ReplicasInfoProtocol = DEFAULT_REPLICAS_INFO,
+    ) -> None:
+        """
+        :param iterable: An iterable object that returns data partitions.
+        :param batch_size: Batch size.
+        :param generator: Random number generator for batch shuffling.
+            If ``None``, shuffling will be disabled. Default: ``None``.
+        :param replicas_info: A connector object capable of fetching total replica count and replica id during runtime.
+            Default: value of ``DEFAULT_REPLICAS_INFO`` - a pre-built connector which assumes standard Torch DDP mode.
+        """
+        super().__init__()
+
+        self.iterable = iterable
+
+        self.batch_size = batch_size
+        self.generator = generator
+        self.replicas_info = replicas_info
+
+    def __iter__(self) -> Iterator[Batch]:
+        for partition in iter(self.iterable):
+            iterable = IterableDataset(
+                named_columns=partition,
+                generator=self.generator,
+                batch_size=self.batch_size,
+                replicas_info=self.replicas_info,
+            )
+
+            yield from iter(iterable)

replay/data/nn/parquet/utils/compute_length.py

@@ -0,0 +1,66 @@
+import warnings
+from collections.abc import Iterable
+from typing import Protocol
+
+from replay.data.nn.parquet.info.partitioning import partitioning_per_replica
+from replay.data.nn.parquet.iterator import BatchesIterator
+
+
+class HasLengthProtocol(Protocol):
+    def __len__(self) -> int: ...
+
+
+def compute_fixed_size_generic_length_from_sizes(
+    partition_sizes: Iterable[int], batch_size: int, num_replicas: int
+) -> int:
+    residue = 0
+    batch_counter = 0
+    for partition_size in partition_sizes:
+        per_replica = partitioning_per_replica(partition_size, num_replicas)
+        batch_count = per_replica // batch_size
+        residue += per_replica % batch_size
+        if batch_size < residue:
+            batch_count += residue // batch_size
+            residue = residue % batch_size
+        batch_counter += batch_count
+    batch_counter += residue > 0
+    return batch_counter
+
+
+def compute_fixed_size_batches_length(iterable: BatchesIterator, batch_size: int, num_replicas: int) -> int:
+    assert isinstance(iterable, BatchesIterator)
+
+    partition_size = iterable.batch_size
+
+    def default_partitions(fragment_size: int) -> list[int]:
+        full_partitions_count = fragment_size // partition_size
+        result = [partition_size] * full_partitions_count
+        if (residue := (fragment_size % partition_size)) > 0:
+            result.append(residue)
+        return result
+
+    partition_sizes = []
+    for fragment in iterable.dataset.get_fragments():
+        fragment_size = fragment.count_rows()
+        partitions = default_partitions(fragment_size)
+        partition_sizes.extend(partitions)
+
+    result = compute_fixed_size_generic_length_from_sizes(
+        partition_sizes=partition_sizes,
+        num_replicas=num_replicas,
+        batch_size=batch_size,
+    )
+
+    return result
+
+
+def compute_fixed_size_generic_length(iterable: Iterable[HasLengthProtocol], batch_size: int, num_replicas: int) -> int:
+    warnings.warn("Generic length computation. This may cause performance issues.", UserWarning, stacklevel=2)
+    return compute_fixed_size_generic_length_from_sizes(map(len, iterable), batch_size, num_replicas)
+
+
+def compute_fixed_size_length(iterable: Iterable[HasLengthProtocol], batch_size: int, num_replicas: int) -> int:
+    if isinstance(iterable, BatchesIterator):
+        return compute_fixed_size_batches_length(iterable, batch_size, num_replicas)
+    else:
+        return compute_fixed_size_generic_length(iterable, batch_size, num_replicas)

replay/data/nn/schema.py

@@ -86,12 +86,14 @@ class TensorFeatureInfo:
             default: ``None``.
         :param feature_sources: columns names and DataFrames feature came from,
             default: ``None``.
-        :param cardinality: cardinality of categorical feature, required for ids columns,
-            optional for others,
-            default: ``None``.
-        :param padding_value: value to pad sequences to desired length
-        :param embedding_dim: embedding dimensions of categorical feature,
-            default: ``None``.
+        :param cardinality: cardinality of categorical feature.
+            number of unique items in vocabulary (catalog).
+            The specified cardinality value must not take into account the padding value.
+            Default: ``None``.
+        :param padding_value: value to pad sequences to desired length.
+            It is recommended to set the padding value for categorical features in the `cardinality` value.
+        :param embedding_dim: embedding dimensions of the feature.
+            Default: ``None`` - it means will be used value of ``DEFAULT_EMBEDDING_DIM``.
         :param tensor_dim: tensor dimensions of numerical feature,
             default: ``None``.
         """
@@ -106,8 +108,8 @@ class TensorFeatureInfo:
             raise ValueError(msg)
         self._feature_type = feature_type
 
-        if feature_type in [FeatureType.NUMERICAL, FeatureType.NUMERICAL_LIST] and (cardinality or embedding_dim):
-            msg = "Cardinality and embedding dimensions are needed only with categorical feature type."
+        if feature_type in [FeatureType.NUMERICAL, FeatureType.NUMERICAL_LIST] and cardinality is not None:
+            msg = "Cardinality is needed only with categorical feature type."
             raise ValueError(msg)
         self._cardinality = cardinality
 
@@ -115,9 +117,8 @@ class TensorFeatureInfo:
             msg = "Tensor dimensions is needed only with numerical feature type."
             raise ValueError(msg)
 
-        if feature_type in [FeatureType.CATEGORICAL, FeatureType.CATEGORICAL_LIST]:
-            self._embedding_dim = embedding_dim or self.DEFAULT_EMBEDDING_DIM
-        else:
+        self._embedding_dim = embedding_dim or self.DEFAULT_EMBEDDING_DIM
+        if feature_type in [FeatureType.NUMERICAL, FeatureType.NUMERICAL_LIST]:
             self._tensor_dim = tensor_dim
 
     @property
@@ -236,9 +237,6 @@ class TensorFeatureInfo:
         """
         :returns: Embedding dimensions of the feature.
         """
-        if not self.is_cat:
-            msg = f"Can not get embedding dimensions because feature type of {self.name} feature is not categorical."
-            raise RuntimeError(msg)
         return self._embedding_dim
 
     def _set_embedding_dim(self, embedding_dim: int) -> None:

replay/data/nn/sequence_tokenizer.py

@@ -10,6 +10,7 @@ import numpy as np
 import polars as pl
 from pandas import DataFrame as PandasDataFrame
 from polars import DataFrame as PolarsDataFrame
+from typing_extensions import deprecated
 
 from replay.data import Dataset, FeatureHint, FeatureSchema, FeatureSource, FeatureType
 from replay.data.dataset_utils import DatasetLabelEncoder
@@ -24,6 +25,7 @@ SequenceDataFrameLike = Union[PandasDataFrame, PolarsDataFrame]
 _T = TypeVar("_T")
 
 
+@deprecated("`SequenceTokenizer` class is deprecated.")
 class SequenceTokenizer:
     """
     Data tokenizer for transformers;
@@ -507,6 +509,7 @@ class SequenceTokenizer:
                 pickle.dump(self, file)
 
 
+@deprecated("`_BaseSequenceProcessor` class is deprecated.", stacklevel=2)
 class _BaseSequenceProcessor(Generic[_T]):
     """
     Base class for sequence processing
@@ -600,6 +603,7 @@ class _BaseSequenceProcessor(Generic[_T]):
         pass
 
 
+@deprecated("`_PandasSequenceProcessor` class is deprecated.", stacklevel=2)
 class _PandasSequenceProcessor(_BaseSequenceProcessor[PandasDataFrame]):
     """
     Class to process sequences of different categorical and numerical features.
@@ -780,6 +784,7 @@ class _PandasSequenceProcessor(_BaseSequenceProcessor[PandasDataFrame]):
         return values
 
 
+@deprecated("`_PolarsSequenceProcessor` class is deprecated.", stacklevel=2)
 class _PolarsSequenceProcessor(_BaseSequenceProcessor[PolarsDataFrame]):
     """
     Class to process sequences of different categorical and numerical features.

replay/data/nn/sequential_dataset.py

@@ -8,11 +8,13 @@ import pandas as pd
 import polars as pl
 from pandas import DataFrame as PandasDataFrame
 from polars import DataFrame as PolarsDataFrame
+from typing_extensions import deprecated
 
 if TYPE_CHECKING:
     from .schema import TensorSchema
 
 
+@deprecated("`SequentialDataset` class is deprecated.", stacklevel=2)
 class SequentialDataset(abc.ABC):
     """
     Abstract base class for sequential dataset
@@ -138,6 +140,7 @@ class SequentialDataset(abc.ABC):
         return df_converted
 
 
+@deprecated("`PandasSequentialDataset` class is deprecated.")
 class PandasSequentialDataset(SequentialDataset):
     """
     Sequential dataset that stores sequences in PandasDataFrame format.
@@ -234,6 +237,7 @@ class PandasSequentialDataset(SequentialDataset):
         return dataset
 
 
+@deprecated("`PolarsSequentialDataset` class is deprecated.")
 class PolarsSequentialDataset(PandasSequentialDataset):
     """
     Sequential dataset that stores sequences in PolarsDataFrame format.

replay/data/nn/torch_sequential_dataset.py

@@ -5,6 +5,7 @@ from typing import TYPE_CHECKING, NamedTuple, Optional, Union, cast
 import numpy as np
 import torch
 from torch.utils.data import Dataset as TorchDataset
+from typing_extensions import deprecated
 
 if TYPE_CHECKING:
     from .schema import TensorFeatureInfo, TensorMap, TensorSchema
@@ -13,6 +14,7 @@ if TYPE_CHECKING:
 
 # We do not use dataclasses as PyTorch default collate
 # function in dataloader supports only namedtuple
+@deprecated("`TorchSequentialBatch` class is deprecated.", stacklevel=2)
 class TorchSequentialBatch(NamedTuple):
     """
     Batch of TorchSequentialDataset
@@ -23,6 +25,7 @@ class TorchSequentialBatch(NamedTuple):
     features: "TensorMap"
 
 
+@deprecated("`TorchSequentialDataset` class is deprecated.")
 class TorchSequentialDataset(TorchDataset):
     """
     Torch dataset for sequential recommender models
@@ -160,6 +163,7 @@ class TorchSequentialDataset(TorchDataset):
                 yield (i, offset_from_seq_beginning)
 
 
+@deprecated("`TorchSequentialValidationBatch` class is deprecated.", stacklevel=2)
 class TorchSequentialValidationBatch(NamedTuple):
     """
     Batch of TorchSequentialValidationDataset
@@ -176,6 +180,7 @@ DEFAULT_GROUND_TRUTH_PADDING_VALUE = -1
 DEFAULT_TRAIN_PADDING_VALUE = -2
 
 
+@deprecated("`TorchSequentialValidationDataset` class is deprecated.")
 class TorchSequentialValidationDataset(TorchDataset):
     """
     Torch dataset for sequential recommender models that additionally stores ground truth

replay/data/utils/batching.py

@@ -0,0 +1,69 @@
+from functools import lru_cache
+from typing import Iterator, Tuple
+
+
+def validate_length(length: int) -> int:
+    if length < 1:
+        msg: str = f"Length is invalid. Got {length}."
+        raise ValueError(msg)
+    return length
+
+
+def validate_batch_size(batch_size: int) -> int:
+    if batch_size < 1:
+        msg: str = f"Batch Size is invalid. Got {batch_size}."
+        raise ValueError(msg)
+    return batch_size
+
+
+def validate_input(length: int, batch_size: int) -> Tuple[int, int]:
+    length = validate_length(length)
+    batch_size = validate_batch_size(batch_size)
+    return (length, batch_size)
+
+
+def uniform_batch_count(length: int, batch_size: int) -> int:
+    @lru_cache
+    def _uniform_batch_count(length: int, batch_size: int) -> int:
+        length, batch_size = validate_input(length, batch_size)
+        batch_count: int = length // batch_size
+        batch_count = batch_count + bool(length % batch_size)
+        assert batch_count >= 1
+        assert length <= batch_count * batch_size
+        assert (batch_count - 1) * batch_size < length
+        return batch_count
+
+    return _uniform_batch_count(length, batch_size)
+
+
+class UniformBatching:
+    def __init__(self, length: int, batch_size: int) -> None:
+        length, batch_size = validate_input(length, batch_size)
+
+        self.length: int = length
+        self.batch_size: int = batch_size
+
+    @property
+    def batch_count(self) -> int:
+        return uniform_batch_count(self.length, self.batch_size)
+
+    def __len__(self) -> int:
+        return self.batch_count
+
+    def get_limits(self, index: int) -> Tuple[int, int]:
+        if (index < 0) or (self.batch_count <= index):
+            msg: str = f"Batching Index is invalid. Got {index}."
+            raise IndexError(msg)
+        first: int = index * self.batch_size
+        last: int = min(self.length, first + self.batch_size)
+        assert (first >= 0) and (first < self.length)
+        assert (first < last) and (last <= self.length)
+        return (first, last)
+
+    def __getitem__(self, index: int) -> Tuple[int, int]:
+        return self.get_limits(index)
+
+    def __iter__(self) -> Iterator[Tuple[int, int]]:
+        index: int
+        for index in range(self.batch_count):
+            yield self.get_limits(index)

replay/data/utils/typing/dtype.py

@@ -0,0 +1,65 @@
+from functools import lru_cache
+
+import numpy as np
+import pyarrow as pa
+import torch
+
+
+@lru_cache
+def _torch_to_numpy(dtype: torch.dtype) -> np.dtype:
+    exemplar: torch.Tensor = torch.asarray([0], dtype=dtype)
+    return exemplar.numpy().dtype
+
+
+def torch_to_numpy(dtype: torch.dtype) -> np.dtype:
+    return _torch_to_numpy(dtype)
+
+
+@lru_cache
+def _numpy_to_torch(dtype: np.dtype) -> torch.dtype:
+    exemplar: np.ndarray = np.asarray([0], dtype=dtype)
+    return torch.from_numpy(exemplar).dtype
+
+
+def numpy_to_torch(dtype: np.dtype) -> torch.dtype:
+    return _numpy_to_torch(dtype)
+
+
+@lru_cache
+def _pyarrow_to_numpy(dtype: pa.DataType) -> np.dtype:
+    exemplar: pa.Array = pa.array([0], type=dtype)
+    return exemplar.to_numpy().dtype
+
+
+def pyarrow_to_numpy(dtype: pa.DataType) -> np.dtype:
+    return _pyarrow_to_numpy(dtype)
+
+
+@lru_cache
+def _numpy_to_pyarrow(dtype: np.dtype) -> pa.DataType:
+    exemplar: np.ndarray = np.asarray([0], dtype=dtype)
+    return pa.array(exemplar).type
+
+
+def numpy_to_pyarrow(dtype: np.dtype) -> pa.DataType:
+    return _numpy_to_pyarrow(dtype)
+
+
+@lru_cache
+def _torch_to_pyarrow(dtype: torch.dtype) -> pa.DataType:
+    np_dtype: np.dtype = torch_to_numpy(dtype)
+    return numpy_to_pyarrow(np_dtype)
+
+
+def torch_to_pyarrow(dtype: torch.dtype) -> pa.DataType:
+    return _torch_to_pyarrow(dtype)
+
+
+@lru_cache
+def _pyarrow_to_torch(dtype: pa.DataType) -> torch.dtype:
+    np_dtype: np.dtype = pyarrow_to_numpy(dtype)
+    return numpy_to_torch(np_dtype)
+
+
+def pyarrow_to_torch(dtype: pa.DataType) -> torch.dtype:
+    return _pyarrow_to_torch(dtype)

replay/metrics/torch_metrics_builder.py

@@ -139,7 +139,9 @@ class _CoverageHelper:
         """
         self._ensure_hists_on_device(train.device)
         flatten_train = train.flatten()
-        filtered_train = torch.masked_select(flatten_train, flatten_train != -2)
+        filtered_train = torch.masked_select(
+            flatten_train, ((flatten_train >= 0) & (flatten_train <= self.item_count - 1))
+        )
         self._train_hist += torch.histc(filtered_train.float(), bins=self.item_count, min=0, max=self.item_count - 1)
 
     def get_metrics(self) -> Mapping[str, float]:
@@ -193,7 +195,7 @@ class _MetricBuilder(abc.ABC):
 
 class TorchMetricsBuilder(_MetricBuilder):
     """
-    Computes specified metrics over multiple batches
+    Computes specified metrics over multiple batches.
     """
 
     def __init__(
@@ -203,12 +205,12 @@ class TorchMetricsBuilder(_MetricBuilder):
         item_count: Optional[int] = None,
     ) -> None:
         """
-        :param metrics: (list[MetricName]): Names of metrics to calculate.
-            Default: `["map", "ndcg", "recall"]`.
-        :param top_k: (list): Consider the highest k scores in the ranking.
-            Default: `[1, 5, 10, 20]`.
-        :param item_count: (optional, int): the total number of items in the dataset.
-            You can omit this parameter if you don't need to calculate the Coverage metric.
+        :param metrics: Names of metrics to calculate.
+            Default: ``["map", "ndcg", "recall"]``.
+        :param top_k: Consider the highest k scores in the ranking.
+            Default: ``[1, 5, 10, 20]``.
+        :param item_count: the total number of items in the dataset.
+            You can omit this parameter if you don't need to calculate the ``Coverage`` metric.
         """
         self._mr = _MetricRequirements.from_metrics(
             set(metrics),
@@ -272,12 +274,16 @@ class TorchMetricsBuilder(_MetricBuilder):
         """
         Add a batch with predictions, ground truth and train set to calculate the metrics.
 
-        :param predictions: (torch.LongTensor): A batch with the same number of recommendations for each user.
-        :param ground_truth: (torch.LongTensor): A batch corresponding to the test set for each user.
-            If users have a test set of different sizes then you need to do the padding using -1.
-        :param train: (optional, int): A batch corresponding to the train set for each user.
-            If users have a train set of different sizes then you need to do the padding using -2.
-            You can omit this parameter if you don't need to calculate the coverage or novelty metrics.
+        :param predictions: A batch with the same number of recommendations for each user.
+        :param ground_truth: A batch corresponding to the test set for each user.
+            If users have a test set of different sizes then you need to do
+            the padding using a value that is not found in the item ID's.
+            For example, these can be negative values.
+        :param train: A batch corresponding to the train set for each user.
+            If users have a train set of different sizes then you need to do
+            the padding using a value that is not found in the item ID's and ``ground_truth``.
+            For example, these can be negative values.
+            You can omit this parameter if you don't need to calculate the ``coverage`` or ``novelty`` metrics.
         """
         self._ensure_constants_on_device(predictions.device)
         metrics_sum = np.array(self._compute_metrics_sum(predictions, ground_truth, train), dtype=np.float64)

replay/models/nn/loss/sce.py

@@ -6,9 +6,9 @@ import torch
 
 @dataclass(frozen=True)
 class SCEParams:
-    """Set of parameters for ScalableCrossEntropyLoss.
+    """
+    Set of parameters for ScalableCrossEntropyLoss.
 
-    Constructor arguments:
     :param n_buckets: Number of buckets into which samples will be distributed.
     :param bucket_size_x: Number of item hidden representations that will be in each bucket.
     :param bucket_size_y: Number of item embeddings that will be in each bucket.
@@ -33,11 +33,6 @@ class ScalableCrossEntropyLoss:
 
         :param SCEParams: Dataclass with ScalableCrossEntropyLoss parameters.
             Dataclass contains following values:
-                :param n_buckets: Number of buckets into which samples will be distributed.
-                :param bucket_size_x: Number of item hidden representations that will be in each bucket.
-                :param bucket_size_y: Number of item embeddings that will be in each bucket.
-                :param mix_x: Whether a randomly generated matrix will be multiplied by the model output matrix or not.
-                    Default: ``False``.
         """
         assert all(
             param is not None for param in sce_params._get_not_none_params()

replay/models/nn/optimizer_utils/__init__.py

@@ -1,4 +1,9 @@
 from replay.utils import TORCH_AVAILABLE
 
 if TORCH_AVAILABLE:
-    from .optimizer_factory import FatLRSchedulerFactory, FatOptimizerFactory, LRSchedulerFactory, OptimizerFactory
+    from .optimizer_factory import (
+        FatLRSchedulerFactory,
+        FatOptimizerFactory,
+        LRSchedulerFactory,
+        OptimizerFactory,
+    )

replay/models/nn/optimizer_utils/optimizer_factory.py

@@ -2,8 +2,13 @@ import abc
 from collections.abc import Iterator
 
 import torch
+from typing_extensions import deprecated
 
 
+@deprecated(
+    "`OptimizerFactory` class is deprecated. Use `replay.nn.lightning.optimizer.BaseOptimizerFactory` instead.",
+    stacklevel=2,
+)
 class OptimizerFactory(abc.ABC):
     """
     Interface for optimizer factory
@@ -20,6 +25,10 @@ class OptimizerFactory(abc.ABC):
         """
 
 
+@deprecated(
+    "`LRSchedulerFactory` class is deprecated. Use `replay.nn.lightning.scheduler.BaseLRSchedulerFactory` instead.",
+    stacklevel=2,
+)
 class LRSchedulerFactory(abc.ABC):
     """
     Interface for learning rate scheduler factory
@@ -36,6 +45,9 @@ class LRSchedulerFactory(abc.ABC):
         """
 
 
+@deprecated(
+    "`FatOptimizerFactory` class is deprecated. Use `replay.nn.lightning.optimizer.OptimizerFactory` instead.",
+)
 class FatOptimizerFactory(OptimizerFactory):
     """
     Factory that creates optimizer depending on passed parameters
@@ -75,6 +87,9 @@ class FatOptimizerFactory(OptimizerFactory):
         raise ValueError(msg)
 
 
+@deprecated(
+    "`FatLRSchedulerFactory` class is deprecated. Use `replay.nn.lightning.scheduler.LRSchedulerFactory` instead.",
+)
 class FatLRSchedulerFactory(LRSchedulerFactory):
     """
     Factory that creates learning rate schedule depending on passed parameters