add docstring for metrics/ranking

pyhealth/metrics/ranking.py

@@ -4,6 +4,33 @@
 def ranking_metrics_fn(qrels: Dict[str, Dict[str, int]],
  results: Dict[str, Dict[str, float]],
  k_values: List[int]) -> Dict[str, float]:
+ """Computes metrics for ranking tasks.
+
+ Args:
+ qrels: Ground truth. A dictionary of query ids and their corresponding
+ relevance judgements. The relevance judgements are a dictionary of
+ document ids and their corresponding relevance scores.
+ results: Ranked results. A dictionary of query ids and their corresponding
+ document scores. The document scores are a dictionary of document ids and
+ their corresponding scores.
+ k_values: A list of integers specifying the cutoffs for the metrics.
+
+ Returns:
+ A dictionary of metrics and their corresponding values.
+
+ Examples:
+ >>> qrels = {
+ ... "q1": {"d1": 1, "d2": 0, "d3": 1},
+ ... "q2": {"d1": 1, "d2": 1, "d3": 0}
+ ... }
+ >>> results = {
+ ... "q1": {"d1": 0.5, "d2": 0.2, "d3": 0.1},
+ ... "q2": {"d1": 0.1, "d2": 0.2, "d3": 0.5}
+ ... }
+ >>> k_values = [1, 2]
+ >>> ranking_metrics_fn(qrels, results, k_values)
+ {'NDCG@1': 0.5, 'MAP@1': 0.25, 'Recall@1': 0.25, 'P@1': 0.5, 'NDCG@2': 0.5, 'MAP@2': 0.375, 'Recall@2': 0.5, 'P@2': 0.5}
+ """
  import pytrec_eval
  ret = {}
 
@@ -36,3 +63,9 @@ def ranking_metrics_fn(qrels: Dict[str, Dict[str, int]],
  ret[f"P@{k}"] = round(ret[f"P@{k}"] / len(scores), 5)
 
  return ret
+
+
+if __name__ == "__main__":
+ import doctest
+
+ doctest.testmod(verbose=True)