add by phenomenon

blimp.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 """Blimp Metric."""
 
+from collections import defaultdict
+
 import datasets
 import evaluate
 import torch
@@ -21,7 +23,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 
 datasets.logging.set_verbosity_error()
 
-BLIMP_PHENOMENA = [
+BLIMP_UIDS = [
     "adjunct_island",
     "anaphor_gender_agreement",
     "anaphor_number_agreement",
@@ -191,30 +193,37 @@ class Blimp(evaluate.Metric):
             # assign one of the special tokens to also be the pad token
             tokenizer.add_special_tokens({"pad_token": existing_special_tokens[0]})
 
-        print("PAD", tokenizer.pad_token_id)
-
         run_all = len(predictions) == 0 or predictions[0] == "*"
         blimp_sets = (
-            BLIMP_PHENOMENA
+            BLIMP_UIDS
             if run_all
-            else [p for p in BLIMP_PHENOMENA if p.lower() in predictions]
+            else [p for p in BLIMP_UIDS if p.lower() in predictions]
         )
 
         assert len(blimp_sets) > 0, "no valid phenomena selected"
 
         results = {}
+        phenom_results = defaultdict(list)
 
-        for phenomenon in logging.tqdm(blimp_sets, desc="Evaluating phenomena..."):
-            dataset = datasets.load_dataset("nyu-mll/blimp", phenomenon)["train"]
+        for category in logging.tqdm(blimp_sets, desc="Evaluating phenomena..."):
+            dataset = datasets.load_dataset("nyu-mll/blimp", category)["train"]
 
             # Prepare batches of good and bad sentences
 
+            phenom = dataset[0]["linguistics_term"]
+
             sents = [(x["sentence_good"], x["sentence_bad"]) for x in dataset]
             good_sents, bad_sents = zip(*sents[: min(1000, samples_per_set)])
 
             # Get probabilities in batches
             good_probs = get_batch_probabilities(
-                model, tokenizer, good_sents, device, batch_size, phenomenon
+                model,
+                tokenizer,
+                good_sents,
+                device,
+                batch_size,
+                category,
+                sent_type="good",
             )
             bad_probs = get_batch_probabilities(
                 model,
@@ -222,19 +231,29 @@ class Blimp(evaluate.Metric):
                 bad_sents,
                 device,
                 batch_size,
-                phenomenon,
+                category,
                 sent_type="bad",
             )
 
             # Compare probabilities
             correct = sum(g > b for g, b in zip(good_probs, bad_probs))
             accuracy = correct / len(good_probs)
-            results[phenomenon] = accuracy
+            results[category] = accuracy
+
+            phenom_results[phenom].append(accuracy)
 
+        phenom_term_averages = {
+            term: sum(accuracies) / len(accuracies)
+            for term, accuracies in phenom_results.items()
+        }
         # Calculate overall accuracy
         overall_accuracy = sum(results.values()) / len(results)
 
-        return {"phenomenon_accuracies": results, "overall_accuracy": overall_accuracy}
+        return {
+            "by_uid": results,
+            "accuracy": overall_accuracy,
+            "by_phenomenon": phenom_term_averages,
+        }
 
 
 def get_batch_probabilities(