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

replay/nn/sequential/twotower/reader.py

@@ -0,0 +1,89 @@
+from typing import Protocol
+
+import numpy as np
+import pandas as pd
+import torch
+
+from replay.data import FeatureSource
+from replay.data.nn import TensorSchema
+
+
+class FeaturesReaderProtocol(Protocol):
+    def __getitem__(self, key: str) -> torch.Tensor: ...
+
+
+class FeaturesReader:
+    """
+    Prepares a dict of item features values that will be used for training and inference of the Item Tower.
+    """
+
+    def __init__(self, schema: TensorSchema, metadata: dict, path: str):
+        """
+        :param schema: the same tensor schema used in TwoTower model.
+        :param metadata: A dictionary of feature names that
+            associated with its shape and padding_value.\n
+            Example: {"item_id" : {"shape": 100, "padding": 7657}}.\n
+            For details, see the section :ref:`parquet-processing`.
+        :param path: path to parquet with dataframe of item features.\n
+            **Note:**\n
+            1. Dataframe columns must be already encoded.\n
+            2. Every feature for item "tower" in `schema` must contain ``feature_sources`` with the names
+               of the source features to create correct inverse mapping.
+               Also, for each such feature one of the requirements must be met: the ``schema`` for the feature must
+               contain ``feature_sources`` with a source of type FeatureSource.ITEM_FEATURES
+               or hint type FeatureHint.ITEM_ID.
+
+        """
+        item_feature_names = [
+            info.feature_source.column
+            for name, info in schema.items()
+            if info.feature_source.source == FeatureSource.ITEM_FEATURES or name == schema.item_id_feature_name
+        ]
+        metadata_names = list(metadata.keys())
+
+        if (unique_metadata_names := set(metadata_names)) != (unique_schema_names := set(item_feature_names)):
+            extra_metadata_names = unique_metadata_names - unique_schema_names
+            if extra_metadata_names:
+                msg = (
+                    "The metadata contains information about the following columns,"
+                    f"which are not described in schema: {extra_metadata_names}."
+                )
+                raise ValueError(msg)
+
+            extra_schema_names = unique_schema_names - unique_metadata_names
+            if extra_schema_names:
+                msg = (
+                    "The schema contains information about the following columns,"
+                    f"which are not described in metadata: {extra_schema_names}."
+                )
+                raise ValueError(msg)
+
+        features = pd.read_parquet(
+            path=path,
+            columns=metadata_names,
+        )
+
+        def add_padding(row: np.array, max_len: int, padding_value: int):
+            return np.concatenate(([padding_value] * (max_len - len(row)), row))
+
+        for k, v in metadata.items():
+            if not v:
+                continue
+            features[k] = features[k].apply(add_padding, args=(v["shape"], v["padding"]))
+
+        inverse_feature_names_mapping = {
+            schema[feature].feature_source.column: feature for feature in item_feature_names
+        }
+        features.rename(columns=inverse_feature_names_mapping, inplace=True)
+        features.sort_values(by=schema.item_id_feature_name, inplace=True)
+        features.reset_index(drop=True, inplace=True)
+
+        self._features = {}
+
+        for k in features.columns:
+            dtype = torch.float32 if schema[k].is_num else torch.int64
+            feature_tensor = torch.asarray(features[k], dtype=dtype)
+            self._features[k] = feature_tensor
+
+    def __getitem__(self, key: str) -> torch.Tensor:
+        return self._features[key]

replay/nn/transform/__init__.py

@@ -0,0 +1,22 @@
+from .copy import CopyTransform
+from .grouping import GroupTransform
+from .negative_sampling import MultiClassNegativeSamplingTransform, UniformNegativeSamplingTransform
+from .next_token import NextTokenTransform
+from .rename import RenameTransform
+from .reshape import UnsqueezeTransform
+from .sequence_roll import SequenceRollTransform
+from .token_mask import TokenMaskTransform
+from .trim import TrimTransform
+
+__all__ = [
+    "CopyTransform",
+    "GroupTransform",
+    "MultiClassNegativeSamplingTransform",
+    "NextTokenTransform",
+    "RenameTransform",
+    "SequenceRollTransform",
+    "TokenMaskTransform",
+    "TrimTransform",
+    "UniformNegativeSamplingTransform",
+    "UnsqueezeTransform",
+]

