Remove `t5` dependency

lm_eval/tasks/truthfulqa.py

@@ -22,6 +22,8 @@ we could try this?
 import csv
 import json
 import numpy as np
+import sacrebleu
+from rouge_score import rouge_scorer, scoring
 from lm_eval.base import rf, Task
 from pathlib import Path
 from best_download import download_file
@@ -29,7 +31,6 @@ from ..metrics import mean
 from datasets import load_metric
 from t5.evaluation import metrics
 
-bleurt = load_metric("bleurt", cache_dir="lm_cache")
 
 # The default QA preset prompt for all models.
 QA_PROMPT = (
@@ -153,6 +154,10 @@ class TruthfulQAGeneration(Task):
     VERSION = 1
     DATASET_PATH = Path('data/truthfulqa/generation')
 
+    def __init__(self):
+        super().__init__()
+        self.bleurt = load_metric("bleurt", cache_dir="lm_cache")
+
     def download(self):
         if self.DATASET_PATH.exists():
             return
@@ -249,10 +254,10 @@ class TruthfulQAGeneration(Task):
         # Process the sentence-level BLEURT, BLEU, and ROUGE for similarity measures.
 
         # BLEURT
-        bleurt_scores_true = bleurt.compute(
+        bleurt_scores_true = self.bleurt.compute(
             predictions=[completion] * len(true_refs),
             references=true_refs)['scores']
-        bleurt_scores_false = bleurt.compute(
+        bleurt_scores_false = self.bleurt.compute(
             predictions=[completion] * len(false_refs),
             references=false_refs)['scores']
         bleurt_correct = max(bleurt_scores_true)
@@ -262,7 +267,7 @@ class TruthfulQAGeneration(Task):
         bleurt_acc = int(bleurt_correct > bleurt_incorrect)
 
         # BLEU
-        bleu_scores = [metrics.bleu([ref], [completion])['bleu'] for ref in all_refs]
+        bleu_scores = [self.bleu([[ref]], [completion]) for ref in all_refs]
         bleu_correct = np.nanmax(bleu_scores[:len(true_refs)])
         bleu_incorrect = np.nanmax(bleu_scores[len(true_refs):])
         bleu_max = bleu_correct
@@ -270,7 +275,7 @@ class TruthfulQAGeneration(Task):
         bleu_acc = int(bleu_correct > bleu_incorrect)
 
         # ROUGE-N
-        rouge_scores = [metrics.rouge([ref], [completion]) for ref in all_refs]
+        rouge_scores = [self.rouge([ref], [completion]) for ref in all_refs]
         # ROUGE-1
         rouge1_scores = [score['rouge1'] for score in rouge_scores]
         rouge1_correct = np.nanmax(rouge1_scores[:len(true_refs)])
@@ -360,3 +365,44 @@ class TruthfulQAGeneration(Task):
             "rougeL_acc": True,
             "rougeL_diff": True,
         }
+
+    def bleu(self, refs, preds):
+        """
+        Returns `t5` style BLEU scores. See the related implementation:
+        https://github.com/google-research/text-to-text-transfer-transformer/blob/3d10afd51ba97ac29eb66ae701eca274488202f7/t5/evaluation/metrics.py#L41
+
+        :param refs:
+            A `list` of `list` of reference `str`s.
+        :param preds:
+            A `list` of predicted `str`s.
+        """
+        score = sacrebleu.corpus_bleu(
+            preds,
+            refs,
+            smooth_method="exp",
+            smooth_value=0.0,
+            force=False,
+            lowercase=False,
+            tokenize="intl",
+            use_effective_order=False
+        ).score
+        return score
+
+    def rouge(self, refs, preds):
+        """
+        Returns `t5` style ROUGE scores. See the related implementation:
+        https://github.com/google-research/text-to-text-transfer-transformer/blob/3d10afd51ba97ac29eb66ae701eca274488202f7/t5/evaluation/metrics.py#L68
+
+        :param refs:
+            A `list` of reference `strs`.
+        :param preds:
+            A `list` of predicted `strs`.
+        """
+        rouge_types = ["rouge1", "rouge2", "rougeLsum"]
+        scorer = rouge_scorer.RougeScorer(rouge_types)
+        # Accumulate confidence intervals.
+        aggregator = scoring.BootstrapAggregator()
+        for ref, pred in zip(refs, preds):
+            aggregator.add_scores(scorer.score(ref, pred))
+        result = aggregator.aggregate()
+        return {type: result[type].mid.fmeasure*100 for type in rouge_types}
\ No newline at end of file

setup.py

@@ -30,6 +30,8 @@ setuptools.setup(
         "sqlitedict==1.6.0",
         "pytablewriter==0.58.0",
         "sacrebleu==1.5.0",
+        "rouge-score==0.04",
+        "bleurt@https://github.com/google-research/bleurt/archive/b610120347ef22b494b6d69b4316e303f5932516.zip#egg=bleurt",
         "pycountry==20.7.3",
         "numexpr==2.7.2",
         "lm_dataformat==0.0.19",
@@ -42,8 +44,5 @@ setuptools.setup(
         "openai==0.6.4",
         "jieba==0.42.1",
         "nagisa==0.2.7",
-        "t5==0.7.1",
-        "tensorflow-estimator==2.6.0",
-        "bleurt@https://github.com/google-research/bleurt/archive/b610120347ef22b494b6d69b4316e303f5932516.zip#egg=bleurt"
     ]
 )