replay-rec 0.18.0__py3-none-any.whl → 0.18.1__py3-none-any.whl

replay/models/nn/sequential/bert4rec/lightning.py

@@ -31,7 +31,7 @@ class Bert4Rec(lightning.LightningModule):
         loss_sample_count: Optional[int] = None,
         negative_sampling_strategy: str = "global_uniform",
         negatives_sharing: bool = False,
-        optimizer_factory: Optional[OptimizerFactory] = None,
+        optimizer_factory: OptimizerFactory = FatOptimizerFactory(),
         lr_scheduler_factory: Optional[LRSchedulerFactory] = None,
     ):
         """
@@ -65,7 +65,7 @@ class Bert4Rec(lightning.LightningModule):
         :param negatives_sharing: Apply negative sharing in calculating sampled logits.
             Default: ``False``.
         :param optimizer_factory: Optimizer factory.
-            Default: ``None``.
+            Default: ``FatOptimizerFactory``.
         :param lr_scheduler_factory: Learning rate schedule factory.
             Default: ``None``.
         """
@@ -95,6 +95,7 @@ class Bert4Rec(lightning.LightningModule):
         item_count = tensor_schema.item_id_features.item().cardinality
         assert item_count
         self._vocab_size = item_count
+        self.candidates_to_score = None
 
     def training_step(self, batch: Bert4RecTrainingBatch, batch_idx: int) -> torch.Tensor:  # noqa: ARG002
         """
@@ -107,33 +108,51 @@ class Bert4Rec(lightning.LightningModule):
         self.log("train_loss", loss, on_step=True, on_epoch=True, prog_bar=True, sync_dist=True)
         return loss
 
+    def predict_step(
+        self, batch: Bert4RecPredictionBatch, batch_idx: int, dataloader_idx: int = 0  # noqa: ARG002
+    ) -> torch.Tensor:
+        """
+        :param batch (Bert4RecPredictionBatch): Batch of prediction data.
+        :param batch_idx (int): Batch index.
+        :param dataloader_idx (int): Dataloader index.
+
+        :returns: Calculated scores on prediction batch.
+        """
+        batch = _prepare_prediction_batch(self._schema, self._model.max_len, batch)
+        return self._model_predict(batch.features, batch.padding_mask, batch.tokens_mask)
+
+    def predict(
+        self,
+        batch: Bert4RecPredictionBatch,
+        candidates_to_score: Optional[torch.LongTensor] = None,
+    ) -> torch.Tensor:
+        """
+        :param batch (Bert4RecPredictionBatch): Batch of prediction data.
+        :param candidates_to_score: Item ids to calculate scores.
+            Default: ``None``.
+
+        :returns: Calculated scores on prediction batch.
+        """
+        batch = _prepare_prediction_batch(self._schema, self._model.max_len, batch)
+        return self._model_predict(batch.features, batch.padding_mask, batch.tokens_mask, candidates_to_score)
+
     def forward(
         self,
         feature_tensors: TensorMap,
         padding_mask: torch.BoolTensor,
         tokens_mask: torch.BoolTensor,
+        candidates_to_score: Optional[torch.LongTensor] = None,
     ) -> torch.Tensor:  # pragma: no cover
         """
         :param feature_tensors:  Batch of features.
         :param padding_mask: Padding mask where 0 - <PAD>, 1 otherwise.
         :param tokens_mask: Token mask where 0 - <MASK> tokens, 1 otherwise.