replay/nn/transform/copy.py

@@ -0,0 +1,38 @@
+from __future__ import annotations
+
+import torch
+
+
+class CopyTransform(torch.nn.Module):
+    """
+    Copies a set of columns according to the provided mapping.
+    All copied columns are detached from the graph to prevent erroneous
+    differentiation.
+
+    Example:
+
+    .. code-block:: python
+
+        >>> input_batch = {"item_id_mask": torch.BoolTensor([False, True, True])}
+        >>> transform = CopyTransform({"item_id_mask" : "padding_id"})
+        >>> output_batch = transform(input_batch)
+        >>> output_batch
+        {'item_id_mask': tensor([False,  True,  True]),
+        'padding_id': tensor([False,  True,  True])}
+
+    """
+
+    def __init__(self, mapping: dict[str, str]) -> None:
+        """
+        :param mapping: A dictionary maps which source tensors will be copied into the batch with new names.
+            Tensors with new names will be copies of original ones, original tensors are stayed in batch.
+        """
+        super().__init__()
+        self.mapping = mapping
+
+    def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        output_batch = dict(batch.items())
+        output_batch |= {
+            out_column: output_batch[in_column].clone().detach() for in_column, out_column in self.mapping.items()
+        }
+        return output_batch

replay/nn/transform/grouping.py

@@ -0,0 +1,39 @@
+import torch
+
+
+class GroupTransform(torch.nn.Module):
+    """
+    Combines existing tensors from a batch moving them to the common groups.
+    The name of the shared keys and the keys to be moved are specified in ``mapping``.
+
+    Example:
+
+    .. code-block:: python
+
+        >>> input_batch = {
+        ...     "item_id": torch.LongTensor([[30, 22, 1]]),
+        ...     "item_feature": torch.LongTensor([[1, 11, 11]])
+        ... }
+        >>> transform = GroupTransform({"feature_tensors" : ["item_id", "item_feature"]})
+        >>> output_batch = transform(input_batch)
+        >>> output_batch
+        {'feature_tensors': {'item_id': tensor([[30, 22,  1]]),
+        'item_feature': tensor([[ 1, 11, 11]])}}
+
+    """
+
+    def __init__(self, mapping: dict[str, list[str]]) -> None:
+        """
+        :param mapping: A dict mapping new names to a list of existing names for grouping.
+        """
+        super().__init__()
+        self.mapping = mapping
+        self._grouped_keys = set().union(*mapping.values())
+
+    def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        output_batch = {k: v for k, v in batch.items() if k not in self._grouped_keys}
+
+        for group_name, feature_names in self.mapping.items():
+            output_batch[group_name] = {name: batch[name] for name in feature_names if name in batch}
+
+        return output_batch

replay/nn/transform/negative_sampling.py

