keras-rs-nightly 0.0.1.dev2025022603__tar.gz → 0.0.1.dev2025022803__tar.gz

{keras_rs_nightly-0.0.1.dev2025022603 → keras_rs_nightly-0.0.1.dev2025022803}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: keras-rs-nightly
-Version: 0.0.1.dev2025022603
+Version: 0.0.1.dev2025022803
 Summary: Multi-backend recommender systems with Keras 3.
 Author-email: Keras RS team <keras-rs@google.com>
 License: Apache License 2.0

{keras_rs_nightly-0.0.1.dev2025022603 → keras_rs_nightly-0.0.1.dev2025022803}/keras_rs/api/layers/__init__.py

@@ -9,3 +9,6 @@ from keras_rs.src.layers.modeling.feature_cross import FeatureCross
 from keras_rs.src.layers.retrieval.brute_force_retrieval import (
     BruteForceRetrieval,
 )
+from keras_rs.src.layers.retrieval.hard_negative_mining import (
+    HardNegativeMining,
+)

{keras_rs_nightly-0.0.1.dev2025022603 → keras_rs_nightly-0.0.1.dev2025022803}/keras_rs/src/layers/modeling/dot_interaction.py

@@ -28,6 +28,7 @@ class DotInteraction(keras.layers.Layer):
             entries will be zeros. Otherwise, the output will be only the lower
             triangular part of the interaction matrix. The latter saves space
             but is much slower.
