use numpy code for simplicity

blimp.py

@@ -18,6 +18,7 @@ from typing import Optional
 
 import datasets
 import evaluate
+import numpy as np
 import torch
 from evaluate import logging
 from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -250,24 +251,18 @@ class Blimp(evaluate.Metric):
                 sent_type="bad",
             )
 
-            # Compare probabilities
-            correct = sum(g > b for g, b in zip(good_probs, bad_probs))
-            accuracy = correct / len(good_probs)
-            results[category] = accuracy
+            # compute accuracy (mean of instances where good prob > bad prob)
+            accuracy = np.mean(good_probs > bad_probs)
 
+            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 {
             "by_uid": results,
-            "accuracy": overall_accuracy,
-            "by_phenomenon": phenom_term_averages,
+            "accuracy": np.mean(list(results.values())),
+            "by_phenomenon": {
+                term: np.mean(acc) for term, acc in phenom_results.items()
+            },
         }
 
 
@@ -307,12 +302,11 @@ def get_batch_probabilities(
 
         if batch_size > 1:
             # mask padding tokens
-            mask = (labels != tokenizer.pad_token_id).float()
-            token_log_probs *= mask
+            token_log_probs.masked_fill_(labels == tokenizer.pad_token_id, 0.0)
 
         # sum log probabilities
         sequence_log_probs = token_log_probs.sum(dim=1)
 
-        probs.extend(sequence_log_probs.cpu().tolist())
+        probs.append(sequence_log_probs.cpu().numpy())
 
-    return probs
+    return np.concatenate(probs)

requirements.txt

@@ -1,4 +1,4 @@
 git+https://github.com/huggingface/evaluate@5aa3982a9a8c86e506860e381d428a64b0cce73b
 torch
-torch
-transformers
+transformers
+numpy