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

replay/nn/ffn.py

@@ -0,0 +1,135 @@
+import contextlib
+from typing import Literal
+
+import torch
+
+from replay.data.nn.schema import TensorMap
+
+from .utils import create_activation
+
+
+class PointWiseFeedForward(torch.nn.Module):
+    """
+    Point wise feed forward network layer.
+
+    Source paper: https://arxiv.org/pdf/1808.09781.pdf
+    """
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        dropout: float,
+        activation: Literal["relu", "gelu"] = "gelu",
+    ) -> None:
+        """
+        :param embedding_dim: Dimension of the input features.
+        :param dropout: probability of an element to be zeroed.
+        :param activation: the name of the activation function.
+            Default: ``"gelu"``.
+        """
+        super().__init__()
+
+        self.conv1 = torch.nn.Conv1d(embedding_dim, embedding_dim, kernel_size=1)
+        self.dropout1 = torch.nn.Dropout(p=dropout)
+        self.activation = create_activation(activation)
+        self.conv2 = torch.nn.Conv1d(embedding_dim, embedding_dim, kernel_size=1)
+        self.dropout2 = torch.nn.Dropout(p=dropout)
+
+    def reset_parameters(self) -> None:
+        for _, param in self.named_parameters():
+            with contextlib.suppress(ValueError):
+                torch.nn.init.xavier_normal_(param.data)
+
+    def forward(self, input_embeddings: torch.LongTensor) -> torch.LongTensor:
+        """
+        :param input_embeddings: Query feature tensor.
+
+        :returns: Output tensors.
+        """
+        x: torch.Tensor = self.conv1(input_embeddings.transpose(-1, -2))
+        x = self.activation(x)
+        x = self.dropout1(x)
+        x = self.conv2(x)
+        x = self.dropout2(x)
+        x = x.transpose(-1, -2)
+        x += input_embeddings
+
+        return x
+
+
+class SwiGLU(torch.nn.Module):
+    """
+    SwiGLU Activation Function.
+    Combines the Swish activation with Gated Linear Units.
+    """
+
+    def __init__(self, embedding_dim: int, hidden_dim: int):
+        """
+        :param embedding_dim: Dimension of the input features.
+        :param hidden_dim: Dimension of hidden layer.
+            According to the original source,
+            it is recommended to set the size of the hidden layer as :math:`2 \\cdot \\text{embedding_dim}`.
+        """
+        super().__init__()
+        # Intermediate projection layers
+        # Typically, SwiGLU splits the computation into two parts
+        self.WG = torch.nn.Linear(embedding_dim, hidden_dim)
+        self.W1 = torch.nn.Linear(embedding_dim, hidden_dim)
+        self.W2 = torch.nn.Linear(hidden_dim, embedding_dim)
+
+    def reset_parameters(self) -> None:
+        for _, param in self.named_parameters():
+            with contextlib.suppress(ValueError):
+                torch.nn.init.xavier_normal_(param.data)
+
+    def forward(
+        self,
+        input_embeddings: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Forward pass for SwiGLU.
+
+        :param input_embeddings: Input tensor of shape ``(batch_size, sequence_length, embedding_dim)``.
+
+        :returns: Output tensor of shape ``(batch_size, sequence_length, embedding_dim)``.
+        """
+        # Apply the gates
+        activation = torch.nn.functional.silu(self.WG(input_embeddings))  # Activation part
+        linear = self.W1(input_embeddings)  # Linear part
+        return self.W2(activation * linear)  # Element-wise multiplication and projection
+
+
+class SwiGLUEncoder(torch.nn.Module):
+    """
+    MLP block consists of SwiGLU Feed-Forward network followed by a RMSNorm layer with skip connection.
+    """
+
+    def __init__(self, embedding_dim: int, hidden_dim: int) -> None:
+        """
+        :param embedding_dim: Dimension of the input features.
+        """
+        super().__init__()
+        self.sw1 = SwiGLU(embedding_dim, hidden_dim)
+        self.norm1 = torch.nn.RMSNorm(embedding_dim)
+        self.sw2 = SwiGLU(embedding_dim, hidden_dim)
+        self.norm2 = torch.nn.RMSNorm(embedding_dim)
+
+    def reset_parameters(self) -> None:
+        self.sw1.reset_parameters()
+        self.sw2.reset_parameters()
+        self.norm1.reset_parameters()
+        self.norm2.reset_parameters()
+
+    def forward(
+        self,
+        feature_tensors: TensorMap,  # noqa: ARG002
+        input_embeddings: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        forward(input_embeddings)
+        :param input_embeddings: Input tensor of shape ``(batch_size, sequence_length, embedding_dim)``.
+        :returns: Output tensor of shape ``(batch_size, sequence_length, embedding_dim)``.
+        """
+        x = self.norm1(self.sw1(input_embeddings) + input_embeddings)
+        x = self.norm2(self.sw2(x) + x)
+        return x

replay/nn/head.py

@@ -0,0 +1,49 @@
+import torch
+
+
+class EmbeddingTyingHead(torch.nn.Module):
+    """
+    The model head for calculating the output logits as a dot product
+    between the model hidden state and the item embeddings.
+    The module supports both 2-d and 3-d tensors for the hidden state and the item embeddings.
+
+    As a result of the work, the scores for each item will be obtained.
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        item_embeddings: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        :param hidden_states: hidden state of shape
+            ``(batch_size, embedding_dim)`` or ``(batch_size, sequence_length, embedding_dim)``.
+        :param item_embeddings: item embeddings of shape
+            ``(num_items, embedding_dim)`` or ``(batch_size, num_items, embedding_dim)``.
+        :return: logits of shape ``(batch_size, num_items)``
+            or ``(batch_size, sequence_length, num_items)``.
+        """
+        if item_embeddings.dim() == 2:
+            item_embeddings = item_embeddings.transpose(-1, -2).contiguous()
+            # hidden_states shape [B, *, E]
+            # item embeddings shape [I, E]
+            # [B, *, E] x [E, I] -> [B, *, I]
+            return hidden_states.matmul(item_embeddings)
+        elif item_embeddings.dim() == 3 and hidden_states.dim() == 2:
+            item_embeddings = item_embeddings.transpose(-1, -2).contiguous()
+            # out_embeddings shape [B, E]
+            # item embeddings shape [B, I, E]
+            # [B, E] x [B, E, I] -> [B, I]
+            hidden_states = hidden_states.unsqueeze(-2)
+            logits = hidden_states.matmul(item_embeddings)
+            return logits.squeeze(-2)
+        # out_embeddings shape: [B, *, E]
+        # item embeddings shape [B, *, E]
+        # [*, 1, E] x [*, E, 1] -> [B, *]
+        return torch.bmm(
+            hidden_states.view(-1, 1, hidden_states.size(-1)),
+            item_embeddings.view(-1, item_embeddings.size(-1), 1),
+        ).view(hidden_states.size(0), *item_embeddings.shape[1:-1])

replay/nn/lightning/__init__.py

@@ -0,0 +1 @@
+from .module import LightningModule

replay/nn/lightning/callback/__init__.py

@@ -0,0 +1,9 @@
+from .metrics_callback import ComputeMetricsCallback
+from .predictions_callback import (
+    HiddenStatesCallback,
+    PandasTopItemsCallback,
+    PolarsTopItemsCallback,
+    SparkTopItemsCallback,
+    TopItemsCallbackBase,
+    TorchTopItemsCallback,
+)

replay/nn/lightning/callback/metrics_callback.py

@@ -0,0 +1,183 @@
+from typing import Any, Optional
+
+import lightning
+import torch
+from lightning.pytorch.utilities.combined_loader import CombinedLoader
+from lightning.pytorch.utilities.rank_zero import rank_zero_only
+
+from replay.metrics.torch_metrics_builder import (
+    MetricName,
+    TorchMetricsBuilder,
+    metrics_to_df,
+)
+from replay.nn.lightning import LightningModule
+from replay.nn.lightning.postprocessor import PostprocessorBase
+from replay.nn.output import InferenceOutput
+
+
+class ComputeMetricsCallback(lightning.Callback):
+    """
+    Callback for validation and testing stages.
+
+    If multiple validation/testing dataloaders are used,
+    the suffix of the metric name will contain the serial number of the dataloader.
+
+    For the correct calculation of metrics inside the callback,
+    the batch must contain the ``ground_truth_column`` key - the padding value of this tensor can be any,
+    the main condition is that the padding value does not overlap with the existing item ID values.
+    For example, these can be negative values.
+
+    To calculate the ``coverage`` and ``novelty`` metrics, the batch must additionally contain the ``train_column`` key.
+    The padding value of this tensor can be any, the main condition is that the padding value does not overlap
+    with the existing item ID values. For example, these can be negative values.
+    """
+
+    def __init__(
+        self,
+        metrics: Optional[list[MetricName]] = None,
+        ks: Optional[list[int]] = None,
+        postprocessors: Optional[list[PostprocessorBase]] = None,
+        item_count: Optional[int] = None,
+        ground_truth_column: str = "ground_truth",
+        train_column: str = "train",
+    ):
+        """
+        :param metrics: Sequence of metrics to calculate.\n
+            Default: ``None``. This means that the default metrics will be used - ``Map``, ``NDCG``, ``Recall``.
+        :param ks: highest k scores in ranking.\n
+            Default: ``None``. This means that the default ``ks`` will be ``[1, 5, 10, 20]``.
+        :param postprocessors: A list of postprocessors for modifying logits from the model.
+            For example, it can be a softmax operation to logits or set the ``-inf`` value for some IDs.
+            Default: ``None``.
+        :param item_count: the total number of items in the dataset, required only for ``Coverage`` calculations.
+            Default: ``None``.
+        :param ground_truth_column: Name of key in batch that contains ground truth items.
+        :param train_column: Name of key in batch that contains items on which the model is trained.
+        """
+        self._metrics = metrics
+        self._ks = ks
+        self._item_count = item_count
+        self._metrics_builders: list[TorchMetricsBuilder] = []
+        self._dataloaders_size: list[int] = []
+        self._postprocessors: list[PostprocessorBase] = postprocessors or []
+        self._ground_truth_column = ground_truth_column
+        self._train_column = train_column
+
+    def _get_dataloaders_size(self, dataloaders: Optional[Any]) -> list[int]:
+        if isinstance(dataloaders, CombinedLoader):
+            return [len(dataloader) for dataloader in dataloaders.flattened]  # pragma: no cover
+        return [len(dataloaders)]
+
+    def on_validation_epoch_start(
+        self,
+        trainer: lightning.Trainer,
+        pl_module: LightningModule,  # noqa: ARG002
+    ) -> None:
+        self._dataloaders_size = self._get_dataloaders_size(trainer.val_dataloaders)
+        self._metrics_builders = [
+            TorchMetricsBuilder(self._metrics, self._ks, self._item_count) for _ in self._dataloaders_size
+        ]
+        for builder in self._metrics_builders:
+            builder.reset()
+
+    def on_test_epoch_start(
+        self,
+        trainer: lightning.Trainer,
+        pl_module: LightningModule,  # noqa: ARG002
+    ) -> None:
+        self._dataloaders_size = self._get_dataloaders_size(trainer.test_dataloaders)
+        self._metrics_builders = [
+            TorchMetricsBuilder(self._metrics, self._ks, self._item_count) for _ in self._dataloaders_size
+        ]
+        for builder in self._metrics_builders:
+            builder.reset()
+
+    def _apply_postproccesors(self, batch: dict, logits: torch.Tensor) -> torch.Tensor:
+        for postprocessor in self._postprocessors:
+            logits = postprocessor.on_validation(batch, logits)
+        return logits
+
+    def on_validation_batch_end(
+        self,
+        trainer: lightning.Trainer,
+        pl_module: LightningModule,
+        outputs: InferenceOutput,
+        batch: dict,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:
+        self._batch_end(
+            trainer,
+            pl_module,
+            outputs,
+            batch,
+            batch_idx,
+            dataloader_idx,
+        )
+
+    def on_test_batch_end(
+        self,
+        trainer: lightning.Trainer,
+        pl_module: LightningModule,
+        outputs: InferenceOutput,
+        batch: dict,
+        batch_idx: int,
+        dataloader_idx: int = 0,
+    ) -> None:  # pragma: no cover
+        self._batch_end(
+            trainer,
+            pl_module,
+            outputs,
+            batch,
+            batch_idx,
+            dataloader_idx,
+        )
+
+    def _batch_end(
+        self,
+        trainer: lightning.Trainer,  # noqa: ARG002
+        pl_module: LightningModule,
+        outputs: InferenceOutput,
+        batch: dict,
+        batch_idx: int,
+        dataloader_idx: int,
+    ) -> None:
+        seen_scores = self._apply_postproccesors(batch, outputs["logits"])
+        sampled_items = torch.topk(seen_scores, k=self._metrics_builders[dataloader_idx].max_k, dim=1).indices
+        self._metrics_builders[dataloader_idx].add_prediction(
+            sampled_items, batch[self._ground_truth_column], batch.get(self._train_column)
+        )
+
+        if batch_idx + 1 == self._dataloaders_size[dataloader_idx]:
+            pl_module.log_dict(
+                self._metrics_builders[dataloader_idx].get_metrics(),
+                on_epoch=True,
+                sync_dist=True,
+                add_dataloader_idx=True,
+            )
+
+    def on_validation_epoch_end(self, trainer: lightning.Trainer, pl_module: LightningModule) -> None:
+        self._epoch_end(trainer, pl_module)
+
+    def on_test_epoch_end(self, trainer: lightning.Trainer, pl_module: LightningModule) -> None:  # pragma: no cover
+        self._epoch_end(trainer, pl_module)
+
+    def _epoch_end(
+        self,
+        trainer: lightning.Trainer,
+        pl_module: LightningModule,  # noqa: ARG002
+    ) -> None:
+        @rank_zero_only
+        def print_metrics() -> None:
+            metrics = {}
+            for name, value in trainer.logged_metrics.items():
+                if "@" in name:
+                    metrics[name] = value.item()
+
+            if metrics:
+                metrics_df = metrics_to_df(metrics)
+
+                print(metrics_df)  # noqa: T201
+                print()  # noqa: T201
+
+        print_metrics()

replay/nn/lightning/callback/predictions_callback.py

@@ -0,0 +1,314 @@
+import abc
+from typing import Generic, Optional, TypeVar
+
+import lightning
+import torch
+
+from replay.nn.lightning import LightningModule
+from replay.nn.lightning.postprocessor import PostprocessorBase
+from replay.nn.output import InferenceOutput
+from replay.utils import (
+    PYSPARK_AVAILABLE,
+    MissingImport,
+    PandasDataFrame,
+    PolarsDataFrame,
+    SparkDataFrame,
+)
+
+if PYSPARK_AVAILABLE:  # pragma: no cover
+    import pyspark.sql.functions as sf
+    from pyspark.sql import SparkSession
+    from pyspark.sql.types import ArrayType, DoubleType, IntegerType, StructType
+else:  # pragma: no cover
+    SparkSession = MissingImport
+
+
+_T = TypeVar("_T")
+
+
+class TopItemsCallbackBase(lightning.Callback, Generic[_T]):
+    """
+    The base class for a callback that records the result at the inference stage via ``LightningModule``.
+    The result consists of top K the highest logit values, IDs of these  top K logit values
+    and corresponding query ids (encoded IDs of users named ``query_id``).
+
+    For the callback to work correctly, the batch is expected to contain the ``query_id`` key.
+    """
+
+    def __init__(
+        self,
+        top_k: int,
+        query_column: str,
+        item_column: str,
+        rating_column: str = "rating",
+        postprocessors: Optional[list[PostprocessorBase]] = None,
+    ) -> None:
+        """
+        :param top_k: Take the ``top_k`` IDs with the highest logit values.
+        :param query_column: The name of the query column in the resulting dataframe.
+        :param item_column: The name of the item column in the resulting dataframe.
+        :param rating_column: The name of the rating column in the resulting dataframe.
+            This column will contain the ``top_k`` items with the highest logit values.
+        :param postprocessors: A list of postprocessors for modifying logits from the model
+            before sorting and taking top K ones.
+            For example, it can be a softmax operation to logits or set the ``-inf`` value for some IDs.
+            Default: ``None``.
+        """
+        super().__init__()
+        self.query_column = query_column
+        self.item_column = item_column
+        self.rating_column = rating_column
+        self._top_k = top_k
+        self._postprocessors: list[PostprocessorBase] = postprocessors or []
+        self._query_batches: list[torch.Tensor] = []
+        self._item_batches: list[torch.Tensor] = []
+        self._item_scores: list[torch.Tensor] = []
+
+    def on_predict_epoch_start(
+        self,
+        trainer: lightning.Trainer,  # noqa: ARG002
+        pl_module: LightningModule,
+    ) -> None:
+        self._query_batches.clear()
+        self._item_batches.clear()
+        self._item_scores.clear()
+
+        candidates = pl_module.candidates_to_score
+        for postprocessor in self._postprocessors:
+            postprocessor.candidates = candidates
+
+    def on_predict_batch_end(
+        self,
+        trainer: lightning.Trainer,  # noqa: ARG002
+        pl_module: LightningModule,
+        outputs: InferenceOutput,
+        batch: dict,
+        batch_idx: int,  # noqa: ARG002
+        dataloader_idx: int = 0,  # noqa: ARG002
+    ) -> None:
+        logits = self._apply_postproccesors(batch, outputs["logits"])
+        top_scores, top_item_ids = torch.topk(logits, k=self._top_k, dim=1)
+        if pl_module.candidates_to_score is not None:
+            top_item_ids = torch.take(pl_module.candidates_to_score, top_item_ids)
+
+        self._query_batches.append(batch["query_id"])
+        self._item_batches.append(top_item_ids)
+        self._item_scores.append(top_scores)
+
+    def get_result(self) -> _T:
+        """
+        :returns: prediction result
+        """
+        prediction = self._ids_to_result(
+            torch.cat(self._query_batches),
+            torch.cat(self._item_batches),
+            torch.cat(self._item_scores),
+        )
+        return prediction
+
+    def _apply_postproccesors(self, batch: dict, logits: torch.Tensor) -> torch.Tensor:
+        for postprocessor in self._postprocessors:
+            logits = postprocessor.on_prediction(batch, logits)
+        return logits
+
+    @abc.abstractmethod
+    def _ids_to_result(
+        self,
+        query_ids: torch.Tensor,
+        item_ids: torch.Tensor,
+        item_scores: torch.Tensor,
+    ) -> _T:  # pragma: no cover
+        pass
+
+
+class PandasTopItemsCallback(TopItemsCallbackBase[PandasDataFrame]):
+    """
+    A callback that records the result of the model's forward function at the inference stage in a Pandas Dataframe.
+    """
+
+    def _ids_to_result(
+        self,
+        query_ids: torch.Tensor,
+        item_ids: torch.Tensor,
+        item_scores: torch.Tensor,
+    ) -> PandasDataFrame:
+        prediction = PandasDataFrame(
+            {
+                self.query_column: query_ids.flatten().cpu().numpy(),
+                self.item_column: list(item_ids.cpu().numpy()),
+                self.rating_column: list(item_scores.cpu().numpy()),
+            }
+        )
+        return prediction.explode([self.item_column, self.rating_column])
+
+
+class PolarsTopItemsCallback(TopItemsCallbackBase[PolarsDataFrame]):
+    """
+    A callback that records the result of the model's forward function at the inference stage in a Polars Dataframe.
+    """
+
+    def _ids_to_result(
+        self,
+        query_ids: torch.Tensor,
+        item_ids: torch.Tensor,
+        item_scores: torch.Tensor,
+    ) -> PolarsDataFrame:
+        prediction = PolarsDataFrame(
+            {
+                self.query_column: query_ids.flatten().cpu().numpy(),
+                self.item_column: list(item_ids.cpu().numpy()),
+                self.rating_column: list(item_scores.cpu().numpy()),
+            }
+        )
+        return prediction.explode([self.item_column, self.rating_column])
+
+
+class SparkTopItemsCallback(TopItemsCallbackBase[SparkDataFrame]):
+    """
+    A callback that records the result of the model's forward function at the inference stage in a Spark Dataframe.
+    """
+
+    def __init__(
+        self,
+        top_k: int,
+        query_column: str,
+        item_column: str,
+        rating_column: str,
+        spark_session: SparkSession,
+        postprocessors: Optional[list[PostprocessorBase]] = None,
+    ) -> None:
+        """
+        :param top_k: Take the ``top_k`` IDs with the highest logit values.
+        :param query_column: The name of the query column in the resulting dataframe.
+        :param item_column: The name of the item column in the resulting dataframe.
+        :param rating_column: The name of the rating column in the resulting dataframe.
+            This column will contain the ``top_k`` items with the highest logit values.
+        :param spark_session: Spark session. Required to create a Spark DataFrame.
+        :param postprocessors: A list of postprocessors for modifying logits from the model
+            before sorting and taking top K ones.
+            For example, it can be a softmax operation to logits or set the ``-inf`` value for some IDs.
+            Default: ``None``.
+        """
+        super().__init__(
+            top_k=top_k,
+            query_column=query_column,
+            item_column=item_column,
+            rating_column=rating_column,
+            postprocessors=postprocessors,
+        )
+        self.spark_session = spark_session
+
+    def _ids_to_result(
+        self,
+        query_ids: torch.Tensor,
+        item_ids: torch.Tensor,
+        item_scores: torch.Tensor,
+    ) -> SparkDataFrame:
+        schema = (
+            StructType()
+            .add(self.query_column, IntegerType(), False)
+            .add(self.item_column, ArrayType(IntegerType()), False)
+            .add(self.rating_column, ArrayType(DoubleType()), False)
+        )
+        prediction = (
+            self.spark_session.createDataFrame(
+                data=list(
+                    zip(
+                        query_ids.flatten().cpu().numpy().tolist(),
+                        item_ids.cpu().numpy().tolist(),
+                        item_scores.cpu().numpy().tolist(),
+                    )
+                ),
+                schema=schema,
+            )
+            .withColumn(
+                "exploded_columns",
+                sf.explode(sf.arrays_zip(self.item_column, self.rating_column)),
+            )
+            .select(
+                self.query_column,
+                f"exploded_columns.{self.item_column}",
+                f"exploded_columns.{self.rating_column}",
+            )
+        )
+        return prediction
+
+
+class TorchTopItemsCallback(TopItemsCallbackBase[tuple[torch.LongTensor, torch.LongTensor, torch.Tensor]]):
+    """
+    A callback that records the result of the model's forward function at the inference stage in a PyTorch Tensors.
+    """
+
+    def __init__(
+        self,
+        top_k: int,
+        postprocessors: Optional[list[PostprocessorBase]] = None,
+    ) -> None:
+        """
+        :param top_k: Take the ``top_k`` IDs with the highest logit values.
+        :param postprocessors: A list of postprocessors for modifying logits from the model
+            before sorting and taking top K.
+            For example, it can be a softmax operation to logits or set the ``-inf`` value for some IDs.
+            Default: ``None``.
+        """
+        super().__init__(
+            top_k=top_k,
+            query_column="query_id",
+            item_column="item_id",
+            rating_column="rating",
+            postprocessors=postprocessors,
+        )
+
+    def _ids_to_result(
+        self,
+        query_ids: torch.Tensor,
+        item_ids: torch.Tensor,
+        item_scores: torch.Tensor,
+    ) -> tuple[torch.LongTensor, torch.LongTensor, torch.Tensor]:
+        return (
+            query_ids.flatten().cpu().long(),
+            item_ids.cpu().long(),
+            item_scores.cpu(),
+        )
+
+
+class HiddenStatesCallback(lightning.Callback):
+    """
+    A callback for getting any hidden state from the model.
+
+    When applying this callback,
+    it is expected that the result of the model's forward function contains the ``hidden_states`` key.
+    """
+
+    def __init__(self, hidden_state_index: int):
+        """
+        :param hidden_state_index: It is expected that the result of the model's forward function
+            contains the ``hidden_states`` key. ``hidden_states`` key contains Tuple of PyTorch Tensors.
+            Therefore, to get a specific hidden state, you need to submit an index from this tuple.
+        """
+        self._hidden_state_index = hidden_state_index
+        self._embeddings_per_batch: list[torch.Tensor] = []
+
+    def on_predict_epoch_start(
+        self,
+        trainer: lightning.Trainer,  # noqa: ARG002
+        pl_module: LightningModule,  # noqa: ARG002
+    ) -> None:
+        self._embeddings_per_batch.clear()
+
+    def on_predict_batch_end(
+        self,
+        trainer: lightning.Trainer,  # noqa: ARG002
+        pl_module: LightningModule,  # noqa: ARG002
+        outputs: InferenceOutput,
+        batch: dict,  # noqa: ARG002
+        batch_idx: int,  # noqa: ARG002
+        dataloader_idx: int = 0,  # noqa: ARG002
+    ) -> None:
+        self._embeddings_per_batch.append(outputs["hidden_states"][self._hidden_state_index].detach().cpu())
+
+    def get_result(self):
+        """
+        :returns: Hidden states through all batches.
+        """
+        return torch.cat(self._embeddings_per_batch)