generalization in eval_args

generation_evaluator.py

@@ -5,6 +5,7 @@ import numpy as np
 import spacy
 import torch
 from alignscore import AlignScore
+from transformers import AutoTokenizer
 
 _CITATION = """\
 @inproceedings{lin-2004-rouge,
@@ -155,9 +156,7 @@ class GenerationEvaluator(evaluate.Metric):
         # Download AlignScore model and move to GPU if possible
         model_path = dl_manager.download(ALIGNSCORE_ARGS["ckpt_path"])
         ALIGNSCORE_ARGS["ckpt_path"] = model_path
-        ALIGNSCORE_ARGS["device"] = (
-            "cuda:0" if torch.cuda.is_available() else "cpu"
-        )
+        ALIGNSCORE_ARGS["device"] = "cuda:0" if torch.cuda.is_available() else "cpu"
         self.align_scorer = AlignScore(**ALIGNSCORE_ARGS)
 
         # Prepare scorers
@@ -167,20 +166,33 @@ class GenerationEvaluator(evaluate.Metric):
         self.bert_scorer = evaluate.load("bertscore")
         self.chrf_scorer = evaluate.load("chrf")
 
-    def _compute(self, predictions, references, tokenizer=None):
+    def _compute(self, predictions, references, **eval_kwargs):
+        tokenizer_name = eval_kwargs.pop("tokenizer_name", None)
+        tokenizer = None
+
+        if tokenizer_name is not None:
+            tks = AutoTokenizer.from_pretrained(tokenizer_name)
+            tokenizer = tks.tokenize
+
         # Compute ROUGE
         rouge_results = self.rouge_scorer.compute(
-            predictions=predictions, references=references, tokenizer=tokenizer
+            predictions=predictions,
+            references=references,
+            tokenizer=tokenizer,
+            **eval_kwargs
         )
 
         # Compute BLEU
         if tokenizer is None:
             bleu_results = self.bleu_scorer.compute(
-                predictions=predictions, references=references
+                predictions=predictions, references=references, **eval_kwargs
             )
         else:
             bleu_results = self.bleu_scorer.compute(
-                predictions=predictions, references=references, tokenizer=tokenizer
+                predictions=predictions,
+                references=references,
+                tokenizer=tokenizer,
+                **eval_kwargs
             )
 
         # Compute Exact Match
@@ -203,9 +215,7 @@ class GenerationEvaluator(evaluate.Metric):
 
         # Compute AlignScore
         align_score = round(
-            np.mean(
-                self.align_scorer.score(contexts=references, claims=predictions)
-            ),
+            np.mean(self.align_scorer.score(contexts=references, claims=predictions)),
             4,
         )
 

requirements.txt

@@ -7,4 +7,5 @@ rouge_score
 numpy
 sacrebleu
 git+https://github.com/yuh-zha/AlignScore.git
-spacy
+spacy
+transformers