select option

rouge_raw.py

@@ -37,7 +37,7 @@ Module for raw ROUGE score calculation from:
 """
 
 import re
-from typing import Sequence
+from typing import Sequence, Optional
 
 import datasets
 import evaluate
@@ -175,6 +175,8 @@ ROCUE RAW metric for list of predictions and references.
 Args:
     predictions: list of predictions to evaluate. Each prediction should be a string with tokens separated by spaces.
     references: list of reference for each prediction. Each reference should be a string with tokens separated by spaces.
+    select: (Optional) string. The name of the metric to return. One of: 'rougeraw1_precision', 'rougeraw1_recall', 'rougeraw1_fmeasure', 'rougeraw2_precision', 'rougeraw2_recall', 'rougeraw2_fmeasure', 'rougerawl_precision', 'rougerawl_recall', 'rougerawl_fmeasure'.
+        If None, all metrics are returned as a dictionary.
 Returns:
     rougeraw1_precision
     rougeraw1_recall
@@ -215,9 +217,9 @@ class RougeRaw(evaluate.Metric):
             ],
         )
 
-    def _compute(self, predictions: Sequence[str], references: Sequence[str]):
+    def _compute(self, predictions: Sequence[str], references: Sequence[str], select: Optional[str] = None):
         res = RougeRawOriginal().corpus(references, predictions)
-        return {
+        res = {
             "rougeraw1_precision": res["1"].p,
             "rougeraw1_recall": res["1"].r,
             "rougeraw1_fmeasure": res["1"].f,
@@ -229,3 +231,7 @@ class RougeRaw(evaluate.Metric):
             "rougerawl_fmeasure": res["L"].f,
         }
 
+        if select is not None:
+            return res[select]
+        return res
+