Feature/Negative sampling (#275)

Added negative sampling customization to transformer-based models

Author: spirinamayya

CHANGELOG.md

@@ -8,11 +8,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## Unreleased
 
 ### Added 
+- `TransformerNegativeSamplerBase` and `CatalogUniformSampler` classes, `negative_sampler_type` and `negative_sampler_kwargs` parameters to transformer-based models ([#275](https://github.com/MobileTeleSystems/RecTools/pull/275))
 - `SimilarityModuleBase`, `DistanceSimilarityModule`, similarity module to `TransformerTorchBackbone` parameters to transformer-based models `similarity_module_type`, `similarity_module_kwargs` ([#272](https://github.com/MobileTeleSystems/RecTools/pull/272))
 - `out_dim` property to `IdEmbeddingsItemNet`, `CatFeaturesItemNet` and `SumOfEmbeddingsConstructor` ([#276](https://github.com/MobileTeleSystems/RecTools/pull/276))
 - `TransformerBackboneBase`, `backbone_type` and `backbone_kwargs` parameters to transformer-based models ([#277](https://github.com/MobileTeleSystems/RecTools/pull/277))
 - `sampled_softmax` loss option for transformer models ([#274](https://github.com/MobileTeleSystems/RecTools/pull/274))
-
 ## [0.12.0] - 24.02.2025
 
 ### Added

rectools/models/nn/transformers/base.py

@@ -40,6 +40,7 @@
 )
 from .data_preparator import TransformerDataPreparatorBase
 from .lightning import TransformerLightningModule, TransformerLightningModuleBase
+from .negative_sampler import CatalogUniformSampler, TransformerNegativeSamplerBase
 from .net_blocks import (
     LearnableInversePositionalEncoding,
     PositionalEncodingBase,
@@ -128,6 +129,16 @@ def _serialize_type_sequence(obj: tp.Sequence[tp.Type]) -> tp.Tuple[str, ...]:
     ),
 ]
 
+TransformerNegativeSamplerType = tpe.Annotated[
+    tp.Type[TransformerNegativeSamplerBase],
+    BeforeValidator(_get_class_obj),
+    PlainSerializer(
+        func=get_class_or_function_full_path,
+        return_type=str,
+        when_used="json",
+    ),
+]
+
 
 ItemNetConstructorType = tpe.Annotated[
     tp.Type[ItemNetConstructorBase],
@@ -204,6 +215,7 @@ class TransformerModelConfig(ModelConfig):
     pos_encoding_type: PositionalEncodingType = LearnableInversePositionalEncoding
     transformer_layers_type: TransformerLayersType = PreLNTransformerLayers
     lightning_module_type: TransformerLightningModuleType = TransformerLightningModule
+    negative_sampler_type: TransformerNegativeSamplerType = CatalogUniformSampler
     similarity_module_type: SimilarityModuleType = DistanceSimilarityModule
     backbone_type: TransformerBackboneType = TransformerTorchBackbone
     get_val_mask_func: tp.Optional[ValMaskCallableSerialized] = None
@@ -262,6 +274,7 @@ def __init__(  # pylint: disable=too-many-arguments, too-many-locals
         item_net_constructor_type: tp.Type[ItemNetConstructorBase] = SumOfEmbeddingsConstructor,
         pos_encoding_type: tp.Type[PositionalEncodingBase] = LearnableInversePositionalEncoding,
         lightning_module_type: tp.Type[TransformerLightningModuleBase] = TransformerLightningModule,
+        negative_sampler_type: tp.Type[TransformerNegativeSamplerBase] = CatalogUniformSampler,
         similarity_module_type: tp.Type[SimilarityModuleBase] = DistanceSimilarityModule,
         backbone_type: tp.Type[TransformerBackboneBase] = TransformerTorchBackbone,
         get_val_mask_func: tp.Optional[ValMaskCallable] = None,
@@ -271,6 +284,7 @@ def __init__(  # pylint: disable=too-many-arguments, too-many-locals
         item_net_constructor_kwargs: tp.Optional[InitKwargs] = None,
         pos_encoding_kwargs: tp.Optional[InitKwargs] = None,
         lightning_module_kwargs: tp.Optional[InitKwargs] = None,
+        negative_sampler_kwargs: tp.Optional[InitKwargs] = None,
         similarity_module_kwargs: tp.Optional[InitKwargs] = None,
         backbone_kwargs: tp.Optional[InitKwargs] = None,
         **kwargs: tp.Any,
@@ -302,6 +316,7 @@ def __init__(  # pylint: disable=too-many-arguments, too-many-locals
         self.item_net_constructor_type = item_net_constructor_type
         self.pos_encoding_type = pos_encoding_type
         self.lightning_module_type = lightning_module_type
+        self.negative_sampler_type = negative_sampler_type
         self.backbone_type = backbone_type
         self.get_val_mask_func = get_val_mask_func
         self.get_trainer_func = get_trainer_func
@@ -310,6 +325,7 @@ def __init__(  # pylint: disable=too-many-arguments, too-many-locals
         self.item_net_constructor_kwargs = item_net_constructor_kwargs
         self.pos_encoding_kwargs = pos_encoding_kwargs
         self.lightning_module_kwargs = lightning_module_kwargs
+        self.negative_sampler_kwargs = negative_sampler_kwargs
         self.similarity_module_kwargs = similarity_module_kwargs
         self.backbone_kwargs = backbone_kwargs
 
@@ -334,6 +350,7 @@ def _init_data_preparator(self) -> None:
             batch_size=self.batch_size,
             dataloader_num_workers=self.dataloader_num_workers,
             train_min_user_interactions=self.train_min_user_interactions,
+            negative_sampler=self._init_negative_sampler() if requires_negatives else None,
             n_negatives=self.n_negatives if requires_negatives else None,
             get_val_mask_func=self.get_val_mask_func,
             shuffle_train=True,
@@ -355,6 +372,12 @@ def _init_trainer(self) -> None:
         else:
             self._trainer = self.get_trainer_func()
 
+    def _init_negative_sampler(self) -> TransformerNegativeSamplerBase:
+        return self.negative_sampler_type(
+            n_negatives=self.n_negatives,
+            **self._get_kwargs(self.negative_sampler_kwargs),
+        )
+
     def _construct_item_net(self, dataset: Dataset) -> ItemNetBase:
         return self.item_net_constructor_type.from_dataset(
             dataset,

rectools/models/nn/transformers/bert4rec.py

@@ -38,6 +38,7 @@
 )
 from .constants import MASKING_VALUE, PADDING_VALUE
 from .data_preparator import TransformerDataPreparatorBase
+from .negative_sampler import CatalogUniformSampler, TransformerNegativeSamplerBase
 from .net_blocks import (
     LearnableInversePositionalEncoding,
     PositionalEncodingBase,
@@ -49,7 +50,29 @@
 
 
 class BERT4RecDataPreparator(TransformerDataPreparatorBase):
-    """Data Preparator for BERT4RecModel."""
+    """Data Preparator for BERT4RecModel.
+
+    Parameters
+    ----------
+    session_max_len : int
+        Maximum length of user sequence.
+    batch_size : int
+        How many samples per batch to load.
+    dataloader_num_workers : int
+        Number of loader worker processes.
+    shuffle_train : bool, default True
+        If ``True``, reshuffles data at each epoch.
+    train_min_user_interactions : int, default 2
+        Minimum length of user sequence. Cannot be less than 2.
+    get_val_mask_func : Callable, default None
+        Function to get validation mask.
+    n_negatives : optional(int), default ``None``
+        Number of negatives for BCE, gBCE and sampled_softmax losses.
+    negative_sampler: optional(TransformerNegativeSamplerBase), default ``None``
+        Negative sampler.
+    mask_prob : float, default 0.15
+        Probability of masking an item in interactions sequence.
+    """
 
     train_session_max_len_addition: int = 0
     item_extra_tokens: tp.Sequence[Hashable] = (PADDING_VALUE, MASKING_VALUE)
@@ -61,6 +84,7 @@ def __init__(
         batch_size: int,
         dataloader_num_workers: int,
         train_min_user_interactions: int,
+        negative_sampler: tp.Optional[TransformerNegativeSamplerBase] = None,
         mask_prob: float = 0.15,
         shuffle_train: bool = True,
         get_val_mask_func: tp.Optional[ValMaskCallable] = None,
@@ -69,6 +93,7 @@ def __init__(
         super().__init__(
             session_max_len=session_max_len,
             n_negatives=n_negatives,
+            negative_sampler=negative_sampler,
             batch_size=batch_size,
             dataloader_num_workers=dataloader_num_workers,
             train_min_user_interactions=train_min_user_interactions,
@@ -119,13 +144,10 @@ def _collate_fn_train(
             yw[i, -len(ses) :] = ses_weights  # ses_weights: [session_len] -> yw[i]: [session_max_len]
 
         batch_dict = {"x": torch.LongTensor(x), "y": torch.LongTensor(y), "yw": torch.FloatTensor(yw)}
-        if self.n_negatives is not None:
-            negatives = torch.randint(
-                low=self.n_item_extra_tokens,
-                high=self.item_id_map.size,
-                size=(batch_size, self.session_max_len, self.n_negatives),
-            )  # [batch_size, session_max_len, n_negatives]
-            batch_dict["negatives"] = negatives
+        if self.negative_sampler is not None:
+            batch_dict["negatives"] = self.negative_sampler.get_negatives(
+                batch_dict, lowest_id=self.n_item_extra_tokens, highest_id=self.item_id_map.size
+            )
         return batch_dict
 
     def _collate_fn_val(self, batch: List[Tuple[List[int], List[float]]]) -> Dict[str, torch.Tensor]:
@@ -147,13 +169,10 @@ def _collate_fn_val(self, batch: List[Tuple[List[int], List[float]]]) -> Dict[st
             yw[i, -1:] = ses_weights[target_idx]  # yw[i]: [1]
 
         batch_dict = {"x": torch.LongTensor(x), "y": torch.LongTensor(y), "yw": torch.FloatTensor(yw)}
-        if self.n_negatives is not None:
-            negatives = torch.randint(
-                low=self.n_item_extra_tokens,
-                high=self.item_id_map.size,
-                size=(batch_size, 1, self.n_negatives),
-            )  # [batch_size, 1, n_negatives]
-            batch_dict["negatives"] = negatives
+        if self.negative_sampler is not None:
+            batch_dict["negatives"] = self.negative_sampler.get_negatives(
+                batch_dict, lowest_id=self.n_item_extra_tokens, highest_id=self.item_id_map.size, session_len_limit=1
+            )
         return batch_dict
 
     def _collate_fn_recommend(self, batch: List[Tuple[List[int], List[float]]]) -> Dict[str, torch.Tensor]:
@@ -213,7 +232,7 @@ class BERT4RecModel(TransformerModelBase[BERT4RecModelConfig]):
     loss : {"softmax", "BCE", "gBCE", "sampled_softmax"}, default "softmax"
         Loss function.
     n_negatives : int, default 1
-        Number of negatives for BCE and gBCE losses.
+        Number of negatives for BCE, gBCE and sampled_softmax losses.
     gbce_t : float, default 0.2
         Calibration parameter for gBCE loss.
     lr : float, default 0.001
@@ -258,6 +277,8 @@ class BERT4RecModel(TransformerModelBase[BERT4RecModelConfig]):
         Type of data preparator used for dataset processing and dataloader creation.
     lightning_module_type : type(TransformerLightningModuleBase), default `TransformerLightningModule`
         Type of lightning module defining training procedure.
+    negative_sampler_type: type(TransformerNegativeSamplerBase), default `CatalogUniformSampler`
+        Type of negative sampler.
     similarity_module_type : type(SimilarityModuleBase), default `DistanceSimilarityModule`
         Type of similarity module.
     backbone_type : type(TransformerBackboneBase), default `TransformerTorchBackbone`
@@ -295,6 +316,9 @@ class BERT4RecModel(TransformerModelBase[BERT4RecModelConfig]):
     lightning_module_kwargs: optional(dict), default ``None``
         Additional keyword arguments to pass during `lightning_module_type` initialization.
         Make sure all dict values have JSON serializable types.
+    negative_sampler_kwargs: optional(dict), default ``None``
+        Additional keyword arguments to pass during `negative_sampler_type` initialization.
+        Make sure all dict values have JSON serializable types.
     similarity_module_kwargs: optional(dict), default ``None``
         Additional keyword arguments to pass during `similarity_module_type` initialization.
         Make sure all dict values have JSON serializable types.
@@ -332,6 +356,7 @@ def __init__(  # pylint: disable=too-many-arguments, too-many-locals
         transformer_layers_type: tp.Type[TransformerLayersBase] = PreLNTransformerLayers,
         data_preparator_type: tp.Type[TransformerDataPreparatorBase] = BERT4RecDataPreparator,
         lightning_module_type: tp.Type[TransformerLightningModuleBase] = TransformerLightningModule,
+        negative_sampler_type: tp.Type[TransformerNegativeSamplerBase] = CatalogUniformSampler,
         similarity_module_type: tp.Type[SimilarityModuleBase] = DistanceSimilarityModule,
         backbone_type: tp.Type[TransformerBackboneBase] = TransformerTorchBackbone,
         get_val_mask_func: tp.Optional[ValMaskCallable] = None,
@@ -346,6 +371,7 @@ def __init__(  # pylint: disable=too-many-arguments, too-many-locals
         item_net_constructor_kwargs: tp.Optional[InitKwargs] = None,
         pos_encoding_kwargs: tp.Optional[InitKwargs] = None,
         lightning_module_kwargs: tp.Optional[InitKwargs] = None,
+        negative_sampler_kwargs: tp.Optional[InitKwargs] = None,
         similarity_module_kwargs: tp.Optional[InitKwargs] = None,
         backbone_kwargs: tp.Optional[InitKwargs] = None,
     ):
@@ -381,6 +407,7 @@ def __init__(  # pylint: disable=too-many-arguments, too-many-locals
             item_net_constructor_type=item_net_constructor_type,
             pos_encoding_type=pos_encoding_type,
             lightning_module_type=lightning_module_type,
+            negative_sampler_type=negative_sampler_type,
             backbone_type=backbone_type,
             get_val_mask_func=get_val_mask_func,
             get_trainer_func=get_trainer_func,
@@ -390,14 +417,17 @@ def __init__(  # pylint: disable=too-many-arguments, too-many-locals
             item_net_constructor_kwargs=item_net_constructor_kwargs,
             pos_encoding_kwargs=pos_encoding_kwargs,
             lightning_module_kwargs=lightning_module_kwargs,
+            negative_sampler_kwargs=negative_sampler_kwargs,
             similarity_module_kwargs=similarity_module_kwargs,
             backbone_kwargs=backbone_kwargs,
         )
 
     def _init_data_preparator(self) -> None:
+        requires_negatives = self.lightning_module_type.requires_negatives(self.loss)
         self.data_preparator: TransformerDataPreparatorBase = self.data_preparator_type(
             session_max_len=self.session_max_len,
-            n_negatives=self.n_negatives if self.loss != "softmax" else None,
+            n_negatives=self.n_negatives if requires_negatives else None,
+            negative_sampler=self._init_negative_sampler() if requires_negatives else None,
             batch_size=self.batch_size,
             dataloader_num_workers=self.dataloader_num_workers,
             train_min_user_interactions=self.train_min_user_interactions,

rectools/models/nn/transformers/data_preparator.py

@@ -29,6 +29,7 @@
 from rectools.dataset.identifiers import IdMap
 
 from .constants import PADDING_VALUE
+from .negative_sampler import TransformerNegativeSamplerBase
 
 
 class SequenceDataset(TorchDataset):
@@ -104,6 +105,10 @@ class TransformerDataPreparatorBase:
         Minimum length of user sequence. Cannot be less than 2.
     get_val_mask_func : Callable, default None
         Function to get validation mask.
+    n_negatives : optional(int), default ``None``
+        Number of negatives for BCE, gBCE and sampled_softmax losses.
+    negative_sampler: optional(TransformerNegativeSamplerBase), default ``None``
+        Negative sampler.
     """
 
     # We sometimes need data preparators to add +1 to actual session_max_len
@@ -119,15 +124,17 @@ def __init__(
         dataloader_num_workers: int,
         shuffle_train: bool = True,
         train_min_user_interactions: int = 2,
-        n_negatives: tp.Optional[int] = None,
         get_val_mask_func: tp.Optional[tp.Callable] = None,
+        n_negatives: tp.Optional[int] = None,
+        negative_sampler: tp.Optional[TransformerNegativeSamplerBase] = None,
         **kwargs: tp.Any,
     ) -> None:
         self.item_id_map: IdMap
         self.extra_token_ids: tp.Dict
         self.train_dataset: Dataset
         self.val_interactions: tp.Optional[pd.DataFrame] = None
         self.session_max_len = session_max_len
+        self.negative_sampler = negative_sampler
         self.n_negatives = n_negatives
         self.batch_size = batch_size
         self.dataloader_num_workers = dataloader_num_workers
@@ -189,6 +196,7 @@ def process_dataset_train(self, dataset: Dataset) -> None:
         if self.get_val_mask_func is not None:
             val_mask = self.get_val_mask_func(raw_interactions)
             interactions = raw_interactions[~val_mask]
+            interactions.reset_index(drop=True, inplace=True)
 
         # Filter train interactions
         interactions = self._filter_train_interactions(interactions)

rectools/models/nn/transformers/lightning.py

@@ -40,8 +40,18 @@ class TransformerLightningModuleBase(LightningModule):  # pylint: disable=too-ma
     ----------
     torch_model : TransformerBackboneBase
         Torch model to make recommendations.
+    model_config: Dict[str, Any]
+        Model config.
+    dataset_schema: DatasetSchemaDict
+        Dataset schema.
+    item_external_ids: ExternalIds
+        External item ids from train dataset.
+    item_extra_tokens : Sequence(Hashable)
+        Elements used for sequence padding.
     lr : float
         Learning rate.
+    gbce_t : float
+        Calibration parameter for gBCE loss.
     loss : str, default "softmax"
         Loss function.
     adam_betas : Tuple[float, float], default (0.9, 0.98)
@@ -240,7 +250,37 @@ def _recommend_i2i(
 
 
 class TransformerLightningModule(TransformerLightningModuleBase):
-    """Lightning module to train transformer models."""
+    """Lightning module to train transformer models.
+
+    Parameters
+    ----------
+    torch_model : TransformerBackboneBase
+        Torch model to make recommendations.
+    model_config: Dict[str, Any]
+        Model config.
+    dataset_schema: DatasetSchemaDict
+        Dataset schema.
+    item_external_ids: ExternalIds
+        External item ids from train dataset.
+    item_extra_tokens : Sequence(Hashable)
+        Elements used for sequence padding.
+    lr : float
+        Learning rate.
+    gbce_t : float
+        Calibration parameter for gBCE loss.
+    loss : str, default "softmax"
+        Loss function.
+    adam_betas : Tuple[float, float], default (0.9, 0.98)
+        Coefficients for running averages of gradient and its square.
+    data_preparator : TransformerDataPreparatorBase
+        Data preparator.
+    verbose : int, default 0
+        Verbosity level.
+    train_loss_name : str, default "train_loss"
+        Name of the training loss.
+    val_loss_name : str, default "val_loss"
+        Name of the training loss.
+    """
 
     i2i_dist = Distance.COSINE
 
@@ -296,7 +336,7 @@ def validation_step(self, batch: tp.Dict[str, torch.Tensor], batch_idx: int) ->
             type_logits = "pos_neg_logits" if self._requires_negatives else "logits"
             outputs = {
                 "loss": loss,
-                type_logits: logits,
+                type_logits: logits.squeeze(),
             }
         else:
             outputs = self._calc_custom_loss_outputs(batch, batch_idx)  # pragma: no cover
@@ -339,7 +379,7 @@ def _get_user_item_embeddings(
             for batch in recommend_dataloader:
                 batch = {k: v.to(device) for k, v in batch.items()}
                 batch_embs = self.torch_model.encode_sessions(batch, item_embs)[:, -1, :]
-                user_embs.append(batch_embs)
+                user_embs.append(batch_embs.cpu())
 
         return torch.cat(user_embs), item_embs