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

replay/nn/loss/login_ce.py

@@ -0,0 +1,373 @@
+from typing import Callable, Optional, TypedDict
+
+import torch
+
+from replay.data.nn import TensorMap
+
+from .base import SampledLossBase, mask_negative_logits
+
+
+class LogInCESampledOutput(TypedDict):
+    positive_logits: torch.Tensor
+    negative_logits: torch.Tensor
+    positive_labels: torch.LongTensor
+    negative_labels: torch.LongTensor
+    target_padding_mask: torch.BoolTensor
+
+
+class LogInCEBase(SampledLossBase):
+    def get_sampled_logits(
+        self,
+        model_embeddings: torch.Tensor,
+        positive_labels: torch.LongTensor,  # [batch_size, seq_len, num_positives]
+        negative_labels: torch.LongTensor,  # [num_negatives] or [batch_size, seq_len, num_negatives]
+        target_padding_mask: torch.BoolTensor,  # [batch_size, seq_len, num_positives]
+    ) -> LogInCESampledOutput:
+        """
+        The function of calculating positive and negative logits in LogInCE losses.
+        Based on the embeddings from the model, positive and negative labels.
+
+        The function supports the calculation of logits for the case of multi-positive labels
+        (there are several labels for each position in the sequence).
+
+        :param model_embeddings: Embeddings from the model. This is usually the last hidden state.
+            Expected shape: ``(batch_size, sequence_length, embedding_dim)``
+        :param positive_labels: a tensor containing labels with positive events.
+            Expected shape: ``(batch_size, sequence_length, num_positives)``
+        :param negative_labels: a tensor containing labels with negative events.
+            Expected shape:
+                - ``(batch_size, sequence_length, num_negatives)``
+                - ``(batch_size, num_negatives)``
+                - ``(num_negatives)`` - a case where the same negative events are used for the entire batch.
+        :param target_padding_mask: Padding mask for ``positive_labels`` (targets).
+            ``False`` value indicates that the corresponding ``key`` value will be ignored.
+            Expected shape: ``(batch_size, sequence_length, num_positives)``
+
+        :returns: LogInCESampledOutput. A dictionary containing positive and negative logits with labels.
+        """
+        ################## SHAPE CHECKING STAGE START ##################
+        batch_size, seq_len, num_positives = positive_labels.size()
+        assert target_padding_mask.size() == (batch_size, seq_len, num_positives)
+        num_negatives = negative_labels.size(-1)
+
+        if negative_labels.size() == (batch_size, num_negatives):
+            # [batch_size, num_negatives] -> [batch_size, 1, num_negatives]
+            negative_labels = negative_labels.unsqueeze(1).repeat(1, seq_len, 1)
+
+        assert negative_labels.size() == (batch_size, seq_len, num_negatives) or negative_labels.dim() == 1
+        ################## SHAPE CHECKING STAGE END ##################
+
+        # Get output embedding for every user event
+        embedding_dim = model_embeddings.size(-1)
+        assert model_embeddings.size() == (batch_size, seq_len, embedding_dim)
+
+        # [batch_size, seq_len, num_positives] -> [batch_size, seq_len]
+        masked_target_padding_mask: torch.BoolTensor = target_padding_mask.sum(-1).bool()
+        masked_batch_size = masked_target_padding_mask.sum().item()
+
+        # Apply target mask
+        # [batch_size, seq_len, emb_dim] -> [masked_batch_size, emb_dim]
+        model_embeddings = model_embeddings[masked_target_padding_mask]
+        assert model_embeddings.size() == (masked_batch_size, embedding_dim)
+
+        # [batch_size, seq_len, num_positives] -> [masked_batch_size, num_positives]
+        positive_labels = positive_labels[masked_target_padding_mask]
+        assert positive_labels.size() == (masked_batch_size, num_positives)
+
+        if negative_labels.dim() > 1:  # pragma: no cover
+            # [batch_size, seq_len, num_negatives] -> [masked_batch_size, num_negatives]
+            negative_labels = negative_labels[masked_target_padding_mask]
+            assert negative_labels.size() == (masked_batch_size, num_negatives)
+
+        positive_logits = self.logits_callback(model_embeddings, positive_labels)
+        assert positive_logits.size() == (masked_batch_size, num_positives)
+
+        negative_logits = self.logits_callback(model_embeddings, negative_labels)
+        assert negative_logits.size() == (masked_batch_size, num_negatives)
+
+        # [batch_size, seq_len, num_positives] -> [masked_batch_size, num_positives]
+        target_padding_mask = target_padding_mask[masked_target_padding_mask]
+        assert target_padding_mask.size() == (masked_batch_size, num_positives)
+
+        return {
+            "positive_logits": positive_logits,
+            "negative_logits": negative_logits,
+            "positive_labels": positive_labels,
+            "negative_labels": negative_labels,
+            "target_padding_mask": target_padding_mask,
+        }
+
+
+class LogInCE(LogInCEBase):
+    """
+    LogInCE loss.
+
+    .. math::
+
+        L_{\\text{InfoNCE}} = -\\log \\frac{\\sum_{p \\in P}
+        \\exp(\\mathrm{sim}(q, p))}{\\sum_{p \\in P}
+        \\exp(\\mathrm{sim}(q, p)) + \\sum_{n \\in N} \\exp(\\mathrm{sim}(q, n))},
+
+    where q -- query embedding, P -- set of positive logits, N -- set of negative logits,
+    :math:`sim(\\cdot, \\cdot)` -- similaruty function.
+
+    The loss supports the calculation of logits for the case of multi-positive labels
+    (there are several labels for each position in the sequence).
+    """
+
+    def __init__(
+        self,
+        cardinality: int,
+        log_epsilon: float = 1e-6,
+        clamp_border: float = 100.0,
+        negative_labels_ignore_index: int = -100,
+    ):
+        """
+        :param cardinality: number of unique items in vocabulary (catalog).
+            The specified cardinality value must not take into account the padding value.
+        :param log_epsilon: correction to avoid zero in the logarithm during loss calculating.
+            Default: ``1e-6``.
+        :param clamp_border: upper bound for clamping loss tensor, lower bound will be setted to ``-clamp_border``.
+            Default: ``100.0``.
+        :param negative_labels_ignore_index: padding value for negative labels.
+            This may be the case when negative labels
+            are formed at the preprocessing level, rather than the negative sampler.
+            The index is ignored and does not contribute to the loss.
+            Default: ``-100``.
+        """
+        super().__init__()
+        self.cardinality = cardinality
+        self.log_epsilon = log_epsilon
+        self.clamp_border = clamp_border
+        self.negative_labels_ignore_index = negative_labels_ignore_index
+        self._logits_callback = None
+
+    @property
+    def logits_callback(
+        self,
+    ) -> Callable[[torch.Tensor, Optional[torch.Tensor]], torch.Tensor]:
+        """
+        Property for calling a function for the logits computation.\n
+
+        This function is expected to receive model's last hidden state
+                    and optionally item IDs, and return a logits tensor.
+
+        It is expected that the corresponding head model method will be used as this function,
+        for example, the ``get_logits`` method of the ``SasRec`` class.
+
+        :return: callable function.
+        """
+        if self._logits_callback is None:
+            msg = "The callback for getting logits is not defined"
+            raise AttributeError(msg)
+        return self._logits_callback
+
+    @logits_callback.setter
+    def logits_callback(self, func: Optional[Callable]) -> None:
+        self._logits_callback = func
+
+    def forward(
+        self,
+        model_embeddings: torch.Tensor,
+        feature_tensors: TensorMap,  # noqa: ARG002
+        positive_labels: torch.LongTensor,
+        negative_labels: torch.LongTensor,  # noqa: ARG002
+        padding_mask: torch.BoolTensor,  # noqa: ARG002
+        target_padding_mask: torch.BoolTensor,
+    ) -> torch.Tensor:
+        """
+        forward(model_embeddings, positive_labels, target_padding_mask)
+        **Note**: At forward pass, the whole catalog of items is used as negatives.
+        Next, negative logits, corresponding to positions where negative labels
+        coincide with positive ones, are masked.
+
+        :param model_embeddings: model output of shape ``(batch_size, sequence_length, embedding_dim)``.
+        :param positive_labels: ground truth labels of positive events
+            of shape (batch_size, sequence_length, num_positives).
+        :param target_padding_mask: padding mask corresponding for ``positive_labels``
+            of shape (batch_size, sequence_length, num_positives).
+        :return: computed loss value.
+        """
+        all_negative_labels = torch.arange(
+            self.cardinality,
+            dtype=torch.long,
+            device=positive_labels.device,
+        )
+        sampled = self.get_sampled_logits(
+            model_embeddings,
+            positive_labels,
+            all_negative_labels,
+            target_padding_mask,
+        )
+        positive_logits = sampled["positive_logits"]  # [masked_batch_size, num_positives]
+        negative_logits = sampled["negative_logits"]  # [masked_batch_size, num_negatives]
+        positive_labels = sampled["positive_labels"]  # [masked_batch_size, num_positives]
+        all_negative_labels = sampled["negative_labels"]  # [masked_batch_size, num_negatives] or [num_negatives]
+        target_padding_mask = sampled["target_padding_mask"]  # [masked_batch_size, num_positives]
+
+        # [masked_batch_size, num_negatives] - assign low values to some negative logits
+        negative_logits = mask_negative_logits(
+            negative_logits,
+            all_negative_labels,
+            positive_labels,
+            self.negative_labels_ignore_index,
+        )
+
+        max_values = torch.max(
+            positive_logits.max(-1, keepdim=True).values,
+            negative_logits.max(-1, keepdim=True).values,
+        )  # [masked_batch_size, 1]
+        positive_logits = positive_logits - max_values
+        negative_logits = negative_logits - max_values
+
+        positive_logits = torch.exp(positive_logits)
+        positive_logits = positive_logits * target_padding_mask
+        # [masked_batch_size, num_positives] -> [masked_batch_size]
+        positive_logits = positive_logits.sum(-1)
+
+        negative_logits = torch.exp(negative_logits)
+        # [masked_batch_size, num_negatives] -> [masked_batch_size]
+        negative_logits = negative_logits.sum(-1)
+
+        probabilities = positive_logits / (positive_logits + negative_logits)
+        loss = -torch.clamp(
+            torch.log(probabilities + self.log_epsilon),
+            -self.clamp_border,
+            self.clamp_border,
+        )
+        return loss.mean()
+
+
+class LogInCESampled(LogInCEBase):
+    """
+    Sampled version of LogInCE (Log InfoNCE) loss (with negative sampling items).
+
+    .. math::
+
+        L_{\\text{InfoNCE}} = -\\log \\frac{\\sum_{p \\in P} \\exp(\\mathrm{sim}(q, p))}{\\sum_{p \\in P}
+        \\exp(\\mathrm{sim}(q, p)) + \\sum_{n \\in N_{\\text{sampled}}} \\exp(\\mathrm{sim}(q, n))},
+
+    where q -- query embedding, P -- set of positive logits, :math:`N_sampled` -- set of negative logits,
+    :math:`sim(\\cdot, \\cdot)` -- similaruty function.\n
+    Same as ``LogInCE``, the difference in the set of negatives.
+
+    The loss supports the calculation of logits for the case of multi-positive labels
+    (there are several labels for each position in the sequence).
+    """
+
+    def __init__(
+        self,
+        log_epsilon: float = 1e-6,
+        clamp_border: float = 100.0,
+        negative_labels_ignore_index: int = -100,
+    ):
+        """
+        :param log_epsilon: correction to avoid zero in the logarithm during loss calculating.
+            Default: 1e-6.
+        :param clamp_border: upper bound for clamping loss tensor, lower bound will be setted to -`clamp_border`.
+            Default: 100.0.
+        :param negative_labels_ignore_index: padding value for negative labels.
+            This may be the case when negative labels
+            are formed at the preprocessing level, rather than the negative sampler.
+            The index is ignored and does not contribute to the loss.
+            Default: ``-100``.
+        """
+        super().__init__()
+        self.log_epsilon = log_epsilon
+        self.clamp_border = clamp_border
+        self.negative_labels_ignore_index = negative_labels_ignore_index
+        self._logits_callback = None
+
+    @property
+    def logits_callback(
+        self,
+    ) -> Callable[[torch.Tensor, Optional[torch.Tensor]], torch.Tensor]:
+        """
+        Property for calling a function for the logits computation.\n
+
+        This function is expected to receive model's last hidden state
+                    and optionally item IDs, and return a logits tensor.
+
+        It is expected that the corresponding head model method will be used as this function,
+        for example, the ``get_logits`` method of the ``SasRec`` class.
+
+        :return: callable function.
+        """
+        if self._logits_callback is None:
+            msg = "The callback for getting logits is not defined"
+            raise AttributeError(msg)
+        return self._logits_callback
+
+    @logits_callback.setter
+    def logits_callback(self, func: Optional[Callable]) -> None:
+        self._logits_callback = func
+
+    def forward(
+        self,
+        model_embeddings: torch.Tensor,
+        feature_tensors: TensorMap,  # noqa: ARG002
+        positive_labels: torch.LongTensor,
+        negative_labels: torch.LongTensor,
+        padding_mask: torch.BoolTensor,  # noqa: ARG002
+        target_padding_mask: torch.BoolTensor,
+    ) -> torch.Tensor:
+        """
+        forward(model_embeddings, positive_labels, negative_labels, target_padding_mask)
+
+        :param model_embeddings: model output of shape ``(batch_size, sequence_length, embedding_dim)``.
+        :param positive_labels: labels of positive events
+            of shape ``(batch_size, sequence_length, num_positives)``.
+        :param negative_labels: labels of sampled negative events.
+
+            Expected shape:
+                    - ``(batch_size, sequence_length, num_negatives)``
+                    - ``(batch_size, num_negatives)``
+                    - ``(num_negatives)`` - a case where the same negative events are used for the entire batch.
+        :param target_padding_mask: padding mask corresponding for ``positive_labels``
+            of shape ``(batch_size, sequence_length, num_positives)``
+
+        :return: computed loss value.
+        """
+        sampled = self.get_sampled_logits(
+            model_embeddings,
+            positive_labels,
+            negative_labels,
+            target_padding_mask,
+        )
+        positive_logits = sampled["positive_logits"]  # [masked_batch_size, num_positives]
+        negative_logits = sampled["negative_logits"]  # [masked_batch_size, num_negatives]
+        positive_labels = sampled["positive_labels"]  # [masked_batch_size, num_positives]
+        negative_labels = sampled["negative_labels"]  # [masked_batch_size, num_negatives] or [num_negatives]
+        target_padding_mask = sampled["target_padding_mask"]  # [masked_batch_size, num_positives]
+
+        # [masked_batch_size, num_negatives] - assign low values to some negative logits
+        negative_logits = mask_negative_logits(
+            negative_logits,
+            negative_labels,
+            positive_labels,
+            self.negative_labels_ignore_index,
+        )
+
+        max_values = torch.max(
+            positive_logits.max(-1, keepdim=True).values,
+            negative_logits.max(-1, keepdim=True).values,
+        )  # [masked_batch_size, 1]
+        positive_logits = positive_logits - max_values
+        negative_logits = negative_logits - max_values
+
+        positive_logits = torch.exp(positive_logits)
+        positive_logits = positive_logits * target_padding_mask
+        # [masked_batch_size, num_positives] -> [masked_batch_size]
+        positive_logits = positive_logits.sum(-1)
+
+        negative_logits = torch.exp(negative_logits)
+        # [masked_batch_size, num_negatives] -> [masked_batch_size]
+        negative_logits = negative_logits.sum(-1)
+
+        probabilities = positive_logits / (positive_logits + negative_logits)
+        loss = -torch.clamp(
+            torch.log(probabilities + self.log_epsilon),
+            -self.clamp_border,
+            self.clamp_border,
+        )
+        return loss.mean()

