Repository: Lightning-AI/torchmetrics

Issue:
Enable `ROUGEScore` to evaluate hypotheses against multiple references.
## 🚀 Feature

Enable `ROUGEScore` to evaluate hypotheses against multiple references.

### Motivation

In the original paper, [Lin (2004)](https://aclanthology.org/W04-1013.pdf) proposes the evaluation of a hypothesis against the multiple references, and eventually, the maximum pairwise score is used.

### Pitch

Enable `ROUGEScore` to evaluate hypotheses against multiple references as it is the case for other text metrics.

### Alternatives

Leave it as it is.

### Additional context

Ideal for #new-contributors



Discussion:
I can take a look at this 
I started working on it. The requirements say: "maximum pairwise score is used". Depending on rouge_types - there are `1, 2, L`, different highest values, and consequently `precision, recall` and `fmeasure.` Which rouge_types' highest value should be used? `1, 2, L` and which metric `precision, recall` and `fmeasure`?

Also, should we have an avg. version also (instead of just maximum pairwise)? It can be passed as an argument during init.
   
> I can take a look at this


@stancld I have written the code for `maximum pairwise`. 
I wanted to know if we can quickly test just the `test_rouge.py` using `pytest`. `make test` seems to be taking a lot of time.
I'll run the full test once `test_rouge.py` works. 
Hi @ashutoshml, thanks a lot for your effort! O:] I'll have a look tomorrow, but you can run:
```
pytest tests/text/test_rouge.py
```
in the project directory to run ROUGE test only :] 

Resolving patch:
diff --git a/torchmetrics/functional/text/rouge.py b/torchmetrics/functional/text/rouge.py
--- a/torchmetrics/functional/text/rouge.py
+++ b/torchmetrics/functional/text/rouge.py
@@ -13,10 +13,11 @@
 # limitations under the License.
 import re
 from collections import Counter
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
 
 import torch
 from torch import Tensor, tensor
+from typing_extensions import Literal
 
 from torchmetrics.utilities.imports import _NLTK_AVAILABLE
 
@@ -33,6 +34,7 @@
     "rougeL": "L",
     "rougeLsum": "Lsum",
 }
+ALLOWED_ACCUMULATE_VALUES = ("avg", "best")
 
 
 def _add_newline_to_end_of_each_sentence(x: str) -> str:
