update

src/backend/tasks/cnndm/task.py

@@ -69,6 +69,7 @@ class CnnDm(Task):
         super().__init__(data_dir=data_dir, cache_dir=cache_dir, download_mode=download_mode, config=config)
         self.factkb_tokenizer = None
         self.factkb_model = None
+        self.bert_score = None
 
     def maybe_init_factkb(self):
         if self.factkb_tokenizer is None or self.factkb_model is None:
@@ -76,6 +77,11 @@ class CnnDm(Task):
             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 maybe_init_bertscore(self):
+        if self.bert_score is None:
+            from evaluate import load
+            self.bert_score = load("bertscore")
+
     def has_training_docs(self):
         return True
 
@@ -153,11 +159,17 @@ class CnnDm(Task):
         factkb_logits = self.factkb_model(**factkb_tokens).logits
         factkb_res = torch.softmax(factkb_logits, dim=1)
 
+        self.maybe_init_bertscore()
+        bert_score_res = self.bert_score.compute(predictions=[completion], references=[gold_summary], model_type="microsoft/deberta-xlarge-mnli", lang="en")
+
         res = {
             "rouge1": rouge1_scores[0],
             "rouge2": rouge2_scores[0],
             "rougeL": rougeL_scores[0],
-            "factKB": float(factkb_res[0][1])
+            "factKB": float(factkb_res[0][1]),
+            "bertscore_precision": float(bert_score_res["precision"][0]),
+            "bertscore_recall": float(bert_score_res["recall"][0]),
+            "bertscore_f1": float(bert_score_res["f1"][0])
         }
 
         return res

src/backend/tasks/xsum/task.py

@@ -153,7 +153,7 @@ class XSum(Task):
         factkb_res = torch.softmax(factkb_logits, dim=1)
 
         self.maybe_init_bertscore()
-        bert_score_res = self.bert_score.compute(predictions=[completion], references=[gold_summary], lang="en")
+        bert_score_res = self.bert_score.compute(predictions=[completion], references=[gold_summary], model_type="microsoft/deberta-xlarge-mnli", lang="en")
 
         res = {
             "rouge1": rouge1_scores[0],