replay/nn/loss/logout_ce.py

@@ -0,0 +1,230 @@
+from typing import Callable, Optional
+
+import torch
+
+from replay.data.nn import TensorMap
+
+from .base import mask_negative_logits
+
+
+class LogOutCE(torch.nn.Module):
+    """
+    LogOutCE loss.
+
+        .. math::
+
+            L_{\\text{InfoNCE}} = - \\sum_{p \\in P} \\log \\frac{ \\exp(\\mathrm{sim}(q, p))}
+            {\\exp(\\mathrm{sim}(q, p))
+            + \\sum_{n \\in N} \\exp(\\mathrm{sim}(q, n))}.
+
+    where q -- query embedding, P -- set of positive logits, N -- set of negative logits,
+    :math:`sim(\\cdot, \\cdot)` -- similaruty function.\n
+
+    The loss supports the calculation of logits for the case of multi-positive labels
+    (there are several labels for each position in the sequence).
+    """
+
+    def __init__(
+        self,
+        cardinality: int,
+        negative_labels_ignore_index: int = -100,
+        **kwargs,
+    ):
+        """
+        To calculate the loss, ``torch.nn.CrossEntropyLoss`` is used.
+        You can pass all parameters for initializing the object via kwargs.
+
+        :param cardinality: number of unique items in vocabulary (catalog).
+            The specified cardinality value must not take into account the padding value.
+        :param negative_labels_ignore_index: padding value for negative labels.
+            This may be the case when negative labels
+            are formed at the preprocessing level, rather than the negative sampler.
+            The index is ignored and does not contribute to the loss.
+            Default: ``-100``.
+        """
+        super().__init__()
+        self.cardinality = cardinality
+        self.negative_labels_ignore_index = negative_labels_ignore_index
+        self._loss = torch.nn.CrossEntropyLoss(**kwargs)
+        self._logits_callback = None
+
+    @property
+    def logits_callback(
+        self,
+    ) -> Callable[[torch.Tensor, Optional[torch.Tensor]], torch.Tensor]:
+        """
+        Property for calling a function for the logits computation.\n
+
+        This function is expected to receive model's last hidden state
+                    and optionally item IDs, and return a logits tensor.
+
+        It is expected that the corresponding head model method will be used as this function,
+        for example, the ``get_logits`` method of the ``SasRec`` class.
+
+        :return: callable function.
+        """
+        if self._logits_callback is None:
+            msg = "The callback for getting logits is not defined"
+            raise AttributeError(msg)
+        return self._logits_callback
+
+    @logits_callback.setter
+    def logits_callback(self, func: Optional[Callable]) -> None:
+        self._logits_callback = func
+
+    def forward(
+        self,
+        model_embeddings: torch.Tensor,
+        feature_tensors: TensorMap,  # noqa: ARG002
+        positive_labels: torch.LongTensor,
+        negative_labels: torch.LongTensor,  # noqa: ARG002
+        padding_mask: torch.BoolTensor,  # noqa: ARG002
+        target_padding_mask: torch.BoolTensor,
+    ) -> torch.Tensor:
+        """
+        forward(model_embeddings, positive_labels, target_padding_mask)
+        **Note**: At forward pass, the whole catalog of items is used as negatives.
+        Next, negative logits, corresponding to positions where negative labels
+        coincide with positive ones, are masked.
+
+        :param model_embeddings: model output of shape ``(batch_size, sequence_length, embedding_dim)``.
+        :param positive_labels: ground truth labels of positive events
+            of shape (batch_size, sequence_length, num_positives).
+        :param target_padding_mask: padding mask corresponding for ``positive_labels``
+            of shape (batch_size, sequence_length, num_positives).
+        :return: computed loss value.
+        """
+        initial_target_padding_mask = target_padding_mask
+        num_positives = target_padding_mask.size(2)
+        # [batch_size, seq_len, num_positives] -> [batch_size * seq_len]
+        if num_positives == 1:
+            target_padding_mask = target_padding_mask.squeeze(-1)
+        else:
+            target_padding_mask = target_padding_mask.sum(-1).bool()
+        masked_batch_size = target_padding_mask.sum().item()
+
+        logits: torch.Tensor = self.logits_callback(model_embeddings)  # [batch_size, seq_len, vocab_size]
+        all_negative_labels = torch.arange(self.cardinality, dtype=torch.long, device=positive_labels.device)
+
+        # [batch_size, seq_len, vocab_size] -> [masked_batch_size, vocab_size]
+        logits = logits[target_padding_mask]
+
+        # [batch_size, seq_len, num_positives] -> [masked_batch_size, num_positives]
+        positive_labels = positive_labels[target_padding_mask]
+
+        # [batch_size, seq_len, num_positives] -> [masked_batch_size, num_positives]
+        target_padding_mask = initial_target_padding_mask[target_padding_mask]
+
+        positive_ids = torch.arange(masked_batch_size, dtype=torch.long, device=positive_labels.device)
+        # [masked_batch_size, vocab_size] -> [masked_batch_size, num_positives]
+        positive_logits = logits[positive_ids.unsqueeze(-1), positive_labels]
+
+        # [masked_batch_size, vocab_size] - assign low values to some negative logits
+        negative_logits = mask_negative_logits(
+            logits,
+            all_negative_labels,
+            positive_labels,
+            self.negative_labels_ignore_index,
+        )
+
+        # [masked_batch_size, num_negatives] -> [masked_batch_size, 1, num_negatives]
+        negative_logits = negative_logits.unsqueeze(-2)
+        # [masked_batch_size, 1, num_negatives] -> [masked_batch_size, num_positives, num_negatives]
+        negative_logits = negative_logits.repeat(1, target_padding_mask.size(-1), 1)
+        # [masked_batch_size, num_positives, num_negatives] -> [masked_batch_size, num_negatives]
+        negative_logits = negative_logits[target_padding_mask]
+        # [masked_batch_size, num_positives] -> [masked_batch_size]
+        positive_logits = positive_logits[target_padding_mask]
+        # [masked_batch_size] -> [masked_batch_size, 1]
+        positive_logits = positive_logits.unsqueeze(-1)
+
+        # [masked_batch_size, 1 + num_negatives] - all logits
+        logits = torch.cat((positive_logits, negative_logits), dim=-1)
+        # [masked_batch_size] - positives are always at 0 position for all recommendation points
+        target = torch.zeros(logits.size(0), dtype=torch.long, device=positive_labels.device)
+        # [masked_batch_size] - loss for all recommendation points
+        loss = self._loss(logits, target)
+        return loss
+
+
+class LogOutCEWeighted(LogOutCE):
+    """
+    LogOutCE loss with sample weights enabling.
+
+        .. math::
+
+            L_{\\text{InfoNCE}} = - \\sum_{p \\in P} \\log \\frac{ \\exp(\\mathrm{sim}(q, p))}
+            {\\exp(\\mathrm{sim}(q, p))
+            + \\sum_{n \\in N} \\exp(\\mathrm{sim}(q, n))}.
+
+    where q -- query embedding, P -- set of positive logits, N -- set of negative logits,
+    :math:`sim(\\cdot, \\cdot)` -- similaruty function.\n
+
+    In addition to calculating the standard loss,
+    weights are applied for each sample.
+    Therefore, it is expected that the sample weights will be in the generated batch,
+    which is fed into the model.
+
+    The loss supports the calculation of logits for the case of multi-positive labels
+    (there are several labels for each position in the sequence).
+    """
+
+    def __init__(
+        self,
+        cardinality: int,
+        feature_name: str,
+        negative_labels_ignore_index: int = -100,
+        **kwargs,
+    ):
+        """
+        To calculate the loss, ``torch.nn.CrossEntropyLoss`` is used with the parameter ``reduction="none"``.
+        You can pass all other parameters for initializing the object via kwargs.
+
+        :param cardinality: number of unique items in vocabulary (catalog).
+            The specified cardinality value must not take into account the padding value.
+        :param feature_name: the name of the key in the batch.
+            The tensor is expected to contain sample weights.
+        :param negative_labels_ignore_index: padding value for negative labels.
+            This may be the case when negative labels
+            are formed at the preprocessing level, rather than the negative sampler.
+            The index is ignored and does not contribute to the loss.
+            Default: ``-100``.
+        """
+        super().__init__(
+            cardinality=cardinality,
+            negative_labels_ignore_index=negative_labels_ignore_index,
+        )
+        self.feature_name = feature_name
+        self._loss = torch.nn.CrossEntropyLoss(reduction="none", **kwargs)
+
+    def forward(
+        self,
+        model_embeddings: torch.Tensor,
+        feature_tensors: TensorMap,
+        positive_labels: torch.LongTensor,
+        negative_labels: torch.LongTensor,  # noqa: ARG002
+        padding_mask: torch.BoolTensor,  # noqa: ARG002
+        target_padding_mask: torch.BoolTensor,
+    ) -> torch.Tensor:
+        """
+        forward(model_embeddings, feature_tensors, positive_labels, target_padding_mask)
+        **Note**: At forward pass, the whole catalog of items is used as negatives.
+        Next, negative logits, corresponding to positions where negative labels
+        coincide with positive ones, are masked.
+
+        :param model_embeddings: model output of shape ``(batch_size, sequence_length, embedding_dim)``.
+        :param feature_tensors: a dictionary of tensors from dataloader.
+            This dictionary is expected to contain a key with the name ``feature_name``,
+            which is specified in the constructor.
+            Expected shape of tensor ``(batch_size, sequence_length, num_positives)``.
+        :param positive_labels: ground truth labels of positive events
+            of shape (batch_size, sequence_length, num_positives).
+        :param target_padding_mask: padding mask corresponding for ``positive_labels``
+            of shape (batch_size, sequence_length, num_positives).
+        :return: computed loss value.
+        """
+        loss: torch.Tensor = super().forward(model_embeddings, None, positive_labels, None, None, target_padding_mask)
+        sample_weight = feature_tensors[self.feature_name]
+        sample_weight = sample_weight[target_padding_mask]
+        loss = (loss * sample_weight).mean()
+        return loss

