keras-rs-nightly 0.0.1.dev2025021903__py3-none-any.whl → 0.3.1.dev202512130338__py3-none-any.whl

keras_rs/src/metrics/ndcg.py

@@ -0,0 +1,197 @@
+from typing import Any, Callable
+
+from keras import ops
+from keras.saving import deserialize_keras_object
+from keras.saving import serialize_keras_object
+
+from keras_rs.src import types
+from keras_rs.src.api_export import keras_rs_export
+from keras_rs.src.metrics.ranking_metric import RankingMetric
+from keras_rs.src.metrics.ranking_metric import (
+    ranking_metric_subclass_doc_string,
+)
+from keras_rs.src.metrics.ranking_metric import (
+    ranking_metric_subclass_doc_string_post_desc,
+)
+from keras_rs.src.metrics.ranking_metrics_utils import compute_dcg
+from keras_rs.src.metrics.ranking_metrics_utils import default_gain_fn
+from keras_rs.src.metrics.ranking_metrics_utils import default_rank_discount_fn
+from keras_rs.src.metrics.ranking_metrics_utils import get_list_weights
+from keras_rs.src.metrics.ranking_metrics_utils import sort_by_scores
+from keras_rs.src.utils.doc_string_utils import format_docstring
+
+
+@keras_rs_export("keras_rs.metrics.NDCG")
+class NDCG(RankingMetric):
+    def __init__(
+        self,
+        k: int | None = None,
+        gain_fn: Callable[[types.Tensor], types.Tensor] = default_gain_fn,
+        rank_discount_fn: Callable[
+            [types.Tensor], types.Tensor
+        ] = default_rank_discount_fn,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(k=k, **kwargs)
+
+        self.gain_fn = gain_fn
+        self.rank_discount_fn = rank_discount_fn
+
+    def compute_metric(
+        self,
+        y_true: types.Tensor,
+        y_pred: types.Tensor,
+        mask: types.Tensor,
+        sample_weight: types.Tensor,
+    ) -> types.Tensor:
+        sorted_y_true, sorted_weights = sort_by_scores(
+            tensors_to_sort=[y_true, sample_weight],
+            scores=y_pred,
+            k=self.k,
+            mask=mask,
+            shuffle_ties=self.shuffle_ties,
+            seed=self.seed_generator,
+        )
+
+        dcg = compute_dcg(
+            y_true=sorted_y_true,
+            sample_weight=sorted_weights,
+            gain_fn=self.gain_fn,
+            rank_discount_fn=self.rank_discount_fn,
+        )
+
+        weighted_gains = ops.multiply(
+            sample_weight,
+            self.gain_fn(y_true),
+        )
+        ideal_sorted_y_true, ideal_sorted_weights = sort_by_scores(
+            tensors_to_sort=[y_true, sample_weight],
+            scores=weighted_gains,
+            k=self.k,
+            mask=mask,
+            shuffle_ties=self.shuffle_ties,
+            seed=self.seed_generator,
+        )
+        ideal_dcg = compute_dcg(
+            y_true=ideal_sorted_y_true,
+            sample_weight=ideal_sorted_weights,
+            gain_fn=self.gain_fn,
+            rank_discount_fn=self.rank_discount_fn,
+        )
+        per_list_ndcg = ops.divide_no_nan(dcg, ideal_dcg)
+
+        per_list_weights = get_list_weights(
+            weights=sample_weight, relevance=self.gain_fn(y_true)
+        )
+
+        return per_list_ndcg, per_list_weights
+
+    def get_config(self) -> dict[str, Any]:
+        config: dict[str, Any] = super().get_config()
+        config.update(
+            {
+                "gain_fn": serialize_keras_object(self.gain_fn),
+                "rank_discount_fn": serialize_keras_object(
+                    self.rank_discount_fn
+                ),
+            }
+        )
+        return config
+
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> "NDCG":
+        config["gain_fn"] = deserialize_keras_object(config["gain_fn"])
+        config["rank_discount_fn"] = deserialize_keras_object(
+            config["rank_discount_fn"]
+        )
+        return cls(**config)
+
+
+concept_sentence = (
+    "It normalizes the Discounted Cumulative Gain (DCG) with the Ideal "
+    "Discounted Cumulative Gain (IDCG) for each list"
+)
+relevance_type = (
+    "graded relevance scores (non-negative numbers where higher values "
+    "indicate greater relevance)"
+)
+score_range_interpretation = (
+    "A normalized score (between 0 and 1) is returned. A score of 1 "
+    "represents the perfect ranking according to true relevance (within the "
+    "top-k), while 0 typically represents a ranking with no relevant items. "
+    "Higher scores indicate better ranking quality relative to the best "
+    "possible ranking"
+)
+
+formula = """
+```
+nDCG@k = DCG@k / IDCG@k
+```
+
+where DCG@k is calculated based on the predicted ranking (`y_pred`):
+
+```
+DCG@k(y') = sum_{i=1}^{k} (gain_fn(y'_i) / rank_discount_fn(i))
+```
+
+And IDCG@k is the Ideal DCG, calculated using the same formula but on items
+sorted perfectly by their *true relevance* (`y_true`):
+
+```
+IDCG@k(y'') = sum_{i=1}^{k} (gain_fn(y''_i) / rank_discount_fn(i))
+```
+
+where:
+- `y'_i`: True relevance of the item at rank `i` in the ranking induced by
+  `y_pred`.
+- `y''_i` True relevance of the item at rank `i` in the *ideal* ranking (sorted
+  by `y_true` descending).
+- `gain_fn` is the user-provided function mapping relevance to gain. The default
+  function (`default_gain_fn`) is typically equivalent to `lambda y: 2**y - 1`.
+- `rank_discount_fn` is the user-provided function mapping rank `i` (1-based) to
+  a discount value. The default function (`default_rank_discount_fn`) is
+  typically equivalent to `lambda rank: 1 / log2(rank + 1)`.
+- If IDCG@k is 0 (e.g., no relevant items), nDCG@k is defined as 0.
+- The final result often aggregates these per-list nDCG scores, potentially
+  involving normalization by list-specific weights, to produce a weighted
+  average."""
+extra_args = """
+        gain_fn: callable. Maps relevance scores (`y_true`) to gain values. The
+            default implements `2**y - 1`.
+        rank_discount_fn: function. Maps rank positions to discount
+            values. The default (`default_rank_discount_fn`) implements
+            `1 / log2(rank + 1)`."""
+example = """
+    >>> batch_size = 2
+    >>> list_size = 5
+    >>> labels = np.random.randint(0, 3, size=(batch_size, list_size))
+    >>> scores = np.random.random(size=(batch_size, list_size))
+    >>> metric = keras_rs.metrics.NDCG()(
+    ...     y_true=labels, y_pred=scores
+    ... )
+
+    Mask certain elements (can be used for uneven inputs):
+
+    >>> batch_size = 2
+    >>> list_size = 5
+    >>> labels = np.random.randint(0, 3, size=(batch_size, list_size))
+    >>> scores = np.random.random(size=(batch_size, list_size))
+    >>> mask = np.random.randint(0, 2, size=(batch_size, list_size), dtype=bool)
+    >>> metric = keras_rs.metrics.NDCG()(
+    ...     y_true={"labels": labels, "mask": mask}, y_pred=scores
+    ... )
+"""
+
+NDCG.__doc__ = format_docstring(
+    ranking_metric_subclass_doc_string,
+    width=80,
+    metric_name="Normalised Discounted Cumulative Gain",
+    metric_abbreviation="nDCG",
+    concept_sentence=concept_sentence,
+    relevance_type=relevance_type,
+    score_range_interpretation=score_range_interpretation,
+    formula=formula,
+    extra_args=extra_args,
+) + ranking_metric_subclass_doc_string_post_desc.format(
+    extra_args=extra_args, example=example
+)

keras_rs/src/metrics/precision_at_k.py

@@ -0,0 +1,117 @@
+from keras import ops
+
+from keras_rs.src import types
+from keras_rs.src.api_export import keras_rs_export
+from keras_rs.src.metrics.ranking_metric import RankingMetric
+from keras_rs.src.metrics.ranking_metric import (
+    ranking_metric_subclass_doc_string,
+)
+from keras_rs.src.metrics.ranking_metric import (
+    ranking_metric_subclass_doc_string_post_desc,
+)
+from keras_rs.src.metrics.ranking_metrics_utils import get_list_weights
+from keras_rs.src.metrics.ranking_metrics_utils import sort_by_scores
+from keras_rs.src.utils.doc_string_utils import format_docstring
+
+
+@keras_rs_export("keras_rs.metrics.PrecisionAtK")
+class PrecisionAtK(RankingMetric):
+    def compute_metric(
+        self,
+        y_true: types.Tensor,
+        y_pred: types.Tensor,
+        mask: types.Tensor,
+        sample_weight: types.Tensor,
+    ) -> types.Tensor:
+        # Assume: `y_true = [0, 0, 1]`, `y_pred = [0.1, 0.9, 0.2]`.
+        # `sorted_y_true = [0, 1, 0]` (sorted in descending order).
+        (sorted_y_true,) = sort_by_scores(
+            tensors_to_sort=[y_true],
+            scores=y_pred,
+            mask=mask,
+            k=self.k,
+            shuffle_ties=self.shuffle_ties,
+            seed=self.seed_generator,
+        )
+
+        # We consider only binary relevance here, anything above 1 is treated
+        # as 1. `relevance = [0., 1., 0.]`.
+        relevance = ops.cast(
+            ops.greater_equal(
+                sorted_y_true, ops.cast(1, dtype=sorted_y_true.dtype)
+            ),
+            dtype=y_pred.dtype,
+        )
+        list_length = ops.shape(sorted_y_true)[1]
+        # TODO: We do not do this for MRR, and the other metrics. Do we need to
+        # do this there too?
+        valid_list_length = ops.minimum(
+            list_length,
+            ops.sum(ops.cast(mask, dtype="int32"), axis=1, keepdims=True),
+        )
+
+        per_list_precision = ops.divide_no_nan(
+            ops.sum(relevance, axis=1, keepdims=True),
+            ops.cast(valid_list_length, dtype=y_pred.dtype),
+        )
+
+        # Get weights.
+        overall_relevance = ops.cast(
+            ops.greater_equal(y_true, ops.cast(1, dtype=y_true.dtype)),
+            dtype=y_pred.dtype,
+        )
+        per_list_weights = get_list_weights(
+            weights=sample_weight, relevance=overall_relevance
+        )
+
+        return per_list_precision, per_list_weights
+
+
+concept_sentence = (
+    "It measures the proportion of relevant items among the top-k "
+    "recommendations"
+)
+relevance_type = "binary indicators (0 or 1) of relevance"
+score_range_interpretation = (
+    "Scores range from 0 to 1, with 1 indicating all top-k items were relevant"
+)
+formula = """```
+P@k(y, s) = 1/k sum_i I[rank(s_i) < k] y_i
+```
+
+where `y_i` is the relevance label (0/1) of the item ranked at position
+`i`, and `I[condition]` is 1 if the condition is met, otherwise 0."""
+extra_args = ""
+example = """
+    >>> batch_size = 2
+    >>> list_size = 5
+    >>> labels = np.random.randint(0, 2, size=(batch_size, list_size))
+    >>> scores = np.random.random(size=(batch_size, list_size))
+    >>> metric = keras_rs.metrics.PrecisionAtK()(
+    ...     y_true=labels, y_pred=scores
+    ... )
+
+    Mask certain elements (can be used for uneven inputs):
+
+    >>> batch_size = 2
+    >>> list_size = 5
+    >>> labels = np.random.randint(0, 2, size=(batch_size, list_size))
+    >>> scores = np.random.random(size=(batch_size, list_size))
+    >>> mask = np.random.randint(0, 2, size=(batch_size, list_size), dtype=bool)
+    >>> metric = keras_rs.metrics.PrecisionAtK()(
+    ...     y_true={"labels": labels, "mask": mask}, y_pred=scores
+    ... )
+"""
+
+PrecisionAtK.__doc__ = format_docstring(
+    ranking_metric_subclass_doc_string,
+    width=80,
+    metric_name="Precision@k",
+    metric_abbreviation="P@k",
+    concept_sentence=concept_sentence,
+    relevance_type=relevance_type,
+    score_range_interpretation=score_range_interpretation,
+    formula=formula,
+) + ranking_metric_subclass_doc_string_post_desc.format(
+    extra_args=extra_args, example=example
+)

keras_rs/src/metrics/ranking_metric.py

@@ -0,0 +1,260 @@
+import abc
+from typing import Any
+
+import keras
+from keras import ops
+
+from keras_rs.src import types
+from keras_rs.src.metrics.utils import standardize_call_inputs_ranks
+from keras_rs.src.utils.keras_utils import check_rank
+from keras_rs.src.utils.keras_utils import check_shapes_compatible
+
+
+class RankingMetric(keras.metrics.Mean, abc.ABC):
+    """Base class for ranking evaluation metrics (e.g., MAP, MRR, DCG, nDCG).
+
+    Ranking metrics are used to evaluate the quality of ranked lists produced
+    by a ranking model. The primary goal in ranking tasks is to order items
+    according to their relevance or utility for a given query or context.
+    These metrics provide quantitative measures of how well a model achieves
+    this goal, typically by comparing the predicted order of items against the
+    ground truth relevance judgments for each list.
+
+    To define your own ranking metric, subclass this class and implement the
+    `compute_metric` method.
+
+    Args:
+        k: int. The number of top-ranked items to consider (the 'k' in 'top-k').
+            Must be a positive integer.
+        shuffle_ties: bool. Whether to randomly shuffle scores before sorting.
+            This is done to break ties. Defaults to `True`.
+        seed: int. Random seed used for shuffling.
+        name: Optional name for the loss instance.
+        dtype: The dtype of the metric's computations. Defaults to `None`, which
+            means using `keras.backend.floatx()`. `keras.backend.floatx()` is a
+            `"float32"` unless set to different value
+            (via `keras.backend.set_floatx()`). If a `keras.DTypePolicy` is
+            provided, then the `compute_dtype` will be utilized.
+    """
+
+    def __init__(
+        self,
+        k: int | None = None,
+        shuffle_ties: bool = True,
+        seed: int | keras.random.SeedGenerator | None = None,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+
+        if k is not None and (not isinstance(k, int) or k < 1):
+            raise ValueError(
+                f"`k` should be a positive integer. Received: `k` = {k}."
+            )
+
+        self.k = k
+        self.shuffle_ties = shuffle_ties
+        self.seed = seed
+
+        # Define `SeedGenerator`. JAX doesn't work, otherwise.
+        self.seed_generator = keras.random.SeedGenerator(seed)
+
+    @abc.abstractmethod
+    def compute_metric(
+        self,
+        y_true: types.Tensor,
+        y_pred: types.Tensor,
+        mask: types.Tensor,
+        sample_weight: types.Tensor,
+    ) -> types.Tensor:
+        """Abstract method, should be implemented by subclasses."""
+        pass
+
+    def update_state(
+        self,
+        y_true: types.Tensor,
+        y_pred: types.Tensor,
+        sample_weight: types.Tensor | None = None,
+    ) -> None:
+        """
+        Accumulates statistics for the ranking metric.
+
+        Args:
+            y_true: tensor or dict. Ground truth values. If tensor, of shape
+                `(list_size)` for unbatched inputs or `(batch_size, list_size)`
+                for batched inputs. If an item has a label of -1, it is ignored
+                in loss computation. If it is a dictionary, it should have two
+                keys: `"labels"` and `"mask"`. `"mask"` can be used to ignore
+                elements in metric computation, i.e., pairs will not be formed
+                with those items. Note that the final mask is an `and` of the
+                passed mask, `labels >= 0`, and `sample_weight > 0`.
+            y_pred: tensor. The predicted values, of shape `(list_size)` for
+                unbatched inputs or `(batch_size, list_size)` for batched
+                inputs. Should be of the same shape as `y_true`.
+            sample_weight: float/tensor. Can be float value, or tensor of
+                shape `(list_size)` or `(batch_size, list_size)`. Defaults to
+                `None`.
+        """
+        # === Process `y_true`, if dict ===
+        passed_mask = None
+        if isinstance(y_true, dict):
+            if "labels" not in y_true:
+                raise ValueError(
+                    '`"labels"` should be present in `y_true`. Received: '
+                    f"`y_true` = {y_true}"
+                )
+
+            passed_mask = y_true.get("mask", None)
+            y_true = y_true["labels"]
+
+        # === Convert to tensors, if list ===
+        # TODO (abheesht): Figure out if we need to cast tensors to
+        # `self.dtype`.
+        y_true = ops.convert_to_tensor(y_true)
+        y_pred = ops.convert_to_tensor(y_pred)
+        if sample_weight is not None:
+            sample_weight = ops.convert_to_tensor(sample_weight)
+        if passed_mask is not None:
+            passed_mask = ops.convert_to_tensor(passed_mask)
+
+        # Cast to the correct dtype.
+        y_true = ops.cast(y_true, dtype=self.dtype)
+        y_pred = ops.cast(y_pred, dtype=self.dtype)
+        if sample_weight is not None:
+            sample_weight = ops.cast(sample_weight, dtype=self.dtype)
+
+        # === Process `sample_weight` ===
+        if sample_weight is None:
+            sample_weight = ops.cast(1, dtype=y_pred.dtype)
+
+        y_true_shape = ops.shape(y_true)
+        y_true_rank = len(y_true_shape)
+        sample_weight_shape = ops.shape(sample_weight)
+        sample_weight_rank = len(sample_weight_shape)
+
+        # Check `y_true_rank` first. Can be 1 for unbatched inputs, 2 for
+        # batched.
+        check_rank(y_true_rank, allowed_ranks=(1, 2), tensor_name="y_true")
+
+        # Check `sample_weight` rank. Should be between 0 and `y_true_rank`.
+        check_rank(
+            sample_weight_rank,
+            allowed_ranks=tuple(range(y_true_rank + 1)),
+            tensor_name="sample_weight",
+        )
+
+        if y_true_rank == 2:
+            # If `sample_weight` rank is 1, it should be of shape
+            # `(batch_size,)`. Otherwise, it should be of shape
+            # `(batch_size, list_size)`.
+            if sample_weight_rank == 1:
+                check_shapes_compatible(sample_weight_shape, (y_true_shape[0],))
+                # Uprank this, so that we get per-list weights here.
+                sample_weight = ops.expand_dims(sample_weight, axis=1)
+            elif sample_weight_rank == 2:
+                check_shapes_compatible(sample_weight_shape, y_true_shape)
+
+        # Reshape `sample_weight` to the shape of `y_true`.
+        sample_weight = ops.multiply(ops.ones_like(y_true), sample_weight)
+
+        # Mask all values less than 0 (since less than 0 implies invalid
+        # labels).
+        valid_mask = ops.greater_equal(y_true, ops.cast(0, y_true.dtype))
+        if passed_mask is not None:
+            valid_mask = ops.logical_and(valid_mask, passed_mask)
+
+        # === Process inputs - shape checking, upranking, etc. ===
+        y_true, y_pred, valid_mask, batched = standardize_call_inputs_ranks(
+            y_true=y_true,
+            y_pred=y_pred,
+            mask=valid_mask,
+            check_y_true_rank=False,
+        )
+
+        # Uprank sample_weight if unbatched.
+        if not batched:
+            sample_weight = ops.expand_dims(sample_weight, axis=0)
+
+        # Get `mask` from `sample_weight`.
+        sample_weight_mask = ops.greater(
+            sample_weight, ops.cast(0, dtype=sample_weight.dtype)
+        )
+        mask = ops.logical_and(valid_mask, sample_weight_mask)
+
+        # === Update "invalid" `y_true`, `y_pred` entries based on mask ===
+
+        # `y_true`: assign 0 for invalid entries
+        y_true = ops.where(mask, y_true, ops.zeros_like(y_true))
+        # `y_pred`: assign a value slightly smaller than the smallest value
+        # so that it gets sorted last.
+        y_pred = ops.where(
+            mask,
+            y_pred,
+            -keras.config.epsilon() * ops.ones_like(y_pred)
+            + ops.amin(y_pred, axis=1, keepdims=True),
+        )
+        sample_weight = ops.where(
+            mask, sample_weight, ops.cast(0, sample_weight.dtype)
+        )
+
+        # === Actual computation ===
+        per_list_metric_values, per_list_metric_weights = self.compute_metric(
+            y_true=y_true, y_pred=y_pred, mask=mask, sample_weight=sample_weight
+        )
+
+        # Chain to `super()` to get mean metric.
+        # TODO(abheesht): Figure out if we want to return unaggregated metric
+        # values too of shape `(batch_size,)` from `result()`.
+        super().update_state(
+            per_list_metric_values, sample_weight=per_list_metric_weights
+        )
+
+    def get_config(self) -> dict[str, Any]:
+        config: dict[str, Any] = super().get_config()
+        config.update(
+            {"k": self.k, "shuffle_ties": self.shuffle_ties, "seed": self.seed}
+        )
+        return config
+
+
+ranking_metric_subclass_doc_string = """
+Computes {metric_name} ({metric_abbreviation}).
+
+This metric evaluates ranking quality. {concept_sentence}. The metric processes
+true relevance labels in `y_true` ({relevance_type}) against predicted scores in
+`y_pred`. The scores in `y_pred` are used to determine the rank order of items,
+by sorting in descending order. {score_range_interpretation}.
+
+For each list of predicted scores `s` in `y_pred` and the corresponding list
+of true labels `y` in `y_true`, the per-query {metric_abbreviation} score is
+calculated as follows:
+{formula}
+
+The final {metric_abbreviation} score reported is typically the weighted
+average of these per-query scores across all queries/lists in the dataset.
+
+Note: `sample_weight` is handled differently for ranking metrics. For
+batched inputs, `sample_weight` can be scalar, 1D, 2D. The scalar case and
+1D case (list-wise weights) are straightforward. The 2D case (item-wise
+weights) is different, in the sense that the sample weights are aggregated
+to get 1D weights. For more details, refer to
+`keras_rs.src.metrics.ranking_metrics_utils.get_list_weights`.
+"""
+
+ranking_metric_subclass_doc_string_post_desc = """
+
+    Args:{extra_args}
+        k: int. The number of top-ranked items to consider (the 'k' in 'top-k').
+            Must be a positive integer.
+        shuffle_ties: bool. Whether to randomly shuffle scores before sorting.
+            This is done to break ties. Defaults to `True`.
+        seed: int. Random seed used for shuffling.
+        name: Optional name for the loss instance.
+        dtype: The dtype of the metric's computations. Defaults to `None`, which
+            means using `keras.backend.floatx()`. `keras.backend.floatx()` is a
+            `"float32"` unless set to different value
+            (via `keras.backend.set_floatx()`). If a `keras.DTypePolicy` is
+            provided, then the `compute_dtype` will be utilized.
+
+    Example:
+{example}
+"""