From e8596a3aae451e78cef3ae8c1f2fd7268ed08fd8 Mon Sep 17 00:00:00 2001 From: Aditya Barua Date: Fri, 11 Feb 2022 14:28:12 -0800 Subject: [PATCH] Update rouge metric to take in an SPM model for tokenization. PiperOrigin-RevId: 428091414 --- t5/evaluation/metrics.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/t5/evaluation/metrics.py b/t5/evaluation/metrics.py index f51ef62f..908cf828 100644 --- a/t5/evaluation/metrics.py +++ b/t5/evaluation/metrics.py @@ -65,20 +65,22 @@ def bleu(targets, predictions): return {"bleu": bleu_score.score} -def rouge(targets, predictions, score_keys=None): +def rouge(targets, predictions, score_keys=None, spm_model=None): """Computes rouge score. Args: targets: list of strings predictions: list of strings score_keys: list of strings with the keys to compute. + spm_model: string, path to SentencePieceModel model. If provided, this model + is used for tokenizing the targets and predictions. 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, spm_model=spm_model) aggregator = scoring.BootstrapAggregator() def _prepare_summary(summary):