+        :param candidates_to_score: Item ids to calculate scores.
+            Default: ``None``.
 
         :returns: Calculated scores.
         """
-        return self._model_predict(feature_tensors, padding_mask, tokens_mask)
-
-    def predict_step(
-        self, batch: Bert4RecPredictionBatch, batch_idx: int, dataloader_idx: int = 0  # noqa: ARG002
-    ) -> torch.Tensor:
-        """
-        :param batch (Bert4RecPredictionBatch): Batch of prediction data.
-        :param batch_idx (int): Batch index.
-        :param dataloader_idx (int): Dataloader index.
-
-        :returns: Calculated scores on prediction batch.
-        """
-        batch = self._prepare_prediction_batch(batch)
-        return self._model_predict(batch.features, batch.padding_mask, batch.tokens_mask)
+        return self._model_predict(feature_tensors, padding_mask, tokens_mask, candidates_to_score)
 
     def validation_step(
         self, batch: Bert4RecValidationBatch, batch_idx: int, dataloader_idx: int = 0  # noqa: ARG002
@@ -150,8 +169,7 @@ class Bert4Rec(lightning.LightningModule):
         """
         :returns: Configured optimizer and lr scheduler.
         """
-        optimizer_factory = self._optimizer_factory or FatOptimizerFactory()
-        optimizer = optimizer_factory.create(self._model.parameters())
+        optimizer = self._optimizer_factory.create(self._model.parameters())
 
         if self._lr_scheduler_factory is None:
             return optimizer
@@ -159,49 +177,20 @@ class Bert4Rec(lightning.LightningModule):
         lr_scheduler = self._lr_scheduler_factory.create(optimizer)
         return [optimizer], [lr_scheduler]
 
-    def _prepare_prediction_batch(self, batch: Bert4RecPredictionBatch) -> Bert4RecPredictionBatch:
-        if batch.padding_mask.shape[1] > self._model.max_len:
-            msg = f"The length of the submitted sequence \
-                must not exceed the maximum length of the sequence. \
-                The length of the sequence is given {batch.padding_mask.shape[1]}, \
-                while the maximum length is {self._model.max_len}"
-            raise ValueError(msg)
-
-        if batch.padding_mask.shape[1] < self._model.max_len:
-            query_id, padding_mask, features, _ = batch
-            sequence_item_count = padding_mask.shape[1]
-            for feature_name, feature_tensor in features.items():
-                if self._schema[feature_name].is_cat:
-                    features[feature_name] = torch.nn.functional.pad(
-                        feature_tensor, (self._model.max_len - sequence_item_count, 0), value=0
-                    )
-                else:
-                    features[feature_name] = torch.nn.functional.pad(
-                        feature_tensor.view(feature_tensor.size(0), feature_tensor.size(1)),
-                        (self._model.max_len - sequence_item_count, 0),
-                        value=0,
-                    ).unsqueeze(-1)
-            padding_mask = torch.nn.functional.pad(
-                padding_mask, (self._model.max_len - sequence_item_count, 0), value=0
-            )
-            shifted_features, shifted_padding_mask, tokens_mask = _shift_features(self._schema, features, padding_mask)
-            batch = Bert4RecPredictionBatch(query_id, shifted_padding_mask, shifted_features, tokens_mask)
-        return batch
-
     def _model_predict(
         self,
         feature_tensors: TensorMap,
         padding_mask: torch.BoolTensor,
         tokens_mask: torch.BoolTensor,
+        candidates_to_score: torch.LongTensor = None,
     ) -> torch.Tensor:
         model: Bert4RecModel
-        if isinstance(self._model, torch.nn.DataParallel):
-            model = cast(Bert4RecModel, self._model.module)  # multigpu
-        else:
-            model = self._model
-        scores = model(feature_tensors, padding_mask, tokens_mask)
-        candidate_scores = scores[:, -1, :]
-        return candidate_scores
+        model = (
+            cast(Bert4RecModel, self._model.module) if isinstance(self._model, torch.nn.DataParallel) else self._model
+        )
+        candidates_to_score = self.candidates_to_score if candidates_to_score is None else candidates_to_score
+        scores = model.predict(feature_tensors, padding_mask, tokens_mask, candidates_to_score)
+        return scores
 
     def _compute_loss(self, batch: Bert4RecTrainingBatch) -> torch.Tensor:
         if self._loss_type == "BCE":
@@ -504,6 +493,50 @@ class Bert4Rec(lightning.LightningModule):
 
         self._set_new_item_embedder_to_model(weights_new, new_vocab_size)
 
+    @property
+    def optimizer_factory(self) -> OptimizerFactory:
+        """
+        Returns current optimizer_factory.
+        """
+        return self._optimizer_factory
+
+    @optimizer_factory.setter
+    def optimizer_factory(self, optimizer_factory: OptimizerFactory) -> None:
+        """
+        Sets new optimizer_factory.
+        :param optimizer_factory: New optimizer factory.
+        """
+        if isinstance(optimizer_factory, OptimizerFactory):
+            self._optimizer_factory = optimizer_factory
+        else:
+            msg = f"Expected optimizer_factory of type OptimizerFactory, got {type(optimizer_factory)}"
+            raise ValueError(msg)
+
+    @property
+    def candidates_to_score(self) -> Union[torch.LongTensor, None]:
+        """
+        Returns tensor of item ids to calculate scores.
+        """
+        return self._candidates_to_score
+
+    @candidates_to_score.setter
+    def candidates_to_score(self, candidates: Optional[torch.LongTensor] = None) -> None:
+        """
+        Sets tensor of item ids to calculate scores.
+        :param candidates: Tensor of item ids to calculate scores.
+        """
+        total_item_count = self._model.item_count
+        if isinstance(candidates, torch.Tensor) and candidates.dtype is torch.long:
+            if 0 < candidates.shape[0] <= total_item_count:
+                self._candidates_to_score = candidates
+            else:
+                msg = f"Expected candidates length to be between 1 and {total_item_count=}"
+                raise ValueError(msg)
+        elif candidates is not None:
+            msg = f"Expected candidates to be of type torch.LongTensor or None, gpt {type(candidates)}"
+            raise ValueError(msg)
+        self._candidates_to_score = candidates
+
     def _set_new_item_embedder_to_model(self, weights_new: torch.nn.Embedding, new_vocab_size: int):
         self._model.item_embedder.cat_embeddings[self._model.schema.item_id_feature_name] = weights_new
         if self._model.enable_embedding_tying is True:
@@ -521,3 +554,37 @@ class Bert4Rec(lightning.LightningModule):
         self._vocab_size = new_vocab_size
         self._model.item_count = new_vocab_size
         self._schema.item_id_features[self._schema.item_id_feature_name]._set_cardinality(new_vocab_size)
+
+
+def _prepare_prediction_batch(
+    schema: TensorSchema, max_len: int, batch: Bert4RecPredictionBatch
+) -> Bert4RecPredictionBatch:
+    if batch.padding_mask.shape[1] > max_len:
+        msg = (
+            f"The length of the submitted sequence "
+            "must not exceed the maximum length of the sequence. "
+            f"The length of the sequence is given {batch.padding_mask.shape[1]}, "
+            f"while the maximum length is {max_len}"
+        )
+        raise ValueError(msg)
+
+    if batch.padding_mask.shape[1] < max_len:
+        query_id, padding_mask, features, _ = batch
+        sequence_item_count = padding_mask.shape[1]
+        for feature_name, feature_tensor in features.items():
+            if schema[feature_name].is_cat:
+                features[feature_name] = torch.nn.functional.pad(
+                    feature_tensor,
+                    (max_len - sequence_item_count, 0),
+                    value=schema[feature_name].padding_value,
+                )
+            else:
+                features[feature_name] = torch.nn.functional.pad(
+                    feature_tensor.view(feature_tensor.size(0), feature_tensor.size(1)),
+                    (max_len - sequence_item_count, 0),
+                    value=schema[feature_name].padding_value,
+                ).unsqueeze(-1)
+        padding_mask = torch.nn.functional.pad(padding_mask, (max_len - sequence_item_count, 0), value=0)
+        shifted_features, shifted_padding_mask, tokens_mask = _shift_features(schema, features, padding_mask)
+        batch = Bert4RecPredictionBatch(query_id, shifted_padding_mask, shifted_features, tokens_mask)
+    return batch

replay/models/nn/sequential/bert4rec/model.py

@@ -98,6 +98,27 @@ class Bert4RecModel(torch.nn.Module):
 
         return all_scores  # [B x L x E]
 
+    def predict(
+        self,
+        inputs: TensorMap,
+        pad_mask: torch.BoolTensor,
+        token_mask: torch.BoolTensor,
+        candidates_to_score: Optional[torch.LongTensor] = None,
+    ) -> torch.Tensor:
+        """
+        :param inputs: Batch of features.
+        :param pad_mask: Padding mask where 0 - <PAD>, 1 otherwise.
+        :param token_mask: Token mask where 0 - <MASK> tokens, 1 otherwise.
+        :param candidates_to_score: Item ids to calculate scores.
+            if `None` predicts for all items
+
+        :returns: Calculated scores among canditates_to_score items.
+        """
+        # final_emb: [B x E]
+        final_emb = self.get_query_embeddings(inputs, pad_mask, token_mask)
+        candidate_scores = self.get_logits(final_emb, candidates_to_score)
+        return candidate_scores
+
     def forward_step(self, inputs: TensorMap, pad_mask: torch.BoolTensor, token_mask: torch.BoolTensor) -> torch.Tensor:
         """
 