@@ -0,0 +1,182 @@
+from typing import Optional
+
+import torch
+
+
+class UniformNegativeSamplingTransform(torch.nn.Module):
+    """
+    Transform for global negative sampling.
+
+    For every batch, transform generates a vector of size ``(num_negative_samples)``
+    consisting of random indices sampeled from a range of ``cardinality``. Unless a custom sample
+    distribution is provided, the indices are weighted equally.
+
+    Example:
+
+    .. code-block:: python
+
+        >>> _ = torch.manual_seed(0)
+        >>> input_batch = {"item_id": torch.LongTensor([[1, 0, 4]])}
+        >>> transform = UniformNegativeSamplingTransform(cardinality=4, num_negative_samples=2)
+        >>> output_batch = transform(input_batch)
+        >>> output_batch
+        {'item_id': tensor([[1, 0, 4]]), 'negative_labels': tensor([2, 1])}
+
+    """
+
+    def __init__(
+        self,
+        cardinality: int,
+        num_negative_samples: int,
+        *,
+        out_feature_name: Optional[str] = "negative_labels",
+        sample_distribution: Optional[torch.Tensor] = None,
+        generator: Optional[torch.Generator] = None,
+    ) -> None:
+        """
+        :param cardinality: number of unique items in vocabulary (catalog).
+            The specified cardinality value must not take into account the padding value.
+        :param num_negative_samples: The size of negatives vector to generate.
+        :param out_feature_name: The name of result feature in batch.
+        :param sample_distribution: The weighs of indices in the vocabulary. If specified, must
+                match the ``cardinality``. Default: ``None``.
+        :param generator: Random number generator to be used for sampling
+                from the distribution. Default: ``None``.
+        """
+        if sample_distribution is not None and sample_distribution.size(-1) != cardinality:
+            msg = (
+                "The sample_distribution parameter has an incorrect size. "
+                f"Got {sample_distribution.size(-1)}, expected {cardinality}."
+            )
+            raise ValueError(msg)
+
+        if num_negative_samples >= cardinality:
+            msg = (
+                "The `num_negative_samples` parameter has an incorrect value."
+                f"Got {num_negative_samples}, expected less than cardinality of items catalog ({cardinality})."
+            )
+            raise ValueError(msg)
+
+        super().__init__()
+
+        self.out_feature_name = out_feature_name
+        self.num_negative_samples = num_negative_samples
+        self.generator = generator
+        if sample_distribution is not None:
+            self.sample_distribution = sample_distribution
+        else:
+            self.sample_distribution = torch.ones(cardinality)
+
+    def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        output_batch = dict(batch.items())
+
+        negatives = torch.multinomial(
+            self.sample_distribution,
+            num_samples=self.num_negative_samples,
+            replacement=False,
+            generator=self.generator,
+        )
+
+        output_batch[self.out_feature_name] = negatives.to(device=next(iter(output_batch.values())).device)
+        return output_batch
+
+
+class MultiClassNegativeSamplingTransform(torch.nn.Module):
+    """
+    Transform for generating negatives using a fixed class-assignment matrix.
+
+    For every batch, transform generates a tensor of size ``(N, num_negative_samples)``, where N is number of classes.
+    This tensor consists of random indices sampled using specified fixed class-assignment matrix.
+
+    Also, transform receives from batch by key a tensor ``negative_selector_name`` of shape (batch size,),
+    where i-th element in [0, N-1] specifies which class of N is used to select from sampled negatives that corresponds
+    to every i-th batch row (user's history sequence).
+
+    The resulting negatives tensor has shape of ``(batch_size, num_negative_samples)``.
+
+    Example:
+
+    .. code-block:: python
+
+        >>> _ = torch.manual_seed(0)
+        >>> sample_mask = torch.tensor([
+        ...     [1, 0, 1, 0, 0, 0],
+        ...     [0, 0, 0, 1, 1, 0],
+        ...     [0, 1, 0, 0, 0, 1],
+        ... ])
+        >>> input_batch = {"negative_selector": torch.tensor([0, 2, 1, 1, 0])}
+        >>> transform = MultiClassNegativeSamplingTransform(
+        ...                 num_negative_samples=2,
+        ...                 sample_mask=sample_mask
+        ... )
+        >>> output_batch = transform(input_batch)
+        >>> output_batch
+        {'negative_selector': tensor([0, 2, 1, 1, 0]),
+         'negative_labels': tensor([[2, 0],
+                 [5, 1],
+                 [3, 4],
+                 [3, 4],
+                 [2, 0]])}
+    """
+
+    def __init__(
+        self,
+        num_negative_samples: int,
+        sample_mask: torch.Tensor,
+        *,
+        negative_selector_name: Optional[str] = "negative_selector",
+        out_feature_name: Optional[str] = "negative_labels",
+        generator: Optional[torch.Generator] = None,
+    ) -> None:
+        """
+        :param num_negative_samples: The size of negatives vector to generate.
+        :param sample_mask: The class-assignment (indicator) matrix of shape: ``(N, number of items in catalog)``,
+            where ``sample_mask[n, i]`` is a weight (or binary indicator) of assigning item i to class n.
+        :param negative_selector_name: name of tensor in batch of shape (batch size,), where i-th element
+            in [0, N-1] specifies which class of N is used to get negatives corresponding to i-th ``query_id`` in batch.
+        :param out_feature_name: The name of result feature in batch.
+        :param generator: Random number generator to be used for sampling from the distribution. Default: ``None``.
+        """
+        if sample_mask.dim() != 2:
+            msg = (
+                "The `sample_mask` parameter has an incorrect shape."
+                f"Got {sample_mask.dim()}, expected shape: (number of classes, number of items in catalog)."
+            )
+            raise ValueError(msg)
+
+        if num_negative_samples >= sample_mask.size(-1):
+            msg = (
+                "The `num_negative_samples` parameter has an incorrect value."
+                f"Got {num_negative_samples}, expected less than cardinality of items catalog ({sample_mask.size(-1)})."
+            )
+            raise ValueError(msg)
+
+        super().__init__()
+
+        self.register_buffer("sample_mask", sample_mask.float())
+
+        self.num_negative_samples = num_negative_samples
+        self.negative_selector_name = negative_selector_name
+        self.out_feature_name = out_feature_name
+        self.generator = generator
+
+    def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        assert self.negative_selector_name in batch
+        assert batch[self.negative_selector_name].dim() == 1
+
+        negative_selector = batch[self.negative_selector_name]  # [batch_size]
+
+        # [N, num_negatives] - shape of negatives
+        negatives = torch.multinomial(
+            input=self.sample_mask,
+            num_samples=self.num_negative_samples,
+            replacement=False,
+            generator=self.generator,
+        )
+
+        # [N, num_negatives] -> [batch_size, num_negatives]
+        selected_negatives = negatives[negative_selector]
+
+        output_batch = dict(batch.items())
+        output_batch[self.out_feature_name] = selected_negatives.to(device=negative_selector.device)
+        return output_batch