@@ -68,7 +70,7 @@ def _compute_metrics(hits_or_lcs: int, pred_len: int, target_len: int) -> Dict[s
     return dict(precision=tensor(precision), recall=tensor(recall), fmeasure=tensor(fmeasure))
 
 
-def _lcs(pred_tokens: List[str], target_tokens: List[str]) -> int:
+def _lcs(pred_tokens: Sequence[str], target_tokens: Sequence[str]) -> int:
     """Common DP algorithm to compute the length of the longest common subsequence.
 
     Args:
@@ -87,7 +89,7 @@ def _lcs(pred_tokens: List[str], target_tokens: List[str]) -> int:
     return LCS[-1][-1]
 
 
-def _normalize_and_tokenize_text(text: str, stemmer: Optional[Any] = None) -> List[str]:
+def _normalize_and_tokenize_text(text: str, stemmer: Optional[Any] = None) -> Sequence[str]:
     """Rouge score should be calculated only over lowercased words and digits. Optionally, Porter stemmer can be
     used to strip word suffixes to improve matching. The text normalization follows the implemantion from `Rouge
     score_Text Normalizition`_
@@ -112,7 +114,7 @@ def _normalize_and_tokenize_text(text: str, stemmer: Optional[Any] = None) -> Li
     return tokens
 
 
-def _rouge_n_score(pred: List[str], target: List[str], n_gram: int) -> Dict[str, Tensor]:
+def _rouge_n_score(pred: Sequence[str], target: Sequence[str], n_gram: int) -> Dict[str, Tensor]:
     """This computes precision, recall and F1 score for the Rouge-N metric.
 
     Args:
@@ -124,7 +126,7 @@ def _rouge_n_score(pred: List[str], target: List[str], n_gram: int) -> Dict[str,
             N-gram overlap.
     """
 
-    def _create_ngrams(tokens: List[str], n: int) -> Counter:
+    def _create_ngrams(tokens: Sequence[str], n: int) -> Counter:
         ngrams: Counter = Counter()
         for ngram in (tuple(tokens[i : i + n]) for i in range(len(tokens) - n + 1)):
             ngrams[ngram] += 1
@@ -140,7 +142,7 @@ def _create_ngrams(tokens: List[str], n: int) -> Counter:
     return _compute_metrics(hits, max(pred_len, 1), max(target_len, 1))
 
 
-def _rouge_l_score(pred: List[str], target: List[str]) -> Dict[str, Tensor]:
+def _rouge_l_score(pred: Sequence[str], target: Sequence[str]) -> Dict[str, Tensor]:
     """This computes precision, recall and F1 score for the Rouge-L or Rouge-LSum metric.
 
     Args:
@@ -158,9 +160,10 @@ def _rouge_l_score(pred: List[str], target: List[str]) -> Dict[str, Tensor]:
 
 
 def _rouge_score_update(
-    preds: List[str],
-    targets: List[str],
+    preds: Sequence[str],
+    targets: Sequence[Sequence[str]],
     rouge_keys_values: List[Union[int, str]],
+    accumulate: str,
     stemmer: Optional[Any] = None,
 ) -> Dict[Union[int, str], List[Dict[str, Tensor]]]:
     """Update the rouge score with the current set of predicted and target sentences.
@@ -169,9 +172,14 @@ def _rouge_score_update(
         preds:
             An iterable of predicted sentences.
         targets:
-            An iterable of target sentences.
+            An iterable of iterable of target sentences.
         rouge_keys_values:
             List of N-grams/'L'/'Lsum' arguments.
+        accumulate:
+            Useful incase of multi-reference rouge score.
+            ``avg`` takes the avg of all references with respect to predictions
+            ``best`` takes the best fmeasure score obtained between prediction and multiple corresponding references.
+            Allowed values are ``avg`` and ``best``.
         stemmer:
             Porter stemmer instance to strip word suffixes to improve matching.
 
@@ -179,7 +187,7 @@ def _rouge_score_update(
         >>> targets = "Is your name John".split()
         >>> preds = "My name is John".split()
         >>> from pprint import pprint
-        >>> score = _rouge_score_update(preds, targets, rouge_keys_values=[1, 2, 3, 'L'])
+        >>> score = _rouge_score_update(preds, targets, rouge_keys_values=[1, 2, 3, 'L'], accumulate='best')
         >>> pprint(score)  # doctest: +SKIP
         {1: [{'fmeasure': tensor(0.), 'precision': tensor(0.), 'recall': tensor(0.)},
             {'fmeasure': tensor(0.), 'precision': tensor(0.), 'recall': tensor(0.)},
@@ -199,24 +207,62 @@ def _rouge_score_update(
             {'fmeasure': tensor(1.), 'precision': tensor(1.), 'recall': tensor(1.)}]}
     """
     results: Dict[Union[int, str], List[Dict[str, Tensor]]] = {rouge_key: [] for rouge_key in rouge_keys_values}
+
     for pred_raw, target_raw in zip(preds, targets):
+        result_inner: Dict[Union[int, str], Dict[str, Tensor]] = {rouge_key: {} for rouge_key in rouge_keys_values}
+        result_avg: Dict[Union[int, str], List[Dict[str, Tensor]]] = {rouge_key: [] for rouge_key in rouge_keys_values}
+        list_results = []
         pred = _normalize_and_tokenize_text(pred_raw, stemmer)
-        target = _normalize_and_tokenize_text(target_raw, stemmer)
+        pred_Lsum = _normalize_and_tokenize_text(_add_newline_to_end_of_each_sentence(pred_raw), stemmer)
 
-        if "Lsum" in rouge_keys_values:
-            # rougeLsum expects "\n" separated sentences within a summary
-            pred_Lsum = _normalize_and_tokenize_text(_add_newline_to_end_of_each_sentence(pred_raw), stemmer)
-            target_Lsum = _normalize_and_tokenize_text(_add_newline_to_end_of_each_sentence(target_raw), stemmer)
+        for target_raw_inner in target_raw:
+            target = _normalize_and_tokenize_text(target_raw_inner, stemmer)
 
-        for rouge_key in rouge_keys_values:
-            if isinstance(rouge_key, int):
-                score = _rouge_n_score(pred, target, rouge_key)
-            else:
-                score = _rouge_l_score(
-                    pred if rouge_key != "Lsum" else pred_Lsum,
-                    target if rouge_key != "Lsum" else target_Lsum,
+            if "Lsum" in rouge_keys_values:
+                # rougeLsum expects "\n" separated sentences within a summary
+                target_Lsum = _normalize_and_tokenize_text(
+                    _add_newline_to_end_of_each_sentence(target_raw_inner), stemmer
                 )
-            results[rouge_key].append(score)
+
+            for rouge_key in rouge_keys_values:
+                if isinstance(rouge_key, int):
+                    score = _rouge_n_score(pred, target, rouge_key)
+                else:
+                    score = _rouge_l_score(
+                        pred if rouge_key != "Lsum" else pred_Lsum,
+                        target if rouge_key != "Lsum" else target_Lsum,
+                    )
+                result_inner[rouge_key] = score
+                result_avg[rouge_key].append(score)
+            list_results.append(result_inner.copy())
+
+        if accumulate == "best":
+            key_curr = rouge_keys_values[0]
+            all_fmeasure = torch.tensor([v[key_curr]["fmeasure"] for v in list_results])
+            highest_idx = int(torch.argmax(all_fmeasure).item())
+
+            for rouge_key in rouge_keys_values:
+                results[rouge_key].append(list_results[highest_idx][rouge_key])
+
+        elif accumulate == "avg":
+            new_result_avg: Dict[Union[int, str], Dict[str, Tensor]] = {
+                rouge_key: {} for rouge_key in rouge_keys_values
+            }
+            for rouge_key, metrics in result_avg.items():
+                _dict_metric_score_batch: Dict[str, List[Tensor]] = {}
+                for metric in metrics:
+                    for _type, value in metric.items():
+                        if _type not in _dict_metric_score_batch:
+                            _dict_metric_score_batch[_type] = []
+                        _dict_metric_score_batch[_type].append(value)
+
+                new_result_avg[rouge_key] = {
+                    _type: torch.tensor(_dict_metric_score_batch[_type]).mean() for _type in _dict_metric_score_batch
+                }
+
+            for rouge_key in rouge_keys_values:
+                results[rouge_key].append(new_result_avg[rouge_key])
+
     return results
 
 
@@ -239,8 +285,9 @@ def _rouge_score_compute(sentence_results: Dict[str, List[Tensor]]) -> Dict[str,
 
 
 def rouge_score(
-    preds: Union[str, List[str]],
-    targets: Union[str, List[str]],
+    preds: Union[str, Sequence[str]],
+    targets: Union[str, Sequence[str], Sequence[Sequence[str]]],
+    accumulate: Literal["avg", "best"] = "best",
     use_stemmer: bool = False,
     rouge_keys: Union[str, Tuple[str, ...]] = ("rouge1", "rouge2", "rougeL", "rougeLsum"),  # type: ignore
 ) -> Dict[str, Tensor]:
@@ -250,7 +297,11 @@ def rouge_score(
         preds:
             An iterable of predicted sentences or a single predicted sentence.
         targets:
-            An iterable of target sentences or a single target sentence.
+            An iterable of iterables of target sentences or an iterable of target sentences or a single target sentence.
+        accumulate:
+            Useful incase of multi-reference rouge score.
+            - ``avg`` takes the avg of all references with respect to predictions
+            - ``best`` takes the best fmeasure score obtained between prediction and multiple corresponding references.
         use_stemmer:
             Use Porter stemmer to strip word suffixes to improve matching.
         rouge_keys:
@@ -303,14 +354,17 @@ def rouge_score(
             raise ValueError(f"Got unknown rouge key {key}. Expected to be one of {list(ALLOWED_ROUGE_KEYS.keys())}")
     rouge_keys_values = [ALLOWED_ROUGE_KEYS[key] for key in rouge_keys]
 
+    if isinstance(targets, list) and all(isinstance(target, str) for target in targets):
+        targets = [targets] if isinstance(preds, str) else [[target] for target in targets]
+
     if isinstance(preds, str):
         preds = [preds]
 
     if isinstance(targets, str):
-        targets = [targets]
+        targets = [[targets]]
 
     sentence_results: Dict[Union[int, str], List[Dict[str, Tensor]]] = _rouge_score_update(
-        preds, targets, rouge_keys_values, stemmer=stemmer
+        preds, targets, rouge_keys_values, stemmer=stemmer, accumulate=accumulate
     )
 
     output: Dict[str, List[Tensor]] = {}
diff --git a/torchmetrics/text/rouge.py b/torchmetrics/text/rouge.py
--- a/torchmetrics/text/rouge.py
+++ b/torchmetrics/text/rouge.py
@@ -11,12 +11,18 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
 
 from torch import Tensor
+from typing_extensions import Literal
 
 from torchmetrics import Metric
-from torchmetrics.functional.text.rouge import ALLOWED_ROUGE_KEYS, _rouge_score_compute, _rouge_score_update
+from torchmetrics.functional.text.rouge import (
+    ALLOWED_ACCUMULATE_VALUES,
+    ALLOWED_ROUGE_KEYS,
+    _rouge_score_compute,
+    _rouge_score_update,
+)
 from torchmetrics.utilities.imports import _NLTK_AVAILABLE
 
 
@@ -27,6 +33,10 @@ class ROUGEScore(Metric):
     Args:
         use_stemmer:
             Use Porter stemmer to strip word suffixes to improve matching.
+        accumulate:
+            Useful incase of multi-reference rouge score.
+            - ``avg`` takes the avg of all references with respect to predictions
+            - ``best`` takes the best fmeasure score obtained between prediction and multiple corresponding references.
         rouge_keys:
             A list of rouge types to calculate.
             Keys that are allowed are ``rougeL``, ``rougeLsum``, and ``rouge1`` through ``rouge9``.
@@ -76,6 +86,7 @@ class ROUGEScore(Metric):
     def __init__(
         self,
         use_stemmer: bool = False,
+        accumulate: Literal["avg", "best"] = "best",
         rouge_keys: Union[str, Tuple[str, ...]] = ("rouge1", "rouge2", "rougeL", "rougeLsum"),  # type: ignore
         compute_on_step: bool = True,
         dist_sync_on_step: bool = False,
@@ -99,31 +110,44 @@ def __init__(
             if key not in ALLOWED_ROUGE_KEYS:
                 raise ValueError(f"Got unknown rouge key {key}. Expected to be one of {ALLOWED_ROUGE_KEYS}")
 
+        if accumulate not in ALLOWED_ACCUMULATE_VALUES:
+            raise ValueError(
+                f"Got unknown accumulate value {accumulate}. Expected to be one of {ALLOWED_ACCUMULATE_VALUES}"
+            )
+
         self.rouge_keys = rouge_keys
         self.rouge_keys_values = [ALLOWED_ROUGE_KEYS[key] for key in rouge_keys]
         self.stemmer = nltk.stem.porter.PorterStemmer() if use_stemmer else None
+        self.accumulate = accumulate
 
         # Adding stated dynamically to prevent IndexError during sync function as some lists can be empty.
         for rouge_key in self.rouge_keys:
             for score in ["fmeasure", "precision", "recall"]:
                 self.add_state(f"{rouge_key}_{score}", [], dist_reduce_fx=None)
 
-    def update(self, preds: Union[str, List[str]], targets: Union[str, List[str]]) -> None:  # type: ignore
+    def update(  # type: ignore
+        self, preds: Union[str, Sequence[str]], targets: Union[str, Sequence[str], Sequence[Sequence[str]]]
+    ) -> None:
         """Compute rouge scores.
 
         Args:
-            preds: An iterable of predicted sentences or a single predicted sentence.
-            targets: An iterable of target sentences or a single target sentence.
+            preds:
+                An iterable of predicted sentences or a single predicted sentence.
+            targets:
+                An iterable of iterable of target sentences or an iterable
+                of target sentences or a single target sentence.
         """
+        if isinstance(targets, list) and all(isinstance(target, str) for target in targets):
+            targets = [targets] if isinstance(preds, str) else [[target] for target in targets]
 
         if isinstance(preds, str):
             preds = [preds]
 
         if isinstance(targets, str):
-            targets = [targets]
+            targets = [[targets]]
 
         output: Dict[Union[int, str], List[Dict[str, Tensor]]] = _rouge_score_update(
-            preds, targets, self.rouge_keys_values, stemmer=self.stemmer
+            preds, targets, self.rouge_keys_values, stemmer=self.stemmer, accumulate=self.accumulate
         )
         for rouge_key, metrics in output.items():
             for metric in metrics:
@@ -147,7 +171,6 @@ def __hash__(self) -> int:
         # override to hash list objects.
         # this is a bug in the upstream pytorch release.
         hash_vals = [self.__class__.__name__]
-
         for key in self._defaults:
             value = getattr(self, key)
             if isinstance(value, list):


Test patch:
diff --git a/tests/text/inputs.py b/tests/text/inputs.py
--- a/tests/text/inputs.py
+++ b/tests/text/inputs.py
@@ -42,6 +42,8 @@
 
 _inputs_multiple_references = Input(preds=TUPLE_OF_HYPOTHESES, targets=TUPLE_OF_REFERENCES)
 
+_inputs_single_sentence_single_reference = Input(preds=HYPOTHESIS_B, targets=REFERENCE_1B)
+
 ERROR_RATES_BATCHES_1 = {
     "preds": [["hello world"], ["what a day"]],
     "targets": [["hello world"], ["what a wonderful day"]],
diff --git a/tests/text/test_rouge.py b/tests/text/test_rouge.py
--- a/tests/text/test_rouge.py
+++ b/tests/text/test_rouge.py
@@ -13,11 +13,13 @@
 # limitations under the License.
 
 from functools import partial
-from typing import List
+from typing import Sequence
 
 import pytest
+import torch
 
 from tests.text.helpers import INPUT_ORDER, TextTester
+from tests.text.inputs import _inputs_multiple_references, _inputs_single_sentence_single_reference
 from torchmetrics.functional.text.rouge import rouge_score
 from torchmetrics.text.rouge import ROUGEScore
 from torchmetrics.utilities.imports import _NLTK_AVAILABLE, _ROUGE_SCORE_AVAILABLE
@@ -30,34 +32,45 @@
 
 ROUGE_KEYS = ("rouge1", "rouge2", "rougeL", "rougeLsum")
 
-SINGLE_SENTENCE_EXAMPLE_PREDS = "The quick brown fox jumps over the lazy dog"
-SINGLE_SENTENCE_EXAMPLE_TARGET = "The quick brown dog jumps on the log."
 
-PREDS = "My name is John"
-TARGETS = "Is your name John"
+def _compute_rouge_score(
+    preds: Sequence[str],
+    targets: Sequence[Sequence[str]],
+    use_stemmer: bool,
+    rouge_level: str,
+    metric: str,
+    accumulate: str,
+):
+    """Evaluates rouge scores from rouge-score package for baseline evaluation."""
+    if isinstance(targets, list) and all(isinstance(target, str) for target in targets):
+        targets = [targets] if isinstance(preds, str) else [[target] for target in targets]
 
-
-BATCHES_1 = {
-    "preds": [["the cat was under the bed"], ["the cat was found under the bed"]],
-    "targets": [["the cat was found under the bed"], ["the tiny little cat was found under the big funny bed "]],
-}
-
-
-BATCHES_2 = {
-    "preds": [["The quick brown fox jumps over the lazy dog"], ["My name is John"]],
-    "targets": [["The quick brown dog jumps on the log."], ["Is your name John"]],
-}
-
-
-def _compute_rouge_score(preds: List[str], targets: List[str], use_stemmer: bool, rouge_level: str, metric: str):
     if isinstance(preds, str):
         preds = [preds]
+
     if isinstance(targets, str):
-        targets = [targets]
+        targets = [[targets]]
+
     scorer = RougeScorer(ROUGE_KEYS, use_stemmer=use_stemmer)
     aggregator = BootstrapAggregator()
-    for pred, target in zip(preds, targets):
-        aggregator.add_scores(scorer.score(target, pred))
+
+    for target_raw, pred_raw in zip(targets, preds):
+        list_results = [scorer.score(target, pred_raw) for target in target_raw]
+        aggregator_avg = BootstrapAggregator()
+
+        if accumulate == "best":
+            key_curr = list(list_results[0].keys())[0]
+            all_fmeasure = torch.tensor([v[key_curr].fmeasure for v in list_results])
+            highest_idx = torch.argmax(all_fmeasure).item()
+            aggregator.add_scores(list_results[highest_idx])
+        elif accumulate == "avg":
+            for _score in list_results:
+                aggregator_avg.add_scores(_score)
+            _score = {rouge_key: scores.mid for rouge_key, scores in aggregator_avg.aggregate().items()}
+            aggregator.add_scores(_score)
+        else:
+            raise ValueError(f"Got unknown accumulate value {accumulate}. Expected to be one of ['best', 'avg']")
+
     rs_scores = aggregator.aggregate()
     rs_result = getattr(rs_scores[rouge_level].mid, metric)
     return rs_result
@@ -84,19 +97,21 @@ def _compute_rouge_score(preds: List[str], targets: List[str], use_stemmer: bool
 @pytest.mark.parametrize(
     ["preds", "targets"],
     [
-        (BATCHES_1["preds"], BATCHES_1["targets"]),
-        (BATCHES_2["preds"], BATCHES_2["targets"]),
+        (_inputs_multiple_references.preds, _inputs_multiple_references.targets),
     ],
 )
+@pytest.mark.parametrize("accumulate", ["avg", "best"])
 class TestROUGEScore(TextTester):
     @pytest.mark.parametrize("ddp", [False, True])
     @pytest.mark.parametrize("dist_sync_on_step", [False, True])
-    def test_rouge_score_class(self, ddp, dist_sync_on_step, preds, targets, pl_rouge_metric_key, use_stemmer):
-        metric_args = {"use_stemmer": use_stemmer}
-
+    def test_rouge_score_class(
+        self, ddp, dist_sync_on_step, preds, targets, pl_rouge_metric_key, use_stemmer, accumulate
+    ):
+        metric_args = {"use_stemmer": use_stemmer, "accumulate": accumulate}
         rouge_level, metric = pl_rouge_metric_key.split("_")
-        rouge_metric = partial(_compute_rouge_score, use_stemmer=use_stemmer, rouge_level=rouge_level, metric=metric)
-
+        rouge_metric = partial(
+            _compute_rouge_score, use_stemmer=use_stemmer, rouge_level=rouge_level, metric=metric, accumulate=accumulate
+        )
         self.run_class_metric_test(
             ddp=ddp,
             preds=preds,
@@ -109,12 +124,13 @@ def test_rouge_score_class(self, ddp, dist_sync_on_step, preds, targets, pl_roug
             key=pl_rouge_metric_key,
         )
 
-    def test_rouge_score_functional(self, preds, targets, pl_rouge_metric_key, use_stemmer):
-        metric_args = {"use_stemmer": use_stemmer}
+    def test_rouge_score_functional(self, preds, targets, pl_rouge_metric_key, use_stemmer, accumulate):
+        metric_args = {"use_stemmer": use_stemmer, "accumulate": accumulate}
 
         rouge_level, metric = pl_rouge_metric_key.split("_")
-        rouge_metric = partial(_compute_rouge_score, use_stemmer=use_stemmer, rouge_level=rouge_level, metric=metric)
-
+        rouge_metric = partial(
+            _compute_rouge_score, use_stemmer=use_stemmer, rouge_level=rouge_level, metric=metric, accumulate=accumulate
+        )
         self.run_functional_metric_test(
             preds,
             targets,
@@ -144,4 +160,9 @@ def test_rouge_metric_wrong_key_value_error():
         ROUGEScore(rouge_keys=key)
 
     with pytest.raises(ValueError):
-        rouge_score(PREDS, TARGETS, rouge_keys=key)
+        rouge_score(
+            _inputs_single_sentence_single_reference.preds,
+            _inputs_single_sentence_single_reference.targets,
+            rouge_keys=key,
+            accumulate="best",
+        )