replay/models/nn/sequential/callbacks/prediction_callbacks.py

@@ -64,6 +64,11 @@ class BasePredictionCallback(lightning.Callback, Generic[_T]):
         self._item_batches.clear()
         self._item_scores.clear()
 
+        candidates = trainer.model.candidates_to_score if hasattr(trainer.model, "candidates_to_score") else None
+        for postprocessor in self._postprocessors:
+            if hasattr(postprocessor, "candidates"):
+                postprocessor.candidates = candidates
+
     def on_predict_batch_end(
         self,
         trainer: lightning.Trainer,  # noqa: ARG002
@@ -88,7 +93,6 @@ class BasePredictionCallback(lightning.Callback, Generic[_T]):
             torch.cat(self._item_batches),
             torch.cat(self._item_scores),
         )
-
         return prediction
 
     def _compute_pipeline(

replay/models/nn/sequential/compiled/__init__.py

@@ -0,0 +1,5 @@
+from replay.utils import OPENVINO_AVAILABLE
+
+if OPENVINO_AVAILABLE:
+    from .bert4rec_compiled import Bert4RecCompiled
+    from .sasrec_compiled import SasRecCompiled

replay/models/nn/sequential/compiled/base_compiled_model.py

@@ -0,0 +1,261 @@
+import pathlib
+import tempfile
+from abc import abstractmethod
+from typing import Any, Dict, List, Literal, Optional, Tuple, Union
+
+import lightning
+import openvino as ov
+import torch
+
+from replay.data.nn import TensorSchema
+
+OptimizedModeType = Literal[
+    "batch",
+    "one_query",
+    "dynamic_batch_size",
+]
+
+
+def _compile_openvino(
+    onnx_path: str,
+    batch_size: int,
+    max_seq_len: int,
+    num_candidates_to_score: int,
+    num_threads: Optional[int],
+) -> ov.CompiledModel:
+    """
+    Method defines compilation strategy for openvino backend.
+
+    :param onnx_path: Path to the model representation in ONNX format.
+    :param batch_size: Defines whether batch will be static or dynamic length.
+    :param max_seq_len: Defines whether sequence will be static or dynamic length.
+    :param num_candidates_to_score: Defines whether candidates will be static or dynamic length.
+    :param num_threads: Defines number of CPU threads for which the model will be compiled by the OpenVino core.
+        If ``None``, then compiler will set this parameter automatically.
+        Default: ``None``.
+
+    :return: Compiled model.
+    """
+    core = ov.Core()
+    if num_threads is not None:
+        core.set_property("CPU", {"INFERENCE_NUM_THREADS": num_threads})
+    model_onnx = core.read_model(model=onnx_path)
+    inputs_names = [inputs.names.pop() for inputs in model_onnx.inputs]
+    del model_onnx
+
+    candidates_input_id = len(inputs_names) - 1 if num_candidates_to_score is not None else len(inputs_names)
+    model_input_scheme = [(input_name, [batch_size, max_seq_len]) for input_name in inputs_names[:candidates_input_id]]
+    if num_candidates_to_score is not None:
+        model_input_scheme += [(inputs_names[candidates_input_id], [num_candidates_to_score])]
+    model_onnx = ov.convert_model(onnx_path, input=model_input_scheme)
+    return core.compile_model(model=model_onnx, device_name="CPU")
+
+
+class BaseCompiledModel:
+    """
+    Base class of CPU-optimized model for inference via OpenVINO.
+    It is recommended to use inherited classes and not to use this one.
+    """
+
+    def __init__(
+        self,
+        compiled_model: ov.CompiledModel,
+        schema: TensorSchema,
+    ) -> None:
+        """
+        :param compiled_model: Compiled model.
+        :param schema: Tensor schema of SasRec model.
+        """
+        self._batch_size: int
+        self._max_seq_len: int
+        self._inputs_names: List[str]
+        self._output_name: str
+
+        self._set_inner_params_from_openvino_model(compiled_model)
+        self._schema = schema
+        self._model = compiled_model
+
+    @abstractmethod
+    def predict(
+        self,
+        batch: Any,
+        candidates_to_score: Optional[torch.LongTensor] = None,
+    ) -> torch.Tensor:
+        """
+        Inference on one batch.
+
+        :param batch: Prediction input.
+        :param candidates_to_score: Item ids to calculate scores.
+            Default: ``None``.
+
+        :return: Tensor with scores.
+        """
+
+    def _validate_candidates_to_score(self, candidates: torch.LongTensor) -> None:
+        """Check if candidates param has proper type"""
+
+        if not (isinstance(candidates, torch.Tensor) and candidates.dtype is torch.long):
+            msg = (
+                "Expected candidates to be of type ``torch.Tensor`` with dtype ``torch.long``, "
+                f"got {type(candidates)} with dtype {candidates.dtype}."
+            )
+            raise ValueError(msg)
+
+    def _valilade_predict_input(self, batch: Any, candidates_to_score: Optional[torch.LongTensor] = None) -> None:
+        if self._num_candidates_to_score is None and candidates_to_score is not None:
+            msg = (
+                "If ``num_candidates_to_score`` is None, "
+                "it is impossible to infer the model with passed ``candidates_to_score``."
+            )
+            raise ValueError(msg)
+
+        if self._batch_size != -1 and batch.padding_mask.shape[0] != self._batch_size:
+            msg = (
+                f"The batch is smaller then defined batch_size={self._batch_size}. "
+                "It is impossible to infer the model with dynamic batch size in ``mode`` = ``batch``. "
+                "Use ``mode`` = ``dynamic_batch_size``."
+            )
+            raise ValueError(msg)
+
+    def _set_inner_params_from_openvino_model(self, compiled_model: ov.CompiledModel) -> None:
+        """Set params for ``predict`` method"""
+
+        input_scheme = compiled_model.inputs
+        self._batch_size = input_scheme[0].partial_shape[0].max_length
+        self._max_seq_len = input_scheme[0].partial_shape[1].max_length
+        self._inputs_names = [input.names.pop() for input in compiled_model.inputs]
+        if "candidates_to_score" in self._inputs_names:
+            self._num_candidates_to_score = input_scheme[-1].partial_shape[0].max_length
+        else:
+            self._num_candidates_to_score = None
+        self._output_name = compiled_model.output().names.pop()
+
+    @staticmethod
+    def _validate_num_candidates_to_score(num_candidates: Union[int, None]) -> Union[int, None]:
+        """Check if num_candidates param is proper"""
+
+        if num_candidates is None:
+            return num_candidates
+        if isinstance(num_candidates, int) and (num_candidates == -1 or num_candidates >= 1):
+            return num_candidates
+
+        msg = (
+            "Expected num_candidates_to_score to be of type ``int``, equal to ``-1``, ``natural number`` or ``None``. "
+            f"Got {num_candidates}."
+        )
+        raise ValueError(msg)
+
+    @staticmethod
+    def _get_input_params(
+        mode: OptimizedModeType,
+        batch_size: Optional[int],
+        num_candidates_to_score: Optional[int],
+    ) -> None:
+        """Get params for model compilation according to compilation mode"""
+
+        if mode == "one_query":
+            batch_size = 1
+
+        if mode == "batch":
+            assert batch_size, f"{mode} mode requires `batch_size`"
+            batch_size = batch_size
+
+        if mode == "dynamic_batch_size":
+            batch_size = -1
+
+        num_candidates_to_score = num_candidates_to_score if num_candidates_to_score else None
+        return batch_size, num_candidates_to_score
+
+    @staticmethod
+    def _run_model_compilation(
+        lightning_model: lightning.LightningModule,
+        model_input_sample: Tuple[Union[torch.Tensor, Dict[str, torch.Tensor]]],
+        model_input_names: List[str],
+        model_dynamic_axes_in_input: Dict[str, Dict],
+        batch_size: int,
+        num_candidates_to_score: Union[int, None],
+        num_threads: Optional[int] = None,
+        onnx_path: Optional[str] = None,
+    ) -> ov.CompiledModel:
+        """
+        Model conversion into ONNX format and compilation with defined engine.
+
+        :param lightning_model: Lightning model to be compiled.
+        :param model_input_sample: An example of model input with proper data type.
+        :param model_input_names: Input tensor names.
+        :param model_dynamic_axes_in_input: Dynamic axes in input.
+        :param batch_size: Defines the size of the axis with index 0 in the input of the compiled model.
+        :param num_candidates_to_score: Defines the size of the candidates in the input of the compiled model.
+        :param num_threads: Number of CPU threads to use.
+            Must be a natural number or ``None``.
+            If ``None``, then compiler will set this parameter automatically.
+            Default: ``None``.
+        :param onnx_path: Save ONNX model to path, if defined.
+            Default: ``None``.
+
+        :return: Compiled model.
+        """
+        max_seq_len = lightning_model._model.max_len
+
+        if onnx_path is None:
+            is_saveble = False
+            onnx_file = tempfile.NamedTemporaryFile(suffix=".onnx")
+            onnx_path = onnx_file.name
+        else:
+            is_saveble = True
+
+        lightning_model.to_onnx(
+            onnx_path,
+            input_sample=model_input_sample,
+            export_params=True,
+            opset_version=torch.onnx._constants.ONNX_DEFAULT_OPSET,
+            do_constant_folding=True,
+            input_names=model_input_names,
+            output_names=["scores"],
+            dynamic_axes=model_dynamic_axes_in_input,
+        )
+        del lightning_model
+
+        compiled_model = _compile_openvino(onnx_path, batch_size, max_seq_len, num_candidates_to_score, num_threads)
+
+        if not is_saveble:
+            onnx_file.close()
+
+        return compiled_model
+
+    @classmethod
+    @abstractmethod
+    def compile(
+        cls,
+        model: Union[lightning.LightningModule, str, pathlib.Path],
+        mode: OptimizedModeType = "one_query",
+        batch_size: Optional[int] = None,
+        num_candidates_to_score: Optional[int] = None,
+        num_threads: Optional[int] = None,
+        onnx_path: Optional[str] = None,
+    ) -> "BaseCompiledModel":
+        """
+        Model compilation.
+
+        :param model: Path to lightning model saved in .ckpt format or the model object itself.
+        :param mode: Inference mode, defines shape of inputs.
+            Could be one of [``one_query``, ``batch``, ``dynamic_batch_size``].\n
+            ``one_query`` - sets input shape to [1, max_seq_len]\n
+            ``batch`` - sets input shape to [batch_size, max_seq_len]\n
+            ``dynamic_batch_size`` - sets batch_size to dynamic range [?, max_seq_len]\n
+            Default: ``one_query``.
+        :param batch_size: Batch size, required for ``batch`` mode.
+            Default: ``None``.
+        :param num_candidates_to_score: Number of item ids to calculate scores.
+            Could be one of [``None``, ``-1``, ``N``].\n
+            ``-1`` - sets candidates_to_score shape to dynamic range [1, ?]\n
+            ``N`` - sets candidates_to_score shape to [1, N]\n
+            ``None`` - disable candidates_to_score usage\n
+            Default: ``None``.
+        :param num_threads: Number of CPU threads to use.
+            Must be a natural number or ``None``.
+            If ``None``, then compiler will set this parameter automatically.
+            Default: ``None``.
+        :param onnx_path: Save ONNX model to path, if defined.
+            Default: ``None``.
+        """