update

src/backend/tasks/xsum/task.py

@@ -3,6 +3,7 @@ from lm_eval.api.instance import Instance
 from lm_eval.api.registry import register_task
 from lm_eval.api.metrics import mean
 
+import torch
 import sacrebleu
 from rouge_score import rouge_scorer, scoring
 
@@ -61,11 +62,11 @@ class XSum(Task):
         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 maybe_init_factkb(self):
+        if self.factkb_tokenizer is None or self.factkb_model is None:
+            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, device_map="auto")
 
     def has_training_docs(self):
         return True
@@ -114,7 +115,8 @@ class XSum(Task):
             Instance(
                 request_type="generate_until",
                 doc=doc,
-                arguments=(ctx, {"until": ["\n", "."]}),
+                # arguments=(ctx, {"until": ["\n", "."]}),
+                arguments=(ctx, {"until": ["\n"]}),
                 idx=0,
                 **kwargs
             )
@@ -123,28 +125,34 @@ class XSum(Task):
     def process_results(self, doc, results):
         completion = results[0]
 
-        # document = doc["document"]
+        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]
 
+        self.maybe_init_factkb()
+        input_factkb = [[completion, document]]
+        factkb_tokens = self.factkb_tokenizer(input_factkb, return_tensors="pt", padding="max_length", truncation=True).to(self.factkb_model.device)
+        factkb_logits = self.factkb_model(**factkb_tokens).logits
+        factkb_res = torch.softmax(factkb_logits, dim=1)
+
         res = {
             "rouge1": rouge1_scores[0],
             "rouge2": rouge2_scores[0],
             "rougeL": rougeL_scores[0],
+            "factKB": float(factkb_res[0][1])
         }
 
+        # breakpoint()
+
         return res
 
     def aggregation(self):
@@ -153,7 +161,7 @@ class XSum(Task):
             A dictionary where keys are the names of submetrics and values are
             functions that aggregate a list of metrics
         """
-        return {k: mean for k in ["rouge1", "rouge2", "rougeL"]}
+        return {k: mean for k in ["rouge1", "rouge2", "rougeL", "factKB"]}
 
     def higher_is_better(self):
         """
@@ -161,4 +169,4 @@ class XSum(Task):
             A dictionary where keys are the names of submetrics and values are
             whether a higher value of the submetric is better
         """
-        return {k: True for k in ["rouge1", "rouge2", "rougeL"]}
+        return {k: True for k in ["rouge1", "rouge2", "rougeL", "factKB"]}