Merge branch 'main' of https://huggingface.co/spaces/ronaldahmed/ccl_win into main

ccl_win.py

@@ -15,6 +15,13 @@
 
 import evaluate
 import datasets
+import numpy as np
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+import getpass
+import pdb
+import os
+import torch
+from rouge_score import scoring
 
 
 # TODO: Add BibTeX citation
@@ -28,7 +35,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 +62,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
@@ -71,8 +79,7 @@ class ccl_win(evaluate.Measurement):
             inputs_description=_KWARGS_DESCRIPTION,
             # This defines the format of each prediction and reference
             features=datasets.Features({
-                'predictions': datasets.Value('int64'),
-                'references': datasets.Value('int64'),
+                'predictions': datasets.Value('string'),
             }),
             # Homepage of the module for documentation
             homepage="http://module.homepage",
@@ -86,10 +93,85 @@ 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="arxiv", batch_size: int = 16, device=None, use_aggregator=True):
         """Returns the scores"""
-        # TODO: Compute the different scores of the module
-        accuracy = sum(i == j for i, j in zip(predictions, references)) / len(predictions)
-        return {
-            "accuracy": accuracy,
-        }
+        MODEL_CACHE_DIR = "/home/rcardena/.cache/huggingface/"
+        BASEDIR = "/gfs/team/nlp/users/rcardena/tools/new_evals/ccl_win"
+        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"
+
+        results = []
+        sent_lens = [len(x.split("\n")) for x in predictions]
+        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))
+        model.to(device)
+        model.eval()
+
+        pred_list,len_by_sample = self.preprocess_adjacent_window(predictions)
+
+        scores = []
+        n_preds = len(pred_list)
+        with torch.no_grad():
+            for b in range(0,n_preds,batch_size):
+                strides = [x.lower() for x in pred_list[b:b+batch_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.extend(probs[:,0].tolist())
+            #
+
+        offset = 0
+        for i,_len in enumerate(len_by_sample):
+            score = float(np.mean(scores[offset:offset+_len])) if sent_lens[i]>1 else 0.
+            if use_aggregator:
+                aggregator.add_scores({"loc_coh_ccl": score})
+            else:
+                results.append(score)
+            offset += _len
+        #
+        outres = {}
+        if use_aggregator:
+            res = aggregator.aggregate()
+            for k in res:   outres[k] = res[k].mid
+        else:
+            outres = {"loc_coh_ccl": results}
+
+        return outres