keras-rs-nightly 0.0.1.dev2025041403__py3-none-any.whl → 0.0.1.dev2025041603__py3-none-any.whl

keras_rs/api/layers/__init__.py

@@ -17,6 +17,7 @@ from keras_rs.src.layers.retrieval.hard_negative_mining import (
 from keras_rs.src.layers.retrieval.remove_accidental_hits import (
     RemoveAccidentalHits,
 )
+from keras_rs.src.layers.retrieval.retrieval import Retrieval
 from keras_rs.src.layers.retrieval.sampling_probability_correction import (
     SamplingProbabilityCorrection,
 )

keras_rs/src/layers/retrieval/brute_force_retrieval.py

@@ -4,10 +4,11 @@ import keras
 
 from keras_rs.src import types
 from keras_rs.src.api_export import keras_rs_export
+from keras_rs.src.layers.retrieval.retrieval import Retrieval
 
 
 @keras_rs_export("keras_rs.layers.BruteForceRetrieval")
-class BruteForceRetrieval(keras.layers.Layer):
+class BruteForceRetrieval(Retrieval):
     """Brute force top-k retrieval.
 
     This layer maintains a set of candidates and is able to exactly retrieve the
@@ -60,11 +61,13 @@ class BruteForceRetrieval(keras.layers.Layer):
         return_scores: bool = True,
         **kwargs: Any,
     ) -> None:
-        super().__init__(**kwargs)
+        # Keep `k`, `return_scores` as separately passed args instead of keeping
+        # them in `kwargs`. This is to ensure the user does not have to hop
+        # to the base class to check which other args can be passed.
+        super().__init__(k=k, return_scores=return_scores, **kwargs)
+
         self.candidate_embeddings = None
         self.candidate_ids = None
-        self.k = k
-        self.return_scores = return_scores
 
         if candidate_embeddings is None:
             if candidate_ids is not None:
@@ -84,36 +87,12 @@ class BruteForceRetrieval(keras.layers.Layer):
 
         Args:
             candidate_embeddings: The candidate embeddings.
-            candidate_ids: The identifiers for the candidates. If `None` the
+            candidate_ids: The identifiers for the candidates. If `None`, the
                 indices of the candidates are returned instead.
         """
-        if candidate_embeddings is None:
-            raise ValueError("`candidate_embeddings` is required")
-
-        if len(candidate_embeddings.shape) != 2:
-            raise ValueError(
-                "`candidate_embeddings` must be a tensor of rank 2 "
-                "(num_candidates, embedding_size), received "
-                "`candidate_embeddings` with shape "
-                f"{candidate_embeddings.shape}"
-            )
-
-        if candidate_embeddings.shape[0] < self.k:
-            raise ValueError(
-                "The number of candidates provided "
-                f"({candidate_embeddings.shape[0]}) is less than the number of "
-                f"candidates to retrieve (k={self.k})."
-            )
-
-        if (
-            candidate_ids is not None
-            and candidate_ids.shape[0] != candidate_embeddings.shape[0]
-        ):
-            raise ValueError(
-                "The `candidate_embeddings` and `candidate_is` tensors must "
-                "have the same number of rows, got tensors of shape "
-                f"{candidate_embeddings.shape} and {candidate_ids.shape}."
-            )
+        self._validate_candidate_embeddings_and_ids(
+            candidate_embeddings, candidate_ids
+        )
 
         if self.candidate_embeddings is not None:
             # Update of existing variables.
@@ -167,31 +146,3 @@ class BruteForceRetrieval(keras.layers.Layer):
             return top_scores, top_ids
         else:
             return top_ids
-
-    def compute_score(
-        self, query_embedding: types.Tensor, candidate_embedding: types.Tensor
-    ) -> types.Tensor:
-        """Computes the standard dot product score from queries and candidates.
-
-        Args:
-            query_embedding: Tensor of query embedding corresponding to the
-                queries for which to retrieve top candidates.
-            candidate_embedding: Tensor of candidate embeddings.
-
-        Returns:
-            The dot product of queries and candidates.
-        """
-
-        return keras.ops.matmul(
-            query_embedding, keras.ops.transpose(candidate_embedding)
-        )
-
-    def get_config(self) -> dict[str, Any]:
-        config: dict[str, Any] = super().get_config()
-        config.update(
-            {
-                "k": self.k,
-                "return_scores": self.compute_score,
-            }
-        )
-        return config

keras_rs/src/layers/retrieval/retrieval.py

@@ -0,0 +1,127 @@
+import abc
+from typing import Any, Optional, Union
+
+import keras
+
+from keras_rs.src import types
+from keras_rs.src.api_export import keras_rs_export
+
+
+@keras_rs_export("keras_rs.layers.Retrieval")
+class Retrieval(keras.layers.Layer, abc.ABC):
+    """Retrieval base abstract class.
+
+    This layer provides a common interface for all retrieval layers. In order
+    to implement a custom retrieval layer, this abstract class should be
+    subclassed.
+
+    Args:
+        k: int. Number of candidates to retrieve.
+        return_scores: bool. When `True`, this layer returns a tuple with the
+            top scores and the top identifiers. When `False`, this layer returns
+            a single tensor with the top identifiers.
+    """
+
+    def __init__(
+        self,
+        k: int = 10,
+        return_scores: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.k = k
+        self.return_scores = return_scores
+
+    def _validate_candidate_embeddings_and_ids(
+        self,
+        candidate_embeddings: types.Tensor,
+        candidate_ids: Optional[types.Tensor] = None,
+    ) -> None:
+        """Validates inputs to `update_candidates()`."""
+
+        if candidate_embeddings is None:
+            raise ValueError("`candidate_embeddings` is required.")
+
+        if len(candidate_embeddings.shape) != 2:
+            raise ValueError(
+                "`candidate_embeddings` must be a tensor of rank 2 "
+                "(num_candidates, embedding_size), received "
+                "`candidate_embeddings` with shape "
+                f"{candidate_embeddings.shape}"
+            )
+
+        if candidate_embeddings.shape[0] < self.k:
+            raise ValueError(
+                "The number of candidates provided "
+                f"({candidate_embeddings.shape[0]}) is less than the number of "
+                f"candidates to retrieve (k={self.k})."
+            )
+
+        if (
+            candidate_ids is not None
+            and candidate_ids.shape[0] != candidate_embeddings.shape[0]
+        ):
+            raise ValueError(
+                "The `candidate_embeddings` and `candidate_is` tensors must "
+                "have the same number of rows, got tensors of shape "
+                f"{candidate_embeddings.shape} and {candidate_ids.shape}."
+            )
+
+    @abc.abstractmethod
+    def update_candidates(
+        self,
+        candidate_embeddings: types.Tensor,
+        candidate_ids: Optional[types.Tensor] = None,
+    ) -> None:
+        """Update the set of candidates and optionally their candidate IDs.
+
+        Args:
+            candidate_embeddings: The candidate embeddings.
+            candidate_ids: The identifiers for the candidates. If `None`, the
+                indices of the candidates are returned instead.
+        """
+        pass
+
+    @abc.abstractmethod
+    def call(
+        self, inputs: types.Tensor
+    ) -> Union[types.Tensor, tuple[types.Tensor, types.Tensor]]:
+        """Returns the top candidates for the query passed as input.
+
+        Args:
+            inputs: the query for which to return top candidates.
+
+        Returns:
+            A tuple with the top scores and the top identifiers if
+            `returns_scores` is True, otherwise a tensor with the top
+            identifiers.
+        """
+        pass
+
+    def compute_score(
+        self, query_embedding: types.Tensor, candidate_embedding: types.Tensor
+    ) -> types.Tensor:
+        """Computes the standard dot product score from queries and candidates.
+
+        Args:
+            query_embedding: Tensor of query embedding corresponding to the
+                queries for which to retrieve top candidates.
+            candidate_embedding: Tensor of candidate embeddings.
+
+        Returns:
+            The dot product of queries and candidates.
+        """
+
+        return keras.ops.matmul(
+            query_embedding, keras.ops.transpose(candidate_embedding)
+        )
+
+    def get_config(self) -> dict[str, Any]:
+        config: dict[str, Any] = super().get_config()
+        config.update(
+            {
+                "k": self.k,
+                "return_scores": self.compute_score,
+            }
+        )
+        return config

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.dev2025041403"
+__version__ = "0.0.1.dev2025041603"
 
 
 @keras_rs_export("keras_rs.version")

{keras_rs_nightly-0.0.1.dev2025041403.dist-info → keras_rs_nightly-0.0.1.dev2025041603.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-rs-nightly
-Version: 0.0.1.dev2025041403
+Version: 0.0.1.dev2025041603
 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.dev2025041403.dist-info → keras_rs_nightly-0.0.1.dev2025041603.dist-info}/RECORD

@@ -1,19 +1,20 @@
 keras_rs/__init__.py,sha256=X3VNKb_6VDEs5GqcbEc_l8mAsefWb5UgSu8krnQdFcM,794
 keras_rs/api/__init__.py,sha256=9Xf-uH9j_SBaTc5RU0pkxrOEgHWPwSKjf4_maySH_nU,272
-keras_rs/api/layers/__init__.py,sha256=VvsmrU1T198D8svbjAk7bysRRfDeT71A-9wv7NyPC28,685
+keras_rs/api/layers/__init__.py,sha256=SB7_QOBPizvbbyQAMb8mPl7vAx0gCxJBPm6V7H67SgU,747
 keras_rs/api/losses/__init__.py,sha256=LGW7eHQh8FbQXdMV1s9zJpbloVlz_Zlo51sorWAvFwE,455
 keras_rs/src/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_rs/src/api_export.py,sha256=RsmG-DvO-cdFeAF9W6LRzms0kvtm-Yp9BAA_d-952zI,510
 keras_rs/src/types.py,sha256=UyOdgjqrqg_b58opnY8n6gTiDHKVR8z_bmEruehERBk,514
-keras_rs/src/version.py,sha256=U6vFDmQdY-2KYLaBQWSFZR7zwHQxNymE9MR_swK3DUU,222
+keras_rs/src/version.py,sha256=wgkFHQtzZaQah52nHJja4pYng_h75BXqZEskD1h29LI,222
 keras_rs/src/layers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_rs/src/layers/feature_interaction/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_rs/src/layers/feature_interaction/dot_interaction.py,sha256=jGHcg0EiWxth6LTxG2yWgHcyx_GXrxvA61uQqpPfnDQ,6900
 keras_rs/src/layers/feature_interaction/feature_cross.py,sha256=5OCSI0vFYzJNmgkKcuHIbVv8U2q3UvS80-qZjPimDjM,8155
 keras_rs/src/layers/retrieval/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-keras_rs/src/layers/retrieval/brute_force_retrieval.py,sha256=mohILOt6PC6jHBztaowDbj3QBnSetuvkq55FmE39PlY,7321
+keras_rs/src/layers/retrieval/brute_force_retrieval.py,sha256=izdppBXxJH0KqYEg7Zsr-SL-SHgAmnFopXMPalEO3uw,5676
 keras_rs/src/layers/retrieval/hard_negative_mining.py,sha256=CY8-3W52ZBIFcEfvjXJxbFltD6ulXl4-sZCRF6stIEc,4119
 keras_rs/src/layers/retrieval/remove_accidental_hits.py,sha256=Z84z2YgKspKeNdc5id8lf9TAyFsbCCz3acJxiKXYipc,3324
+keras_rs/src/layers/retrieval/retrieval.py,sha256=hVOBF10SF2q_TgJdVUqztbnw5qQF-cxVRGdJbOKoL9M,4191
 keras_rs/src/layers/retrieval/sampling_probability_correction.py,sha256=80vgOPfBiF-PC0dSyqS57IcIxOxi_Q_R7eSXHn1G0yI,1437
 keras_rs/src/losses/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_rs/src/losses/pairwise_hinge_loss.py,sha256=vqDGd-OnZxiqdeE6vuabE8BKDfill3D2GM0lW5JUmsg,922
@@ -24,7 +25,7 @@ keras_rs/src/losses/pairwise_soft_zero_one_loss.py,sha256=XBej5nybFXEQ-Vp6GLvNmq
 keras_rs/src/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_rs/src/utils/keras_utils.py,sha256=IjWSRieBkv7UX12qgUoI1tcOeISstCLRSTqSHpT06yE,1275
 keras_rs/src/utils/pairwise_loss_utils.py,sha256=6eF4CTJubCySO8M5nd3_gdTlJsta_YMnwDCcdqWYGHA,3435
-keras_rs_nightly-0.0.1.dev2025041403.dist-info/METADATA,sha256=6Wr1cV03cJFymTfdf6IVSDEPeraj65YKLqTo2QlvX-s,3547
-keras_rs_nightly-0.0.1.dev2025041403.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-keras_rs_nightly-0.0.1.dev2025041403.dist-info/top_level.txt,sha256=pWs8X78Z0cn6lfcIb9VYOW5UeJ-TpoaO9dByzo7_FFo,9
-keras_rs_nightly-0.0.1.dev2025041403.dist-info/RECORD,,
+keras_rs_nightly-0.0.1.dev2025041603.dist-info/METADATA,sha256=9EV3mNVpTEuZyIu5Ihha5KhjRwRJnstN4vm1iXMVvQA,3547
+keras_rs_nightly-0.0.1.dev2025041603.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+keras_rs_nightly-0.0.1.dev2025041603.dist-info/top_level.txt,sha256=pWs8X78Z0cn6lfcIb9VYOW5UeJ-TpoaO9dByzo7_FFo,9
+keras_rs_nightly-0.0.1.dev2025041603.dist-info/RECORD,,