replay/nn/mask.py

@@ -0,0 +1,87 @@
+from abc import ABC, abstractmethod
+from typing import Protocol
+
+import torch
+
+from replay.data.nn.schema import TensorMap
+
+
+class AttentionMaskProto(Protocol):
+    def __call__(self, feature_tensor: TensorMap, padding_mask: torch.BoolTensor) -> torch.Tensor: ...
+
+
+class AttentionMaskBase(ABC):
+    def __init__(
+        self,
+        num_heads: int,
+    ) -> None:
+        self.num_heads = num_heads
+
+    def __call__(
+        self,
+        feature_tensor: TensorMap,
+        padding_mask: torch.BoolTensor,
+    ) -> torch.FloatTensor:
+        """
+        :param feature_tensor: dict of features tensors.
+        :param padding_mask: Padding mask where ``0`` - ``<PAD>``, ``1`` - otherwise.
+        :returns: Float attention mask of shape ``(B * num_heads, L, L)``,
+            where ``-inf`` for ``<PAD>``, ``0`` - otherwise.
+        """
+        attention_mask = self._get_attention_mask(feature_tensor)
+
+        diagonal_attention_mask = torch.diag(
+            torch.ones(padding_mask.size(1), dtype=torch.bool, device=padding_mask.device)
+        )
+        # (B, L) -> (B, 1, 1, L)
+        key_padding_mask = padding_mask.unsqueeze(1).unsqueeze(1)
+        # (B, 1, 1, L) -> (B, 1, L, L), where 0 - PAD, 1 - otherwise
+        key_padding_mask = key_padding_mask | diagonal_attention_mask
+
+        attention_mask = (attention_mask & key_padding_mask).float()
+        attention_mask = attention_mask.masked_fill(attention_mask == 0, float("-inf")).masked_fill(
+            attention_mask == 1, 0.0
+        )
+        if attention_mask.size(1) != self.num_heads and attention_mask.shape[1] == 1:
+            # for default attention_mask of shape (L, L) it becomes (B, 1, L, L)
+            # (B, 1, L, L) -> (B, num_heads, L, L)
+            attention_mask = attention_mask.repeat(1, self.num_heads, 1, 1)
+        # (B, num_heads, L, L) -> (B * num_heads, L, L)
+        attention_mask = attention_mask.reshape(-1, *attention_mask.shape[-2:])
+        return attention_mask
+
+    @abstractmethod
+    def _get_attention_mask(self, feature_tensor: TensorMap) -> torch.Tensor:
+        raise NotImplementedError()  # pragma: no cover
+
+
+class DefaultAttentionMask(AttentionMaskBase):
+    """
+    Constructs a float lower-triangular attenstion mask
+    of shape ``(batch_size * num_heads, sequence_length, sequence_length)``,
+    where ``-inf`` for ``<PAD>``, ``0`` - otherwise.
+    """
+
+    def __init__(
+        self,
+        reference_feature_name: str,
+        num_heads: int,
+    ) -> None:
+        """
+        :param reference_feature_name: To build a mask, you need a reference tensor.
+            So you need to pass the name of the tensor, which will definitely be in the dictionary of feature tensors.
+            The second dimension (1 in zero indexing) of the tensor will be used to construct the attention mask.
+        :param num_heads: Number of attention heads.
+        """
+        super().__init__(num_heads)
+        self._feature_name = reference_feature_name
+
+    def _get_attention_mask(self, feature_tensor: TensorMap) -> torch.Tensor:
+        input_sequence = feature_tensor[self._feature_name]
+        return torch.tril(
+            torch.ones(
+                (input_sequence.size(1), input_sequence.size(1)),
+                dtype=torch.bool,
+                device=input_sequence.device,
+            )
+        )