replay/nn/transform/next_token.py

@@ -0,0 +1,100 @@
+from typing import List, Union
+
+import torch
+
+from replay.data.nn.parquet.impl.masking import DEFAULT_MASK_POSTFIX
+
+
+class NextTokenTransform(torch.nn.Module):
+    """
+    For the tensor specified by key ``label_field`` (typically "item_id") in the batch, this transform creates
+    a corresponding "labels" tensor with a key ``out_feature_name`` in the batch, shifted forward
+    by the specified ``shift`` value. This "labels" tensor are a target that model predicts.
+    Padding mask for "labels" is also created. For all the other features excepted ``query_features``,
+    last ``shift`` elements are truncated.
+
+    This transform is required for the sequential models optimizing next token prediction task.
+
+    **WARNING**: In order to facilitate the shifting, this transform
+    requires extra elements in the sequence. Therefore, when utilizing this
+    transform, ensure you're reading at least ``sequence_length`` + ``shift``
+    elements from your dataset. The resulting batch will have the relevant fields
+    trimmed to ``sequence_length``.
+
+    Example:
+
+    .. code-block:: python
+
+        >>> input_batch = {
+        ...     "user_id": torch.LongTensor([111]),
+        ...     "item_id": torch.LongTensor([[5, 0, 7, 4]]),
+        ...     "item_id_mask": torch.BoolTensor([[0, 1, 1, 1]])
+        ... }
+        >>> transform = NextTokenTransform(label_field="item_id", shift=1, query_features="user_id")
+        >>> output_batch = transform(input_batch)
+        >>> output_batch
+        {'user_id': tensor([111]),
+        'item_id': tensor([[5, 0, 7]]),
+        'item_id_mask': tensor([[False,  True,  True]]),
+        'positive_labels': tensor([[0, 7, 4]]),
+        'positive_labels_mask': tensor([[True, True, True]])}
+
+    """
+
+    def __init__(
+        self,
+        label_field: str,
+        shift: int = 1,
+        query_features: Union[List[str], str] = ["query_id", "query_id_mask"],
+        out_feature_name: str = "positive_labels",
+        mask_postfix: str = DEFAULT_MASK_POSTFIX,
+    ) -> None:
+        """
+        :param label_field: Name of target feature tensor to convert into labels.
+        :param shift: Number of sequence items to shift by. Default: `1`.
+        :param query_features: Name of the query column or list of user features.
+            These columns will be excepted from the shifting and will be stayed unchanged.
+            Default: ``["query_id", "query_id_mask"]``.
+        :param out_feature_name: The name of result feature in batch. Default: ``"positive_labels"``.
+        :param mask_postfix: Postfix to append to the mask feature corresponding to resulting feature.
+            Default: ``"_mask"``.
+        """
+        super().__init__()
+        self.label_field = label_field
+        self.shift = shift
+        self.query_features = [query_features] if isinstance(query_features, str) else query_features
+        self.out_feature_name = out_feature_name
+        self.mask_postfix = mask_postfix
+
+    def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        if batch[self.label_field].dim() < 2:
+            msg = (
+                f"Transform expects batch feature {self.label_field} to be sequential "
+                f"but tensor of shape {batch[self.label_field].shape} found."
+            )
+            raise ValueError(msg)
+
+        max_len = batch[self.label_field].shape[1]
+        if self.shift >= max_len:
+            msg = (
+                f"Transform with shift={self.shift} cannot be applied to sequences of length {max_len}."
+                "Decrease value of `shift` parameter in transform"
+            )
+            raise ValueError(msg)
+
+        target = {feature_name: batch[feature_name] for feature_name in self.query_features}
+        features = {key: value for key, value in batch.items() if key not in self.query_features}
+
+        sequentilal_features = [feature_name for feature_name, feature in features.items() if feature.dim() > 1]
+        for feature_name in features:
+            if feature_name in sequentilal_features:
+                target[feature_name] = batch[feature_name][:, : -self.shift, ...].clone()
+            else:
+                target[feature_name] = batch[feature_name]
+
+        target[self.out_feature_name] = batch[self.label_field][:, self.shift :, ...].clone()
+        target[f"{self.out_feature_name}{self.mask_postfix}"] = batch[f"{self.label_field}{self.mask_postfix}"][
+            :, self.shift :, ...
+        ].clone()
+
+        return target