+        **kwargs: Args to pass to the base class.
 
     References:
     - [M. Naumov et al.](https://arxiv.org/abs/1906.00091)

{keras_rs_nightly-0.0.1.dev2025022603 → keras_rs_nightly-0.0.1.dev2025022803}/keras_rs/src/layers/modeling/feature_cross.py

@@ -52,6 +52,7 @@ class FeatureCross(keras.layers.Layer):
             Regularizer to use for the kernel matrix.
         bias_regularizer: string or `keras.regularizer` regularizer.
             Regularizer to use for the bias vector.
+        **kwargs: Args to pass to the base class.
 
     Example:
 

keras_rs_nightly-0.0.1.dev2025022803/keras_rs/src/layers/retrieval/hard_negative_mining.py

@@ -0,0 +1,111 @@
+from typing import Any
+
+import keras
+import numpy as np
+from keras import ops
+
+from keras_rs.src import types
+from keras_rs.src.api_export import keras_rs_export
+
+MAX_FLOAT = np.finfo(np.float32).max / 100.0
+
+
+def _gather_elements_along_row(
+    data: types.Tensor, column_indices: types.Tensor
+) -> types.Tensor:
+    """Gathers elements from a 2D tensor given the column indices of each row.
+
+    First, gets the flat 1D indices to gather from. Then flattens the data to 1D
+    and uses `ops.take()` to generate 1D output and finally reshapes the output
+    back to 2D.
+
+    Args:
+        data: A [N, M] 2D `Tensor`.
+        column_indices: A [N, K] 2D `Tensor` denoting for each row, the K column
+            indices to gather elements from the data `Tensor`.
+
+    Returns:
+        A [N, K] `Tensor` including output elements gathered from data `Tensor`.
+
+    Raises:
+        ValueError: if the first dimensions of data and column_indices don't
+            match.
+    """
+    num_row, num_column, *_ = ops.shape(data)
+    num_gathered = ops.shape(column_indices)[1]
+    row_indices = ops.tile(
+        ops.expand_dims(ops.arange(num_row), -1), [1, num_gathered]
+    )
+    flat_data = ops.reshape(data, [-1])
+    flat_indices = ops.reshape(
+        ops.add(ops.multiply(row_indices, num_column), column_indices), [-1]
+    )
+    return ops.reshape(
+        ops.take(flat_data, flat_indices), [num_row, num_gathered]
+    )
+
+
+@keras_rs_export("keras_rs.layers.HardNegativeMining")
+class HardNegativeMining(keras.layers.Layer):
+    """Transforms logits and labels to return hard negatives.
+
+    Args:
+        num_hard_negatives: How many hard negatives to return.
+        **kwargs: Args to pass to the base class.
+    """
+
+    def __init__(self, num_hard_negatives: int, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self._num_hard_negatives = num_hard_negatives
+        self.built = True
+
+    def call(
+        self, logits: types.Tensor, labels: types.Tensor
+    ) -> tuple[types.Tensor, types.Tensor]:
+        """Filters logits and labels with per-query hard negative mining.
+
+        The result will include logits and labels for `num_hard_negatives`
+        negatives as well as the positive candidate.
+
+        Args:
+            logits: `[batch_size, number_of_candidates]` tensor of logits.
+            labels: `[batch_size, number_of_candidates]` one-hot tensor of
+                labels.
+
+        Returns:
+            tuple containing:
+            - logits: `[batch_size, num_hard_negatives + 1]` tensor of logits.
+            - labels: `[batch_size, num_hard_negatives + 1]` one-hot tensor of
+                  labels.
+        """
+
+        # Number of sampled logits, i.e, the number of hard negatives to be
+        # sampled (k) + number of true logit (1) per query, capped by batch
+        # size.
+        num_logits = ops.shape(logits)[1]
+        if isinstance(num_logits, int):
+            num_sampled = min(self._num_hard_negatives + 1, num_logits)
+        else:
+            num_sampled = ops.minimum(self._num_hard_negatives + 1, num_logits)
+        # To gather indices of top k negative logits per row (query) in logits,
+        # true logits need to be excluded. First replace the true logits
+        # (corresponding to positive labels) with a large score value and then
+        # select the top k + 1 logits from each row so that selected indices
+        # include the indices of true logit + top k negative logits. This
+        # approach is to avoid using inefficient masking when excluding true
+        # logits.
+
+        # For each query, get the indices of the logits which have the highest
+        # k + 1 logit values, including the highest k negative logits and one
+        # true logit.
+        _, col_indices = ops.top_k(
+            ops.add(logits, ops.multiply(labels, MAX_FLOAT)),
+            k=num_sampled,
+            sorted=False,
+        )
+
+        # Gather sampled logits and corresponding labels.
+        logits = _gather_elements_along_row(logits, col_indices)
+        labels = _gather_elements_along_row(labels, col_indices)
+
+        return logits, labels

{keras_rs_nightly-0.0.1.dev2025022603 → keras_rs_nightly-0.0.1.dev2025022803}/keras_rs/src/version.py

@@ -1,7 +1,7 @@
 from keras_rs.src.api_export import keras_rs_export
 
 # Unique source of truth for the version number.
-__version__ = "0.0.1.dev2025022603"
+__version__ = "0.0.1.dev2025022803"
 
 
 @keras_rs_export("keras_rs.version")

{keras_rs_nightly-0.0.1.dev2025022603 → keras_rs_nightly-0.0.1.dev2025022803}/keras_rs_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: keras-rs-nightly
-Version: 0.0.1.dev2025022603
+Version: 0.0.1.dev2025022803
 Summary: Multi-backend recommender systems with Keras 3.
 Author-email: Keras RS team <keras-rs@google.com>
 License: Apache License 2.0

{keras_rs_nightly-0.0.1.dev2025022603 → keras_rs_nightly-0.0.1.dev2025022803}/keras_rs_nightly.egg-info/SOURCES.txt

@@ -13,6 +13,7 @@ keras_rs/src/layers/modeling/dot_interaction.py
 keras_rs/src/layers/modeling/feature_cross.py
 keras_rs/src/layers/retrieval/__init__.py
 keras_rs/src/layers/retrieval/brute_force_retrieval.py
+keras_rs/src/layers/retrieval/hard_negative_mining.py
 keras_rs/src/utils/__init__.py
 keras_rs/src/utils/keras_utils.py
 keras_rs_nightly.egg-info/PKG-INFO