How to calculate NDCG with binary relevances using sklearn?

Question

I'm trying to calculate the NDCG score for binary relevances:

from sklearn.metrics import ndcg_score
y_true = [0, 1, 0]
y_pred = [0, 1, 0]
ndcg_score(y_true, y_pred)

And getting:

ValueError: Only ('multilabel-indicator', 'continuous-multioutput',  
'multiclass-multioutput') formats are supported. Got binary instead

Is there a way to make this work?

Answer 1

Please try:

from sklearn.metrics import ndcg_score
y_true = [[0, 1, 0]]
y_pred = [[0, 1, 0]]
ndcg_score(y_true, y_pred)
1.0

Note the expected shapes in the docs:

y_true: ndarray, shape (n_samples, n_labels)
y_score: ndarray, shape (n_samples, n_labels)