replay/nn/transform/rename.py

@@ -0,0 +1,33 @@
+import torch
+
+
+class RenameTransform(torch.nn.Module):
+    """
+    Renames specific feature columns into new ones. Changes names in original dict, not creates a new dict.
+    Example:
+
+    .. code-block:: python
+
+        >>> input_batch = {"item_id_mask": torch.BoolTensor([False, True, True])}
+        >>> transform = RenameTransform({"item_id_mask" : "padding_id"})
+        >>> output_batch = transform(input_batch)
+        >>> output_batch
+        {'padding_id': tensor([False,  True,  True])}
+
+    """
+
+    def __init__(self, mapping: dict[str, str]) -> None:
+        """
+        :param mapping: A dict mapping existing names into new ones.
+        """
+        super().__init__()
+        self.mapping = mapping
+
+    def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        output_batch = {}
+
+        for original_name, tensor in batch.items():
+            target_name = self.mapping.get(original_name, original_name)
+            output_batch[target_name] = tensor
+
+        return output_batch

replay/nn/transform/reshape.py

@@ -0,0 +1,41 @@
+import torch
+
+
+class UnsqueezeTransform(torch.nn.Module):
+    """
+    Unsqueeze specified tensor along specified dimension.
+
+    Example:
+
+    .. code-block:: python
+
+        >>> input_batch = {"padding_id": torch.BoolTensor([False, True, True])}
+        >>> transform = UnsqueezeTransform("padding_id", dim=0)
+        >>> output_batch = transform(input_batch)
+        >>> output_batch
+        {'padding_id': tensor([[False,  True,  True]])}
+
+    """
+
+    def __init__(self, column_name: str, dim: int) -> None:
+        """
+        :param column_name: Name of tensor to be unsqueezed.
+        :param dim: Dimension along which tensor will be unsqueezed.
+        """
+        super().__init__()
+        self.column_name = column_name
+        self.dim = dim
+
+    def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        if self.dim > batch[self.column_name].ndim - 1:
+            msg = (
+                "The dim parameter is incorrect."
+                f"Expected unsqueezing by {self.dim} dimension,"
+                f"but got the tensor with {batch[self.column_name].ndim} dimensions."
+            )
+            raise ValueError(msg)
+
+        output_batch = {k: v for k, v in batch.items() if k != self.column_name}
+        output_batch[self.column_name] = batch[self.column_name].unsqueeze(self.dim)
+
+        return output_batch

replay/nn/transform/sequence_roll.py

