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

replay/nn/normalization.py

@@ -0,0 +1,9 @@
+from typing import Protocol
+
+import torch
+
+
+class NormalizerProto(Protocol):
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor: ...
+
+    def reset_parameters(self) -> None: ...

replay/nn/output.py

@@ -0,0 +1,37 @@
+from typing import TypedDict
+
+import torch
+from typing_extensions import NotRequired
+
+
+class TrainOutput(TypedDict):
+    """
+    Storing outputs from models training stage.
+
+    :param loss: a tensor containing the calculated loss.\n
+        It is important that the tensor contains a gradient to call back propagation from the outside.
+    :param hidden_states: Tuple of `torch.Tensor`.\n
+        One for the output of the embeddings, if the model has an embedding layer, +
+        one for the output of each layer.\n
+        Expected shape: ``(batch_size, sequence_length, hidden_size)``.
+    """
+
+    loss: torch.Tensor
+    hidden_states: NotRequired[tuple[torch.Tensor, ...]]
+
+
+class InferenceOutput(TypedDict):
+    """
+    Storing outputs from models inference stage.
+
+    :param logits:
+        Sequence of hidden-states at the output of the last layer of the model.\n
+        Expected shape: ``(batch_size, sequence_length, hidden_size)``.
+    :param hidden_states: Tuple of `torch.Tensor`
+        (one for the output of the embeddings, if the model has an embedding layer, +
+        one for the output of each layer).\n
+        Expected shape: ``(batch_size, sequence_length, hidden_size)``.
+    """
+
+    logits: torch.Tensor
+    hidden_states: NotRequired[tuple[torch.Tensor, ...]]

replay/nn/sequential/__init__.py

@@ -0,0 +1,9 @@
+from .sasrec import (
+    DiffTransformerBlock,
+    DiffTransformerLayer,
+    PositionAwareAggregator,
+    SasRec,
+    SasRecBody,
+    SasRecTransformerLayer,
+)
+from .twotower import ItemTower, QueryTower, TwoTower, TwoTowerBody

replay/nn/sequential/sasrec/__init__.py

@@ -0,0 +1,7 @@
+from .agg import PositionAwareAggregator
+from .diff_transformer import (
+    DiffTransformerBlock,
+    DiffTransformerLayer,
+)
+from .model import SasRec, SasRecBody
+from .transformer import SasRecTransformerLayer

replay/nn/sequential/sasrec/agg.py

@@ -0,0 +1,53 @@
+import contextlib
+
+import torch
+
+from replay.data.nn.schema import TensorMap
+from replay.nn.agg import AggregatorProto
+
+
+class PositionAwareAggregator(torch.nn.Module):
+    """
+    The layer for aggregating embeddings and adding positional encoding.
+    """
+
+    def __init__(
+        self,
+        embedding_aggregator: AggregatorProto,
+        max_sequence_length: int,
+        dropout: float,
+    ) -> None:
+        """
+        :param embedding_aggregator: An object of a class that performs the logic of aggregating multiple embeddings.\n
+            For example, it can be a ``sum``, a ``mean``, or a ``concatenation``.
+        :param max_sequence_length: Max length of sequence.
+        :param dropout: probability of an element to be zeroed.
+        """
+        super().__init__()
+        self.embedding_aggregator = embedding_aggregator
+        self.pe = torch.nn.Embedding(max_sequence_length, self.embedding_aggregator.embedding_dim)
+        self.dropout = torch.nn.Dropout(p=dropout)
+
+    def reset_parameters(self) -> None:
+        self.embedding_aggregator.reset_parameters()
+        for _, param in self.pe.named_parameters():
+            with contextlib.suppress(ValueError):
+                torch.nn.init.xavier_normal_(param.data)
+
+    def forward(self, feature_tensors: TensorMap) -> torch.Tensor:
+        """
+        :param feature_tensors: a dictionary of tensors to pass into ``embedding_aggregator``.
+
+        :returns: Aggregated embeddings with positional encoding.
+        """
+        seqs: torch.Tensor = self.embedding_aggregator(feature_tensors)
+        assert seqs.dim() == 3
+        batch_size, seq_len, embedding_dim = seqs.size()
+        assert (
+            seq_len <= self.pe.num_embeddings
+        ), f"Sequence length = {seq_len} is greater then positional embedding num = {self.pe.num_embeddings}"
+
+        seqs *= embedding_dim**0.5
+        seqs += self.pe.weight[:seq_len].unsqueeze(0).repeat(batch_size, 1, 1)
+        seqs = self.dropout(seqs)
+        return seqs

