update

src/backend/tasks/xsum/task.py

@@ -3,7 +3,51 @@ from lm_eval.api.instance import Instance
 from lm_eval.api.registry import register_task
 from lm_eval.api.metrics import mean
 
-from src.backend.tasks.xsum import utils
+import sacrebleu
+from rouge_score import rouge_scorer, scoring
+
+
+def bleu(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(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)
+    # Add newlines between sentences to correctly compute `rougeLsum`.
+
+    def _prepare_summary(summary):
+        summary = summary.replace(" . ", ".\n")
+        return summary
+
+    # Accumulate confidence intervals.
+    aggregator = scoring.BootstrapAggregator()
+    for ref, pred in zip(refs, preds):
+        ref = _prepare_summary(ref)
+        pred = _prepare_summary(pred)
+        aggregator.add_scores(scorer.score(ref, pred))
+    result = aggregator.aggregate()
+    return {type: result[type].mid.fmeasure * 100 for type in rouge_types}
 
 
 @register_task("xsum")
@@ -14,7 +58,14 @@ class XSum(Task):
 
     def __init__(self, data_dir=None, cache_dir=None, download_mode=None, config=None):
         super().__init__(data_dir=data_dir, cache_dir=cache_dir, download_mode=download_mode, config=config)
-        print('XXX XSum!')
+        self.factkb_tokenizer = None
+        self.factkb_model = None
+
+    def init_factkb(self):
+        from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+        self.factkb_tokenizer = AutoTokenizer.from_pretrained("roberta-base", padding="max_length", truncation=True)
+        self.factkb_model = AutoModelForSequenceClassification.from_pretrained("bunsenfeng/FactKB", num_labels=2)
 
     def has_training_docs(self):
         return True
@@ -70,7 +121,31 @@ class XSum(Task):
         ]
 
     def process_results(self, doc, results):
-        return utils.process_results(doc, results)
+        completion = results[0]
+
+        # document = doc["document"]
+        true_refs = [doc["summary"]]
+        all_refs = true_refs
+
+        # ROUGE-N
+        rouge_scores = [rouge([ref], [completion]) for ref in all_refs]
+
+        # ROUGE-1
+        rouge1_scores = [score["rouge1"] for score in rouge_scores]
+
+        # ROUGE-2
+        rouge2_scores = [score["rouge2"] for score in rouge_scores]
+
+        # ROUGE-L
+        rougeL_scores = [score["rougeLsum"] for score in rouge_scores]
+
+        res = {
+            "rouge1": rouge1_scores[0],
+            "rouge2": rouge2_scores[0],
+            "rougeL": rougeL_scores[0],
+        }
+
+        return res
 
     def aggregation(self):
         """
@@ -87,4 +162,3 @@ class XSum(Task):
             whether a higher value of the submetric is better
         """
         return {k: True for k in ["rouge1", "rouge2", "rougeL"]}
-

src/backend/tasks/xsum/utils.py

@@ -1,89 +0,0 @@
-import sacrebleu
-import numpy as np
-
-from rouge_score import rouge_scorer, scoring
-
-
-def process_results(doc, results):
-    # (Pdb)doc.keys()
-    # dict_keys(['document', 'summary', 'id'])
-    # (Pdb++) results
-    # [' The Welsh Government has announced
-
-    # breakpoint()
-
-    completion = results[0]
-    # true_refs, false_refs = doc["correct_answers"], doc["incorrect_answers"]
-    # all_refs = true_refs + false_refs
-
-    document = doc["document"]
-    true_refs = [doc["summary"]]
-    all_refs = true_refs
-
-    # ROUGE-N
-    rouge_scores = [rouge([ref], [completion]) for ref in all_refs]
-    # ROUGE-1
-    rouge1_scores = [score["rouge1"] for score in rouge_scores]
-    # ROUGE-2
-    rouge2_scores = [score["rouge2"] for score in rouge_scores]
-    # ROUGE-L
-    rougeL_scores = [score["rougeLsum"] for score in rouge_scores]
-
-    res = {
-        "rouge1": rouge1_scores[0],
-        "rouge2": rouge2_scores[0],
-        "rougeL": rougeL_scores[0],
-    }
-
-    return res
-
-
-def bleu(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(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)
-    # Add newlines between sentences to correctly compute `rougeLsum`.
-
-    def _prepare_summary(summary):
-        summary = summary.replace(" . ", ".\n")
-        return summary
-
-    # Accumulate confidence intervals.
-    aggregator = scoring.BootstrapAggregator()
-    for ref, pred in zip(refs, preds):
-        ref = _prepare_summary(ref)
-        pred = _prepare_summary(pred)
-        aggregator.add_scores(scorer.score(ref, pred))
-    result = aggregator.aggregate()
-    return {type: result[type].mid.fmeasure * 100 for type in rouge_types}