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

replay/nn/lightning/module.py

@@ -0,0 +1,123 @@
+import inspect
+from typing import Any, Optional, Union
+
+import lightning
+import torch
+from typing_extensions import override
+
+from replay.nn.lightning.optimizer import BaseOptimizerFactory, OptimizerFactory
+from replay.nn.lightning.scheduler import BaseLRSchedulerFactory
+from replay.nn.output import InferenceOutput, TrainOutput
+
+
+class LightningModule(lightning.LightningModule):
+    """
+    A universal wrapper class above the PyTorch model for working with Lightning library.\n
+    Pay attention to the format of the ``forward`` function's return value.
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        optimizer_factory: Optional[BaseOptimizerFactory] = None,
+        lr_scheduler_factory: Optional[BaseLRSchedulerFactory] = None,
+    ) -> None:
+        """
+        :param model: Initialized model.\n
+            Expected result of the model's ``forward`` function
+            is an object of the ``TrainOutput`` class after training stage
+            and ``InferenceOutput`` after inference stage.
+        :param optimizer_factory: Optimizer factory.
+            Default: ``None``.
+        :param lr_scheduler_factory: Learning rate schedule factory.
+            Default: ``None``.
+        """
+        super().__init__()
+        self.save_hyperparameters(ignore=["model"])
+        self.model = model
+
+        self._optimizer_factory = optimizer_factory
+        self._lr_scheduler_factory = lr_scheduler_factory
+        self.candidates_to_score = None
+
+    def forward(self, batch: dict) -> Union[TrainOutput, InferenceOutput]:
+        """
+        Implementation of the forward function.
+
+        :param batch: A dictionary containing all the necessary information to run the forward function on the model.
+            The dictionary keys must match the names of the arguments in the model's forward function.
+            Keys that do not match the arguments of the model's forward function are filtered out.
+            If the model supports calculating logits for custom candidates on the inference stage,
+            then you can submit them inside the batch or using the ``candidates_to_score`` field.
+        :returns: During training, the model will return an object
+            of the ``TrainOutput`` container class or its successor.
+            At the inference stage, the ``InferenceOutput`` class or its successor will be returned.
+        """
+        if "candidates_to_score" not in batch and self.candidates_to_score is not None and not self.training:
+            batch["candidates_to_score"] = self.candidates_to_score
+        # select only args for model.forward
+        modified_batch = {k: v for k, v in batch.items() if k in inspect.signature(self.model.forward).parameters}
+        return self.model(**modified_batch)
+
+    def training_step(self, batch: dict) -> torch.Tensor:
+        model_output: TrainOutput = self(batch)
+        loss = model_output["loss"]
+        lr = self.optimizers().param_groups[0]["lr"]  # Get current learning rate
+        self.log("learning_rate", lr, on_step=True, on_epoch=True, prog_bar=True)
+        self.log(
+            "train_loss",
+            loss,
+            on_step=True,
+            on_epoch=True,
+            prog_bar=True,
+            sync_dist=True,
+        )
+        return loss
+
+    @override
+    def predict_step(self, batch: dict, batch_idx: int, dataloader_idx: int = 0) -> torch.Tensor:
+        model_output: InferenceOutput = self(batch)
+        return model_output
+
+    @override
+    def test_step(self, batch: dict, batch_idx: int, dataloader_idx: int = 0) -> torch.Tensor:
+        model_output: InferenceOutput = self(batch)
+        return model_output
+
+    @override
+    def validation_step(self, batch: dict, batch_idx: int, dataloader_idx: int = 0) -> torch.Tensor:
+        model_output: InferenceOutput = self(batch)
+        return model_output
+
+    def configure_optimizers(self) -> Any:
+        """
+        Returns:
+            Tuple[List[torch.optim.Optimizer], List[torch.optim.lr_scheduler._LRScheduler]]:
+                Configured optimizer and lr scheduler.
+        """
+        optimizer_factory = self._optimizer_factory or OptimizerFactory()
+        optimizer = optimizer_factory.create(self.model.parameters())
+
+        if self._lr_scheduler_factory is None:
+            return optimizer
+
+        lr_scheduler = self._lr_scheduler_factory.create(optimizer)
+        return [optimizer], [lr_scheduler]
+
+    @property
+    def candidates_to_score(self) -> Optional[torch.LongTensor]:
+        """
+        :getter: Returns a tensor containing the candidate IDs.
+            The tensor will be used during the inference stage of the model.\n
+            If the parameter was not previously set, ``None`` will be returned.
+        :setter: A one-dimensional tensor containing candidate IDs is expected.
+        """
+        return self._candidates_to_score
+
+    @candidates_to_score.setter
+    def candidates_to_score(self, candidates: Optional[torch.LongTensor] = None) -> None:
+        if (candidates is not None) and bool(candidates.unique().numel() != candidates.numel()):
+            msg = "The tensor of candidates to score must be unique."
+            raise ValueError(msg)
+
+        self._candidates_to_score = candidates

replay/nn/lightning/optimizer.py

@@ -0,0 +1,60 @@
+import abc
+from collections.abc import Iterator
+from typing import Literal
+
+import torch
+
+
+class BaseOptimizerFactory(abc.ABC):
+    """
+    Interface for optimizer factory
+    """
+
+    @abc.abstractmethod
+    def create(self, parameters: Iterator[torch.nn.Parameter]) -> torch.optim.Optimizer:  # pragma: no cover
+        """
+        Creates optimizer based on parameters.
+
+        :param parameters: torch parameters to initialize optimizer
+
+        :returns: torch optimizer
+        """
+
+
+class OptimizerFactory(BaseOptimizerFactory):
+    """
+    Factory that creates optimizer depending on passed parameters
+    """
+
+    def __init__(
+        self,
+        optimizer: Literal["adam", "sgd"] = "adam",
+        learning_rate: float = 0.001,
+        weight_decay: float = 0.0,
+        sgd_momentum: float = 0.0,
+        betas: tuple[float, float] = (0.9, 0.98),
+    ) -> None:
+        super().__init__()
+        self.optimizer = optimizer
+        self.learning_rate = learning_rate
+        self.weight_decay = weight_decay
+        self.sgd_momentum = sgd_momentum
+        self.betas = betas
+
+    def create(self, parameters: Iterator[torch.nn.Parameter]) -> torch.optim.Optimizer:
+        """
+        Creates optimizer based on parameters.
+
+        :param parameters: torch parameters to initialize optimizer
+
+        :returns: torch optimizer
+        """
+        if self.optimizer == "adam":
+            return torch.optim.Adam(parameters, lr=self.learning_rate, weight_decay=self.weight_decay, betas=self.betas)
+        if self.optimizer == "sgd":
+            return torch.optim.SGD(
+                parameters, lr=self.learning_rate, weight_decay=self.weight_decay, momentum=self.sgd_momentum
+            )
+
+        msg = "Unexpected optimizer"
+        raise ValueError(msg)

replay/nn/lightning/postprocessor/__init__.py

@@ -0,0 +1,2 @@
+from ._base import PostprocessorBase
+from .seen_items import SeenItemsFilter

replay/nn/lightning/postprocessor/_base.py

@@ -0,0 +1,51 @@
+import abc
+from typing import Optional, Union
+
+import torch
+
+
+class PostprocessorBase(abc.ABC):  # pragma: no cover
+    """
+    Abstract base class for post processor
+    """
+
+    @abc.abstractmethod
+    def on_validation(self, batch: dict, logits: torch.Tensor) -> torch.Tensor:
+        """
+        The method is called externally inside the callback at the validation stage.
+
+        :param batch: the batch sent to the model from the dataloader
+        :param logits: logits from the model
+
+        :returns: modified logits
+        """
+
+    @abc.abstractmethod
+    def on_prediction(self, batch: dict, logits: torch.Tensor) -> torch.Tensor:
+        """
+        The method is called externally inside the callback at the prediction (inference) stage.
+
+        :param batch: the batch sent to the model from the dataloader
+        :param logits: logits from the model
+
+        :returns: modified logits
+        """
+
+    @property
+    def candidates(self) -> Union[torch.LongTensor, None]:
+        """
+        Returns tensor of item ids to calculate scores.
+        """
+        return self._candidates
+
+    @candidates.setter
+    def candidates(self, candidates: Optional[torch.LongTensor] = None) -> None:
+        """
+        Sets tensor of item ids to calculate scores.
+        :param candidates: Tensor of item ids to calculate scores.
+        """
+        if (candidates is not None) and bool(candidates.unique().numel() != candidates.numel()):
+            msg = "The tensor of candidates to score must be unique."
+            raise ValueError(msg)
+
+        self._candidates = candidates

replay/nn/lightning/postprocessor/seen_items.py

@@ -0,0 +1,83 @@
+import torch
+
+from replay.data.nn import TensorMap
+
+from ._base import PostprocessorBase
+
+
+class SeenItemsFilter(PostprocessorBase):
+    """
+    Masks (sets logits value to ``-inf``) the items that already have been seen in the given dataset
+    (i.e. in the sequence of items for that logits are calculated).\n
+    Should be used in Lightning callbacks for inferencing or metrics computing.
+
+    .. rubric:: Input example:
+
+    logits [B=2 users, I=3 items]::
+
+        logits =
+        [[0.1, 0.2, 0.3],    # user0
+        [-0.1, -0.2, -0.3]]  # user1
+
+    Seen items per user::
+
+        seen_items =
+        user0: [1, 0]
+        user1: [1, 2, 1]
+
+    .. rubric:: Output example:
+
+    SeenItemsFilter sets logits of seen items to ``-inf``::
+
+        processed_logits =
+        [[   -inf,    -inf,  0.3000], # user0
+        [-0.1000,    -inf,    -inf]]  # user1
+
+    """
+
+    def __init__(self, item_count: int, seen_items_column="seen_ids") -> None:
+        """
+        :param item_count: Total number of items that the model knows about (``cardinality``).
+            It is recommended to take this value from ``TensorSchema``. \n
+            Please note that values ​​outside the range [0, `item_count-1`] are filtered out (considered as padding).
+        :param seen_items_column: Name of the column in batch that contains users' interactions (seen item ids).
+        """
+        super().__init__()
+        self.item_count = item_count
+        self.seen_items_column = seen_items_column
+        self._candidates = None
+
+    def on_validation(self, batch: dict, logits: torch.Tensor) -> torch.Tensor:
+        return self._compute_scores(batch, logits.detach().clone())
+
+    def on_prediction(self, batch: dict, logits: torch.Tensor) -> torch.Tensor:
+        return self._compute_scores(batch, logits.detach().clone())
+
+    def _compute_scores(
+        self,
+        batch: TensorMap,
+        logits: torch.Tensor,
+    ) -> torch.Tensor:
+        seen_ids_padded = batch[self.seen_items_column]
+        padding_mask = (seen_ids_padded < self.item_count) & (seen_ids_padded >= 0)
+
+        batch_factors = torch.arange(0, logits.size(0), device=logits.device) * self.item_count
+        factored_ids = seen_ids_padded + batch_factors.unsqueeze(1)
+        seen_ids_flat = factored_ids[padding_mask]
+
+        if self._candidates is not None:
+            _logits = torch.full((logits.size(0), self.item_count), -torch.inf)
+            _logits[:, self._candidates] = torch.reshape(logits, _logits[:, self.candidates].shape)
+            logits = _logits
+
+        if logits.is_contiguous():
+            logits.view(-1)[seen_ids_flat] = -torch.inf
+        else:
+            flat_scores = logits.flatten()
+            flat_scores[seen_ids_flat] = -torch.inf
+            logits = flat_scores.reshape(logits.shape)
+
+        if self._candidates is not None:
+            logits = logits[:, self._candidates]
+
+        return logits

replay/nn/lightning/scheduler.py

@@ -0,0 +1,91 @@
+import abc
+import warnings
+from typing import Literal
+
+import torch
+
+
+class BaseLRSchedulerFactory(abc.ABC):
+    """
+    Interface for learning rate scheduler factory
+    """
+
+    @abc.abstractmethod
+    def create(self, optimizer: torch.optim.Optimizer) -> torch.optim.lr_scheduler._LRScheduler:  # pragma: no cover
+        """
+        Creates learning rate scheduler based on optimizer.
+
+        :param optimizer: torch optimizer
+
+        :returns: torch LRScheduler
+        """
+
+
+class LRSchedulerFactory(BaseLRSchedulerFactory):
+    """
+    Factory that creates learning rate schedule depending on passed parameters
+    """
+
+    def __init__(self, decay_step: int = 25, gamma: float = 1.0) -> None:
+        super().__init__()
+        self.decay_step = decay_step
+        self.gamma = gamma
+
+    def create(self, optimizer: torch.optim.Optimizer) -> torch.optim.lr_scheduler._LRScheduler:
+        """
+        Creates learning rate scheduler based on optimizer.
+
+        :param optimizer: torch optimizer
+
+        :returns: torch LRScheduler
+        """
+        return torch.optim.lr_scheduler.StepLR(optimizer, step_size=self.decay_step, gamma=self.gamma)
+
+
+class LambdaLRSchedulerFactory(BaseLRSchedulerFactory):
+    """
+    Factory that creates learning rate schedule depending on passed parameters
+    """
+
+    def __init__(
+        self,
+        warmup_steps: int,
+        warmup_lr: float = 1.0,
+        normal_lr: float = 0.1,
+        update_interval: Literal["epoch", "step"] = "epoch",
+    ) -> None:
+        super().__init__()
+
+        if normal_lr <= 0.0:
+            msg = f"Normal LR must be positive. Got {normal_lr}"
+            raise ValueError(msg)
+        if warmup_lr <= 0.0:
+            msg = f"Warmup LR must be positive. Got {warmup_lr}"
+            raise ValueError(msg)
+        if normal_lr >= warmup_lr:
+            msg = f"Suspicious LR pair: {normal_lr=}, {warmup_lr=}"
+            warnings.warn(msg, stacklevel=2)
+
+        self.warmup_lr = warmup_lr
+        self.normal_lr = normal_lr
+        self.warmup_steps = warmup_steps
+        self.update_interval = update_interval
+
+    def create(self, optimizer: torch.optim.Optimizer) -> torch.optim.lr_scheduler._LRScheduler:
+        """
+        Creates learning rate scheduler based on optimizer.
+
+        :param optimizer: torch optimizer
+
+        :returns: torch LambdaLR
+        """
+        return {
+            "scheduler": torch.optim.lr_scheduler.LambdaLR(optimizer, self.lr_lambda),
+            "interval": self.update_interval,
+            "frequency": 1,
+        }
+
+    def lr_lambda(self, current_step: int) -> float:
+        if current_step >= self.warmup_steps:
+            return self.normal_lr
+        return self.warmup_lr

replay/nn/loss/__init__.py

@@ -0,0 +1,22 @@
+from .base import LossProto
+from .bce import BCE, BCESampled
+from .ce import CE, CESampled, CESampledWeighted, CEWeighted
+from .login_ce import LogInCE, LogInCESampled
+from .logout_ce import LogOutCE, LogOutCEWeighted
+
+LogOutCESampled = CE
+
+__all__ = [
+    "BCE",
+    "CE",
+    "BCESampled",
+    "CESampled",
+    "CESampledWeighted",
+    "CEWeighted",
+    "LogInCE",
+    "LogInCESampled",
+    "LogOutCE",
+    "LogOutCESampled",
+    "LogOutCEWeighted",
+    "LossProto",
+]

replay/nn/loss/base.py

@@ -0,0 +1,197 @@
+from typing import Callable, Optional, Protocol, TypedDict
+
+import torch
+
+from replay.data.nn import TensorMap
+
+
+class LossProto(Protocol):
+    """Class-protocol for working with losses inside models"""
+
+    @property
+    def logits_callback(
+        self,
+    ) -> Callable[[torch.Tensor, Optional[torch.Tensor]], torch.Tensor]: ...
+
+    @logits_callback.setter
+    def logits_callback(self, func: Optional[Callable]) -> None: ...
+
+    def forward(
+        self,
+        model_embeddings: torch.Tensor,
+        feature_tensors: TensorMap,
+        positive_labels: torch.LongTensor,
+        negative_labels: torch.LongTensor,
+        padding_mask: torch.BoolTensor,
+        target_padding_mask: torch.BoolTensor,
+    ) -> torch.Tensor: ...
+
+
+class SampledLossOutput(TypedDict):
+    """A class containing result of the `get_sampled_logits` function in sampled losses"""
+
+    positive_logits: torch.Tensor
+    negative_logits: torch.Tensor
+    positive_labels: torch.LongTensor
+    negative_labels: torch.LongTensor
+
+
+class SampledLossBase(torch.nn.Module):
+    """The base class for calculating sampled losses"""
+
+    @property
+    def logits_callback(
+        self,
+    ) -> Callable[[torch.Tensor, Optional[torch.Tensor]], torch.Tensor]:
+        raise NotImplementedError()  # pragma: no cover
+
+    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]
+    ) -> SampledLossOutput:
+        """
+        The function of calculating positive and negative logits.
+        Based on the model last hidden state, 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: SampledLossOutput. A dictionary containing positive and negative logits with labels.
+        """
+
+        initial_positive_labels = positive_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)
+
+        if negative_labels.dim() == 3:  # pragma: no cover
+            # [batch_size, seq_len, num_negatives] -> [batch_size, seq_len, 1, num_negatives]
+            negative_labels = negative_labels.unsqueeze(-2)
+            if num_positives != 1:
+                # [batch_size, seq_len, num_negatives] -> [batch_size, seq_len, num_positives, num_negatives]
+                negative_labels = negative_labels.repeat((1, 1, num_positives, 1))
+        assert (
+            negative_labels.size() == (batch_size, seq_len, num_positives, 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, emb_dim] ->  [batch_size, seq_len, 1, emb_dim]
+        model_embeddings = model_embeddings.unsqueeze(-2)
+        if num_positives != 1:  # multti positive branch
+            model_embeddings = model_embeddings.repeat((1, 1, num_positives, 1))
+        assert model_embeddings.size() == (
+            batch_size,
+            seq_len,
+            num_positives,
+            embedding_dim,
+        )
+
+        # Apply target mask
+        # [batch_size, seq_len, num_positives] -> [batch_size, seq_len]
+        masked_batch_size = target_padding_mask.sum().item()
+
+        # [batch_size, seq_len, num_positives] -> [masked_batch_size, 1]
+        positive_labels = positive_labels[target_padding_mask].unsqueeze(-1)
+        assert positive_labels.size() == (masked_batch_size, 1)
+
+        if negative_labels.dim() != 1:  # pragma: no cover
+            # [batch_size, seq_len, num_positives, num_negatives] -> [masked_batch_size, num_negatives]
+            negative_labels = negative_labels[target_padding_mask]
+            assert negative_labels.size() == (masked_batch_size, num_negatives)
+
+        # [batch_size, seq_len, num_positives, emb_dim] -> [masked_batch_size, emb_dim]
+        model_embeddings = model_embeddings[target_padding_mask]
+        assert model_embeddings.size() == (masked_batch_size, embedding_dim)
+
+        # Get positive and negative logits
+        positive_logits = self.logits_callback(model_embeddings, positive_labels)
+        assert positive_logits.size() == (masked_batch_size, 1)
+
+        negative_logits = self.logits_callback(model_embeddings, negative_labels)
+        assert negative_logits.size() == (masked_batch_size, num_negatives)
+
+        if num_positives != 1:
+            # [batch_size, seq_len, num_positives] -> [batch_size * seq_len]
+            masked_target_padding_mask = target_padding_mask.sum(-1).view(-1)
+            # [batch_size, seq_len, num_positives] -> [masked_batch_size, num_positives]
+            positive_labels = torch.repeat_interleave(
+                initial_positive_labels.view(-1, num_positives),
+                masked_target_padding_mask,
+                dim=0,
+            )
+
+        return {
+            "positive_logits": positive_logits,
+            "negative_logits": negative_logits,
+            "positive_labels": positive_labels,
+            "negative_labels": negative_labels,
+        }
+
+
+def mask_negative_logits(
+    negative_logits: torch.Tensor,
+    negative_labels: torch.LongTensor,
+    positive_labels: torch.LongTensor,
+    negative_labels_ignore_index: int,
+) -> torch.Tensor:
+    """
+    Assign very small values in negative logits
+    for positions where positive labels equal to negative ones.
+
+    :param negative_logits: Logits from the model for ``negative labels``.
+        Expected shape: (masked_batch_size, num_negatives)
+    :param negative_labels: a tensor containing labels with negative events.
+        Expected shape:
+            - (masked_batch_size, num_negatives)
+            - (num_negatives) - a case where the same negative events are used for the entire batch
+    :param positive_labels: a tensor containing labels with positive events.
+        Expected shape: (masked_batch_size, num_positives)
+    :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.
+
+    :returns: Negative logits with modified elements in those positions
+        where positive labels are equal to negative ones.
+    """
+
+    if negative_labels_ignore_index >= 0:
+        negative_logits.masked_fill_(negative_labels == negative_labels_ignore_index, -1e9)
+
+    if negative_labels.dim() > 1:  # pragma: no cover
+        # [masked_batch_size, num_negatives] -> [masked_batch_size, 1, num_negatives]
+        negative_labels = negative_labels.unsqueeze(-2)
+
+    # [masked_batch_size, num_positives] -> [masked_batch_size, num_positives, 1]
+    positive_labels = positive_labels.unsqueeze(-1)
+    negative_mask = positive_labels == negative_labels  # [masked_batch_size, num_positives, num_negatives]
+
+    # [masked_batch_size, num_positives, num_negatives] -> [masked_batch_size, num_negatives]
+    negative_mask = negative_mask.sum(-2).bool()
+    negative_logits.masked_fill_(negative_mask, -1e9)
+    return negative_logits