Source code for

# Copyright (c) 2021, NVIDIA CORPORATION.
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Optional

import tensorflow as tf

from import EmbeddingWithMetadata, ItemSampler
from import TabularData

[docs]@tf.keras.utils.register_keras_serializable(package="merlin.models") class InBatchSampler(ItemSampler): """Provides in-batch sampling [1]_ for two-tower item retrieval models. The implementation is very simple, as it just returns the current item embeddings and metadata, but it is necessary to have `InBatchSampler` under the same interface of other more advanced samplers (e.g. `CachedCrossBatchSampler`). In a nutshell, for a given (user,item) embeddings pair, the other in-batch item embeddings are used as negative items, rather than computing different embeddings exclusively for negative items. This is a popularity-biased sampling as popular items are observed more often in training batches. P.s. Ignoring the false negatives (negative items equal to the positive ones) is managed by `ItemRetrievalScorer(..., sampling_downscore_false_negatives=True)` References ---------- .. [1] Yi, Xinyang, et al. "Sampling-bias-corrected neural modeling for large corpus item recommendations." Proceedings of the 13th ACM Conference on Recommender Systems. 2019. Parameters ---------- batch_size : int, optional The batch size. If not set it is inferred when the layer is built (first call()) """
[docs] def __init__(self, batch_size: Optional[int] = None, **kwargs): super().__init__(max_num_samples=batch_size, **kwargs) self._last_batch_items_embeddings: tf.Tensor = None # type: ignore self._last_batch_items_metadata: TabularData = {} self.set_batch_size(batch_size)
@property def batch_size(self) -> int: return self._batch_size
[docs] def set_batch_size(self, value): self._batch_size = value if value is not None: self.set_max_num_samples(value)
[docs] def build(self, input_shapes: TabularData) -> None: if self._batch_size is None: self.set_batch_size(input_shapes["embeddings"][0])
[docs] def call(self, inputs: TabularData, training=True) -> EmbeddingWithMetadata: """Returns the item embeddings and item metadata from the current batch. The implementation is very simple, as it just returns the current item embeddings and metadata, but it is necessary to have `InBatchSampler` under the same interface of other more advanced samplers (e.g. `CachedCrossBatchSampler`). Parameters ---------- inputs : TabularData Dict with two keys: "items_embeddings": Items embeddings tensor "items_metadata": Dict like `{"<feature name>": "<feature tensor>"}` which contains features that might be relevant for the sampler. The `InBatchSampler` does not use metadata features specifically, but "item_id" is required when using in combination with `ItemRetrievalScorer(..., sampling_downscore_false_negatives=True)`, so that false negatives are identified and downscored. training : bool, optional Flag indicating if on training mode, by default True Returns ------- EmbeddingWithMetadata Value object with the sampled item embeddings and item metadata """ self.add(inputs, training) items_embeddings = self.sample() return items_embeddings
[docs] def add(self, inputs: TabularData, training=True) -> None: # type: ignore self._check_inputs_batch_sizes(inputs) self._last_batch_items_embeddings = inputs["embeddings"] self._last_batch_items_metadata = inputs["metadata"]
[docs] def sample(self) -> EmbeddingWithMetadata: return EmbeddingWithMetadata( self._last_batch_items_embeddings, self._last_batch_items_metadata )
[docs] def get_config(self): config = super(InBatchSampler, self).get_config() config["batch_size"] = self.batch_size return config