@@ -0,0 +1,48 @@
+import torch
+
+
+class SequenceRollTransform(torch.nn.Module):
+    """
+    Rolls the data along axis 1 by the specified amount
+    and fills the remaining positions by specified padding value.
+
+    Example:
+
+    .. code-block:: python
+
+        >>> input_tensor = {"item_id": torch.LongTensor([[2, 3, 1]])}
+        >>> transform = SequenceRollTransform("item_id", roll=-1, padding_value=10)
+        >>> output_tensor = transform(input_tensor)
+        >>> output_tensor
+        {'item_id': tensor([[ 3,  1, 10]])}
+
+    """
+
+    def __init__(
+        self,
+        field_name: str,
+        roll: int = -1,
+        padding_value: int = 0,
+    ) -> None:
+        """
+        :param field_name: Name of the target column from the batch to be rolled.
+        :param roll: Number of positions to roll by. Default: ``-1``.
+        :param padding_value: The value to use as padding for the sequence. Default: ``0``.
+        """
+        super().__init__()
+        self.field_name = field_name
+        self.roll = roll
+        self.padding_value = padding_value
+
+    def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+        output_batch = {k: v for k, v in batch.items() if k != self.field_name}
+
+        rolled_seq = batch[self.field_name].roll(self.roll, dims=1)
+
+        if self.roll > 0:
+            rolled_seq[:, : self.roll, ...] = self.padding_value
+        else:
+            rolled_seq[:, self.roll :, ...] = self.padding_value
+
+        output_batch[self.field_name] = rolled_seq
+        return output_batch

replay/nn/transform/template/__init__.py

@@ -0,0 +1,2 @@
+from .sasrec import make_default_sasrec_transforms
+from .twotower import make_default_twotower_transforms

replay/nn/transform/template/sasrec.py

@@ -0,0 +1,53 @@
+import copy
+
+import torch
+
+from replay.data.nn import TensorSchema
+from replay.nn.transform import GroupTransform, NextTokenTransform, RenameTransform, UnsqueezeTransform
+
+
+def make_default_sasrec_transforms(
+    tensor_schema: TensorSchema, query_column: str = "query_id"
+) -> dict[str, list[torch.nn.Module]]:
+    """
+    Creates a valid transformation pipeline for SasRec data batches.
+
+    Generated pipeline expects input dataset to contain the following columns:
+        1) Query ID column, specified by ``query_column``.
+        2) Item ID column, specified in the tensor schema.
+
+    :param tensor_schema: TensorSchema used to infer feature columns.
+    :param query_column: Name of the column containing query IDs. Default: ``"query_id"``.
+    :return: dict of transforms specified for every dataset split (train, validation, test, predict).
+    """
+    item_column = tensor_schema.item_id_feature_name
+    train_transforms = [
+        NextTokenTransform(label_field=item_column, query_features=query_column, shift=1),
+        RenameTransform(
+            {
+                query_column: "query_id",
+                f"{item_column}_mask": "padding_mask",
+                "positive_labels_mask": "target_padding_mask",
+            }
+        ),
+        UnsqueezeTransform("target_padding_mask", -1),
+        UnsqueezeTransform("positive_labels", -1),
+        GroupTransform({"feature_tensors": [item_column]}),
+    ]
+
+    val_transforms = [
+        RenameTransform({query_column: "query_id", f"{item_column}_mask": "padding_mask"}),
+        GroupTransform({"feature_tensors": [item_column]}),
+    ]
+    test_transforms = copy.deepcopy(val_transforms)
+
+    predict_transforms = copy.deepcopy(val_transforms)
+
+    transforms = {
+        "train": train_transforms,
+        "validate": val_transforms,
+        "test": test_transforms,
+        "predict": predict_transforms,
+    }
+
+    return transforms

replay/nn/transform/template/twotower.py

@@ -0,0 +1,22 @@
+import torch
+
+from replay.data.nn import TensorSchema
+
+from .sasrec import make_default_sasrec_transforms
+
+
+def make_default_twotower_transforms(
+    tensor_schema: TensorSchema, query_column: str = "query_id"
+) -> dict[str, list[torch.nn.Module]]:
+    """
+    Creates a valid transformation pipeline for TwoTower data batches.
+
+    Generated pipeline expects input dataset to contain the following columns:
+        1) Query ID column, specified by ``query_column``.
+        2) Item ID column, specified in the tensor schema.
+
+    :param tensor_schema: TensorSchema used to infer feature columns.
+    :param query_column: Name of the column containing query IDs. Default: ``"query_id"``.
+    :return: dict of transforms specified for every dataset split (train, validation, test, predict).
+    """
+    return make_default_sasrec_transforms(tensor_schema=tensor_schema, query_column=query_column)