replay/nn/sequential/sasrec/diff_transformer.py

@@ -0,0 +1,125 @@
+import math
+
+import torch
+
+from replay.data.nn import TensorMap
+from replay.nn.attention import MultiHeadDifferentialAttention
+from replay.nn.ffn import SwiGLU
+
+
+class DiffTransformerBlock(torch.nn.Module):
+    """
+    Single Block of the DiffTransformer Architecture.
+    Consists of Multi-Head Differential Attention followed by a SwiGLU Feed-Forward Network.
+
+    Source paper: https://arxiv.org/pdf/2410.05258
+    """
+
+    def __init__(self, embedding_dim: int, num_heads: int, lambda_init: float):
+        """
+        :param embedding_dim: Total dimension of the model. Must be divisible by ``num_heads``.
+        :param num_heads: Number of parallel attention heads.
+        :param lambda_init: Initial value for lambda.
+        """
+        super().__init__()
+        self.attn_norm = torch.nn.RMSNorm(embedding_dim)
+        self.attn = MultiHeadDifferentialAttention(embedding_dim, num_heads, lambda_init, vdim=2 * embedding_dim)
+        self.ff_norm = torch.nn.RMSNorm(embedding_dim)
+        self.ff = SwiGLU(embedding_dim, 2 * embedding_dim)
+
+    def reset_parameters(self) -> None:
+        self.attn_norm.reset_parameters()
+        self.attn.reset_parameters()
+        self.ff_norm.reset_parameters()
+        self.ff.reset_parameters()
+
+    def forward(
+        self,
+        input_embeddings: torch.Tensor,
+        attention_mask: torch.FloatTensor,
+    ) -> torch.Tensor:
+        """
+        Forward pass for a single differential transformer block.
+
+        :param input_embeddings: Input tensor of shape ``(batch_size, sequence_length, embedding_dim)``.
+        :param attention_mask: Causal-like mask for attention pattern, where ``-inf`` for ``PAD``, ``0`` - otherwise.\n
+            Possible shapes:\n
+            1. ``(batch_size * num_heads, sequence_length, sequence_length)``
+            2. ``(batch_size, num_heads, sequence_length, sequence_length)``
+        :returns: Output tensor after processing through the block.
+        """
+        # Apply Multi-Head Differential Attention with residual connection
+        attent_emb = self.attn(
+            input_embeddings,
+            input_embeddings,
+            input_embeddings,
+            attention_mask,
+        )
+        attention_block_out = self.attn_norm(attent_emb + input_embeddings)
+
+        # Apply SwiGLU Feed-Forward Network with residual connection
+        ff_out = self.ff(input_embeddings=attention_block_out)
+        feedforward_block_out = self.ff_norm(ff_out + attention_block_out)
+        return feedforward_block_out
+
+
+class DiffTransformerLayer(torch.nn.Module):
+    """
+    Stacked blocks of the DiffTransformer Architecture.
+    Single block consists of Multi-Head Differential Attention followed by a SwiGLU Feed-Forward Network.
+
+    Source paper: https://arxiv.org/pdf/2410.05258\n
+    Reference: https://github.com/nanowell/Differential-Transformer-PyTorch/blob/main/DiffTransformer.py
+    """
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        num_heads: int,
+        num_blocks: int,
+    ) -> None:
+        """
+        :param embedding_dim: Total dimension of the model. Must be divisible by num_heads.
+        :param num_heads: Number of parallel attention heads.
+        :param num_blocks: Number of Transformer blocks.
+        """
+        torch.nn.MultiheadAttention
+        super().__init__()
+        self.layers = torch.nn.ModuleList(
+            [
+                DiffTransformerBlock(
+                    embedding_dim=embedding_dim,
+                    num_heads=num_heads,
+                    lambda_init=0.8 - 0.6 * math.exp(-0.3 * block_num),
+                )
+                for block_num in range(num_blocks)
+            ]
+        )
+
+    def reset_parameters(self) -> None:
+        for layer in self.layers:
+            layer.reset_parameters()
+
+    def forward(
+        self,
+        feature_tensors: TensorMap,  # noqa: ARG002
+        input_embeddings: torch.Tensor,
+        padding_mask: torch.BoolTensor,  # noqa: ARG002
+        attention_mask: torch.FloatTensor,
+    ) -> torch.Tensor:
+        """
+        forward(input_embeddings, attention_mask)
+        :param input_embeddings: Input tensor of shape ``(batch_size, sequence_length, embedding_dim)``.
+        :param attention_mask: Causal-like mask for attention pattern, where ``-inf`` for ``PAD``, ``0`` - otherwise.\n
+            Possible shapes:\n
+            1. ``(batch_size * num_heads, sequence_length, sequence_length)``
+            2. ``(batch_size, num_heads, sequence_length, sequence_length)``
+        :returns: Output tensor after processing through the layer.
+        """
+        seqs = input_embeddings
+        for layer in self.layers:
+            seqs = layer(
+                input_embeddings=seqs,
+                attention_mask=attention_mask,
+            )
+        return seqs

