coh mech

ccl_win.py

@@ -21,6 +21,8 @@ import getpass
 import pdb
 import os
 import torch
+from rouge_score import scoring
+
 
 # TODO: Add BibTeX citation
 _CITATION = """\
@@ -113,7 +115,7 @@ class ccl_win(evaluate.Measurement):
 
 
 
-    def _compute(self, predictions, dataset, batch_size: int = 16, device=None):
+    def _compute(self, predictions, dataset, batch_size: int = 16, device=None, use_aggregator=True):
         """Returns the scores"""
         MODEL_CACHE_DIR = "/home/rcardena/.cache/huggingface/"
         BASEDIR = "/gfs/team/nlp/users/rcardena/tools/new_evals/ccl_win"
@@ -129,6 +131,12 @@ class ccl_win(evaluate.Measurement):
         else:
             device = "cuda" if torch.cuda.is_available() else "cpu"
 
+        results = []
+        aggregator = None
+        if use_aggregator:
+            np.random.seed(42)
+            aggregator = scoring.BootstrapAggregator()
+
         tokenizer = AutoTokenizer.from_pretrained("roberta-large")
 
         model = AutoModelForSequenceClassification.from_pretrained(os.path.join(BASEDIR,dataset))
@@ -148,12 +156,21 @@ class ccl_win(evaluate.Measurement):
                 probs = torch.softmax(output.logits,dim=-1).detach().cpu().numpy()
                 scores.extend(probs[:,0].tolist())
             #
-        results = []
+
         offset = 0
         for _len in len_by_sample:
-            results.append( float(np.mean(scores[offset:offset+_len])) )
+            score = float(np.mean(scores[offset:offset+_len]))
+            if use_aggregator:
+                aggregator.add_score({"loc_coh_ccl": score})
+            else:
+                results.append(score)
             offset += _len
         #
-        return {
-            "loc_coh_ccl": results,
-        }
+        outres = {}
+        if use_aggregator:
+            res = aggregator.aggregate()
+            for k in res:   outres[k] = res[k].mid
+        else:
+            outres = {"loc_coh_ccl": results}
+
+        return outres