MAINT Parameters validation for metrics.ndcg_score (scikit-learn#25885)

sortofamudkip · web-flow · commit 5a17650a185c · 2023-03-17T13:20:27.000+01:00
diff --git a/sklearn/metrics/_ranking.py b/sklearn/metrics/_ranking.py
@@ -1620,6 +1620,15 @@ def _ndcg_sample_scores(y_true, y_score, k=None, ignore_ties=False):
     return gain
 
 
+@validate_params(
+    {
+        "y_true": ["array-like"],
+        "y_score": ["array-like"],
+        "k": [Interval(Integral, 1, None, closed="left"), None],
+        "sample_weight": ["array-like", None],
+        "ignore_ties": ["boolean"],
+    }
+)
 def ndcg_score(y_true, y_score, *, k=None, sample_weight=None, ignore_ties=False):
     """Compute Normalized Discounted Cumulative Gain.
 
@@ -1633,15 +1642,15 @@ def ndcg_score(y_true, y_score, *, k=None, sample_weight=None, ignore_ties=False
 
     Parameters
     ----------
-    y_true : ndarray of shape (n_samples, n_labels)
+    y_true : array-like of shape (n_samples, n_labels)
         True targets of multilabel classification, or true scores of entities
         to be ranked. Negative values in `y_true` may result in an output
         that is not between 0 and 1.
 
         .. versionchanged:: 1.2
             These negative values are deprecated, and will raise an error in v1.4.
 
-    y_score : ndarray of shape (n_samples, n_labels)
+    y_score : array-like of shape (n_samples, n_labels)
         Target scores, can either be probability estimates, confidence values,
         or non-thresholded measure of decisions (as returned by
         "decision_function" on some classifiers).
@@ -1650,7 +1659,7 @@ def ndcg_score(y_true, y_score, *, k=None, sample_weight=None, ignore_ties=False
         Only consider the highest k scores in the ranking. If `None`, use all
         outputs.
 
-    sample_weight : ndarray of shape (n_samples,), default=None
+    sample_weight : array-like of shape (n_samples,), default=None
         Sample weights. If `None`, all samples are given the same weight.
 
     ignore_ties : bool, default=False
diff --git a/sklearn/tests/test_public_functions.py b/sklearn/tests/test_public_functions.py
@@ -180,6 +180,7 @@ def _check_function_param_validation(
     "sklearn.metrics.median_absolute_error",
     "sklearn.metrics.multilabel_confusion_matrix",
     "sklearn.metrics.mutual_info_score",
+    "sklearn.metrics.ndcg_score",
     "sklearn.metrics.pairwise.additive_chi2_kernel",
     "sklearn.metrics.precision_recall_curve",
     "sklearn.metrics.precision_recall_fscore_support",
