Skip to content

Commit

Permalink
Update rouge function to take in a tokenizer.
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 429416862
  • Loading branch information
Aditya Barua authored and t5-copybara committed Feb 18, 2022
1 parent e1ffda6 commit a14d43a
Showing 1 changed file with 5 additions and 2 deletions.
7 changes: 5 additions & 2 deletions t5/evaluation/metrics.py
Original file line number Diff line number Diff line change
Expand Up @@ -65,20 +65,23 @@ def bleu(targets, predictions):
return {"bleu": bleu_score.score}


def rouge(targets, predictions, score_keys=None):
def rouge(targets, predictions, score_keys=None, tokenizer=None):
"""Computes rouge score.
Args:
targets: list of strings
predictions: list of strings
score_keys: list of strings with the keys to compute.
tokenizer: Tokenizer to use. Should be an instance of
third_party/google_research/google_research/rouge/tokenizers.Tokenizer. If
not specified, tokenizers.DefaultTokenizer is used.
Returns:
dict with score_key: rouge score across all targets and predictions
"""

if score_keys is None:
score_keys = ["rouge1", "rouge2", "rougeLsum"]
scorer = rouge_scorer.RougeScorer(score_keys)
scorer = rouge_scorer.RougeScorer(score_keys, tokenizer=tokenizer)
aggregator = scoring.BootstrapAggregator()

def _prepare_summary(summary):
Expand Down

0 comments on commit a14d43a

Please sign in to comment.