coh mech

ccl_win.py

@@ -15,7 +15,11 @@
 
 import evaluate
 import datasets
-
+import numpy as np
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+import getpass
+import pdb
+import os
 
 # TODO: Add BibTeX citation
 _CITATION = """\
@@ -28,7 +32,7 @@ year={2020}
 
 # TODO: Add description of the module here
 _DESCRIPTION = """\
-This new module is designed to solve this great ML task and is crafted with a lot of care.
+local coherecence with  classifier trained on the shuffle task, window=3 sentences
 """
 
 
@@ -55,11 +59,12 @@ Examples:
 
 # TODO: Define external resources urls if needed
 BAD_WORDS_URL = "http://url/to/external/resource/bad_words.txt"
-
+WINDOW_SIZE = 3
 
 @evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
 class ccl_win(evaluate.Measurement):
     """TODO: Short description of my evaluation module."""
+    
 
     def _info(self):
         # TODO: Specifies the evaluate.EvaluationModuleInfo object
@@ -86,10 +91,65 @@ class ccl_win(evaluate.Measurement):
         # TODO: Download external resources if needed
         pass
 
-    def _compute(self, predictions, references):
+    def preprocess_adjacent_window(self,preds):
+        pred_list = []
+        lens = []
+        for pred in preds:
+            sents = pred.split("\n")
+            ns = len(sents)
+            if ns <= WINDOW_SIZE:
+                pred_list.append(pred)
+                lens.append(1)
+            else:
+                llen = 0
+                for i in range(0,ns-WINDOW_SIZE+1):
+                    sss = sents[i:i+WINDOW_SIZE]
+                    ss = "\n".join(sss)
+                    pred_list.append(ss)
+                    llen += 1
+                lens.append(llen)
+        #
+        return pred_list,lens
+
+
+
+    def _compute(self, predictions, dataset, device=None):
         """Returns the scores"""
-        # TODO: Compute the different scores of the module
-        accuracy = sum(i == j for i, j in zip(predictions, references)) / len(predictions)
+        MODEL_CACHE_DIR = "/home/rcardena/.cache/huggingface/"
+        if getpass.getuser() == "s1987051":
+            MODEL_CACHE_DIR="/disk/ocean/rcardenas/tools/huggingface/"
+        elif getpass.getuser() == "rcardena":
+            MODEL_CACHE_DIR="/gfs/team/nlp/users/rcardena/tools/huggingface/"
+        
+        if device is not None:
+            assert device in ["gpu", "cpu", "cuda"], "device should be either gpu or cpu."
+            if device == "gpu":
+                device = "cuda"
+        else:
+            device = "cuda" if torch.cuda.is_available() else "cpu"
+
+        tokenizer = AutoTokenizer.from_pretrained("roberta-large",cache_dir=MODEL_CACHE_DIR,use_fast="cnn_dailymail" not in dataset)
+        model = transformers.AutoModelForSequenceClassification.from_pretrained(f"./{dataset}/", num_labels=2,cache_dir=MODEL_CACHE_DIR)
+        model.to(device)
+
+        pred_list,len_by_sample = preprocess_adjacent_window(preds)
+
+        scores = []
+        for text in pred_list:
+            sents = text.lower().split("\n")
+            strides = ["\n".join(sents[i:i+WINDOW_SIZE]) for i in range(0,len(sents),WINDOW_SIZE)]
+            tinput = tokenizer(strides,padding=True,truncation=True,max_length=512,return_tensors="pt")
+            tinput = {k:v.to(device) for k,v in tinput.items()}
+            output = model(**tinput)
+            probs = torch.softmax(output.logits,dim=-1).detach().cpu().numpy()
+            scores.append(probs[:,0].mean())
+        #
+        results = []
+        offset = 0
+        for _len in len_by_sample:
+            results.append( float(np.mean(scores[offset:offset+_len])) )
+            offset += _len
+        #
         return {
-            "accuracy": accuracy,
+            "loc_coh_ccl": results,
         }