replay/nn/sequential/sasrec/model.py

@@ -0,0 +1,377 @@
+from collections.abc import Sequence
+from typing import Literal, Optional, Protocol, Union
+
+import torch
+
+from replay.data.nn import TensorMap, TensorSchema
+from replay.nn.agg import AggregatorProto
+from replay.nn.head import EmbeddingTyingHead
+from replay.nn.loss import LossProto
+from replay.nn.mask import AttentionMaskProto
+from replay.nn.normalization import NormalizerProto
+from replay.nn.output import InferenceOutput, TrainOutput
+from replay.nn.utils import warning_is_not_none
+
+
+class EmbedderProto(Protocol):
+    def get_item_weights(
+        self,
+        indices: Optional[torch.LongTensor],
+    ) -> torch.Tensor: ...
+
+    def forward(
+        self,
+        feature_tensors: TensorMap,
+        feature_names: Optional[Sequence[str]] = None,
+    ) -> TensorMap: ...
+
+    def reset_parameters(self) -> None: ...
+
+
+class EncoderProto(Protocol):
+    def forward(
+        self,
+        feature_tensors: TensorMap,
+        input_embeddings: torch.Tensor,
+        padding_mask: torch.BoolTensor,
+        attention_mask: torch.Tensor,
+    ) -> torch.Tensor: ...
+
+    def reset_parameters(self) -> None: ...
+
+
+class SasRecBody(torch.nn.Module):
+    """
+    Implementation of the architecture of the SasRec model.\n
+    It can include various self-written blocks for modifying the model,
+    but the sequence of applying layers is fixed in accordance with the original architecture.
+    """
+
+    def __init__(
+        self,
+        embedder: EmbedderProto,
+        embedding_aggregator: AggregatorProto,
+        attn_mask_builder: AttentionMaskProto,
+        encoder: EncoderProto,
+        output_normalization: NormalizerProto,
+    ):
+        """
+        :param embedder: An object of a class that performs the logic of
+            generating embeddings from an input set of tensors.
+        :param embedding_aggregator: An object of a class that performs the logic of aggregating multiple embeddings.\n
+            For example, it can be a ``sum``, a ``mean``, or a ``concatenation``.
+        :param attn_mask_builder: An object of a class that performs the logic of
+            generating an attention mask based on the features and padding mask given to the model.
+        :param encoder: An object of a class that performs the logic of generating
+            a hidden embedding representation based on
+            features, padding masks, attention mask, and aggregated embedding.
+        :param output_normalization: An object of a class that performs the logic of
+            normalization of the hidden state obtained from the encoder.\n
+            For example, it may be a ``torch.nn.LayerNorm`` or ``torch.nn.RMSNorm``.
+        """
+        super().__init__()
+        self.embedder = embedder
+        self.attn_mask_builder = attn_mask_builder
+        self.embedding_aggregator = embedding_aggregator
+        self.encoder = encoder
+        self.output_normalization = output_normalization
+
+    def reset_parameters(self) -> None:
+        self.embedder.reset_parameters()
+        self.embedding_aggregator.reset_parameters()
+        self.encoder.reset_parameters()
+        self.output_normalization.reset_parameters()
+
+    def forward(
+        self,
+        feature_tensors: TensorMap,
+        padding_mask: torch.BoolTensor,
+    ) -> torch.Tensor:
+        """
+        :param feature_tensors: a dictionary of tensors to generate embeddings.
+        :param padding_mask: A mask of shape ``(batch_size, sequence_length)``
+            indicating which elements within ``key`` to ignore for the purpose of attention (i.e. treat as "padding").
+            ``False`` value indicates that the corresponding ``key`` value will be ignored.
+        :returns: The final hidden state.\n
+            Expected shape: ``(batch_size, sequence_length, embedding_dim)``
+        """
+        embeddings = self.embedder(feature_tensors)
+        agg_emb: torch.Tensor = self.embedding_aggregator(embeddings)
+        assert agg_emb.dim() == 3
+
+        attn_mask = self.attn_mask_builder(feature_tensors, padding_mask)
+
+        hidden_state: torch.Tensor = self.encoder(
+            feature_tensors=feature_tensors,
+            input_embeddings=agg_emb,
+            padding_mask=padding_mask,
+            attention_mask=attn_mask,
+        )
+        assert agg_emb.size() == hidden_state.size()
+
+        hidden_state = self.output_normalization(hidden_state)
+        return hidden_state
+
+
+class SasRec(torch.nn.Module):
+    """
+    A model using the SasRec architecture as a hidden state generator.
+    The hidden states are multiplied by the item embeddings,
+    resulting in logits for each of the items.
+
+    Source paper: https://arxiv.org/pdf/1808.09781.
+
+    Example:
+
+    .. code-block:: python
+
+        from replay.data import FeatureHint, FeatureSource, FeatureType
+        from replay.data.nn import TensorFeatureInfo, TensorFeatureSource, TensorSchema
+        from replay.nn.agg import SumAggregator
+        from replay.nn.embedding import SequenceEmbedding
+        from replay.nn.mask import DefaultAttentionMask
+        from replay.nn.loss import CESampled
+        from replay.nn.sequential import PositionAwareAggregator, SasRecTransformerLayer
+
+        tensor_schema = TensorSchema(
+            [
+                TensorFeatureInfo(
+                    "item_id",
+                    is_seq=True,
+                    feature_type=FeatureType.CATEGORICAL,
+                    embedding_dim=256,
+                    padding_value=NUM_UNIQUE_ITEMS,
+                    cardinality=NUM_UNIQUE_ITEMS+1,
+                    feature_hint=FeatureHint.ITEM_ID,
+                    feature_sources=[TensorFeatureSource(FeatureSource.INTERACTIONS, "item_id")]
+                ),
+            ]
+        )
+
+        body = SasRecBody(
+            embedder=SequenceEmbedding(
+                schema=tensor_schema,
+            ),
+            embedding_aggregator=PositionAwareAggregator(
+                embedding_aggregator=SumAggregator(embedding_dim=256),
+                max_sequence_length=100,
+                dropout=0.2,
+            ),
+            attn_mask_builder=DefaultAttentionMask(
+                reference_feature_name=tensor_schema.item_id_feature_name,
+                num_heads=2,
+            ),
+            encoder=SasRecTransformerLayer(
+                embedding_dim=256,
+                num_heads=2,
+                num_blocks=2,
+                dropout=0.3,
+                activation="relu",
+            ),
+            output_normalization=torch.nn.LayerNorm(256),
+        )
+        sasrec = SasRec(
+            body=body,
+            loss=CESampled(padding_idx=tensor_schema.item_id_features.item().padding_value)
+        )
+
+    """
+
+    def __init__(
+        self,
+        body: SasRecBody,
+        loss: LossProto,
+    ):
+        """
+        :param body: An instance of SasRecBody.
+        :param loss: An object of a class that performs loss calculation
+            based on hidden states from the model, positive and optionally negative labels.
+        """
+        super().__init__()
+        self.body = body
+        self.head = EmbeddingTyingHead()
+        self.loss = loss
+        self.loss.logits_callback = self.get_logits
+
+        self.reset_parameters()
+
+    @classmethod
+    def from_params(
+        cls,
+        schema: TensorSchema,
+        embedding_dim: int = 192,
+        num_heads: int = 4,
+        num_blocks: int = 2,
+        max_sequence_length: int = 50,
+        dropout: float = 0.3,
+        excluded_features: Optional[list[str]] = None,
+        categorical_list_feature_aggregation_method: Literal["sum", "mean", "max"] = "sum",
+    ) -> "SasRec":
+        from replay.nn.agg import SumAggregator
+        from replay.nn.embedding import SequenceEmbedding
+        from replay.nn.loss import CE
+        from replay.nn.mask import DefaultAttentionMask
+
+        from .agg import PositionAwareAggregator
+        from .transformer import SasRecTransformerLayer
+
+        excluded_features = [
+            schema.query_id_feature_name,
+            schema.timestamp_feature_name,
+            *(excluded_features or []),
+        ]
+        excluded_features = list(set(excluded_features))
+
+        body = SasRecBody(
+            embedder=SequenceEmbedding(
+                schema=schema,
+                categorical_list_feature_aggregation_method=categorical_list_feature_aggregation_method,
+                excluded_features=excluded_features,
+            ),
+            embedding_aggregator=PositionAwareAggregator(
+                embedding_aggregator=SumAggregator(embedding_dim=embedding_dim),
+                max_sequence_length=max_sequence_length,
+                dropout=dropout,
+            ),
+            attn_mask_builder=DefaultAttentionMask(
+                reference_feature_name=schema.item_id_feature_name,
+                num_heads=num_heads,
+            ),
+            encoder=SasRecTransformerLayer(
+                embedding_dim=embedding_dim,
+                num_heads=num_heads,
+                num_blocks=num_blocks,
+                dropout=dropout,
+                activation="relu",
+            ),
+            output_normalization=torch.nn.LayerNorm(embedding_dim),
+        )
+        return cls(
+            body=body,
+            loss=CE(ignore_index=schema.item_id_features.item().padding_value),
+        )
+
+    def reset_parameters(self) -> None:
+        self.body.reset_parameters()
+
+    def get_logits(
+        self,
+        model_embeddings: torch.Tensor,
+        candidates_to_score: Optional[torch.LongTensor] = None,
+    ) -> torch.Tensor:
+        item_embeddings: torch.Tensor = self.body.embedder.get_item_weights(candidates_to_score)
+        logits: torch.Tensor = self.head(model_embeddings, item_embeddings)
+        return logits
+
+    def forward_train(
+        self,
+        feature_tensors: TensorMap,
+        padding_mask: torch.BoolTensor,
+        positive_labels: torch.LongTensor,
+        negative_labels: torch.LongTensor,
+        target_padding_mask: torch.BoolTensor,
+    ) -> TrainOutput:
+        hidden_states: torch.Tensor = self.body(feature_tensors, padding_mask)
+        assert hidden_states.dim() == 3
+
+        loss: torch.Tensor = self.loss(
+            model_embeddings=hidden_states,
+            feature_tensors=feature_tensors,
+            positive_labels=positive_labels,
+            negative_labels=negative_labels,
+            padding_mask=padding_mask,
+            target_padding_mask=target_padding_mask,
+        )
+
+        return {
+            "loss": loss,
+            "hidden_states": (hidden_states,),
+        }
+
+    def forward_inference(
+        self,
+        feature_tensors: TensorMap,
+        padding_mask: torch.BoolTensor,
+        candidates_to_score: Optional[torch.LongTensor] = None,
+    ) -> InferenceOutput:
+        hidden_states: torch.Tensor = self.body(feature_tensors, padding_mask)
+        assert hidden_states.dim() == 3
+
+        last_hidden_state = hidden_states[:, -1, :].contiguous()
+        logits = self.get_logits(last_hidden_state, candidates_to_score)
+
+        return {
+            "logits": logits,
+            "hidden_states": (hidden_states,),
+        }
+
+    def forward(
+        self,
+        feature_tensors: TensorMap,
+        padding_mask: torch.BoolTensor,
+        candidates_to_score: Optional[torch.LongTensor] = None,
+        positive_labels: Optional[torch.LongTensor] = None,
+        negative_labels: Optional[torch.LongTensor] = None,
+        target_padding_mask: Optional[torch.BoolTensor] = None,
+    ) -> Union[TrainOutput, InferenceOutput]:
+        """
+        :param feature_tensors: a dictionary of tensors to generate embeddings.
+        :param padding_mask: A mask of shape ``(batch_size, sequence_length)``
+            indicating which elements within ``key`` to ignore for the purpose of attention (i.e. treat as "padding").
+            ``False`` value indicates that the corresponding ``key`` value will be ignored.
+        :param candidates_to_score: a tensor containing item IDs
+            for which you need to get logits at the inference stage.\n
+            **Note:** you must take into account the padding value when creating the tensor.\n
+            The tensor participates in calculations only on the inference stage.
+            You don't have to submit an argument at training stage,
+            but if it is submitted, then no effect will be provided.\n
+            Default: ``None``.
+        :param positive_labels: a tensor containing positive labels for calculating the loss.\n
+            You don't have to submit an argument at inference stage,
+            but if it is submitted, then no effect will be provided.\n
+            Default: ``None``.
+        :param negative_labels: a tensor containing negative labels for calculating the loss.\n
+            **Note:** Before run make sure that your loss supports calculations with negative labels.\n
+            You don't have to submit an argument at inference stage,
+            but if it is submitted, then no effect will be provided.\n
+            Default: ``None``.
+        :param target_padding_mask: A mask of shape ``(batch_size, sequence_length, num_positives)``
+            indicating elements from ``positive_labels`` to ignore during loss calculation.
+            ``False`` value indicates that the corresponding value will be ignored.\n
+            You don't have to submit an argument at inference stage,
+            but if it is submitted, then no effect will be provided.\n
+            Default: ``None``.
+        :returns: During training, the model will return an object
+            of the ``TrainOutput`` container class.
+            At the inference stage, the ``InferenceOutput`` class will be returned.
+        """
+        if self.training:
+            all(
+                map(
+                    warning_is_not_none("Variable `{}` is not None. This will have no effect at the training stage."),
+                    [(candidates_to_score, "candidates_to_score")],
+                )
+            )
+            return self.forward_train(
+                feature_tensors=feature_tensors,
+                padding_mask=padding_mask,
+                positive_labels=positive_labels,
+                negative_labels=negative_labels,
+                target_padding_mask=target_padding_mask,
+            )
+
+        all(
+            map(
+                warning_is_not_none("Variable `{}` is not None. This will have no effect at the inference stage."),
+                [
+                    (positive_labels, "positive_labels"),
+                    (negative_labels, "negative_labels"),
+                    (target_padding_mask, "target_padding_mask"),
+                ],
+            )
+        )
+        return self.forward_inference(
+            feature_tensors=feature_tensors,
+            padding_mask=padding_mask,
+            candidates_to_score=candidates_to_score,
+        )