remove numpy, switch to torch (avoid to/from cpu as much)

blimp.py

@@ -18,7 +18,6 @@ from typing import Optional
 
 import datasets
 import evaluate
-import numpy as np
 import torch
 from evaluate import logging
 from transformers import AutoModelForCausalLM, AutoTokenizer
@@ -175,7 +174,8 @@ class Blimp(evaluate.Metric):
                 else ("mps" if torch.mps.is_available() else "cpu")
             )
 
-        if samples_per_set is None or samples_per_set <= 0:
+        samples_per_set = 1000 if samples_per_set is None else samples_per_set
+        if samples_per_set <= 0 or samples_per_set > 1000:
             samples_per_set = 1000
 
         model = AutoModelForCausalLM.from_pretrained(
@@ -226,13 +226,11 @@ class Blimp(evaluate.Metric):
 
             # 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)])
+            good_sents, bad_sents = zip(*sents[:samples_per_set])
 
             # Get probabilities in batches
-            good_probs = get_batch_probabilities(
+            good_probs = _get_batch_probabilities(
                 model,
                 tokenizer,
                 good_sents,
@@ -241,7 +239,7 @@ class Blimp(evaluate.Metric):
                 category,
                 sent_type="good",
             )
-            bad_probs = get_batch_probabilities(
+            bad_probs = _get_batch_probabilities(
                 model,
                 tokenizer,
                 bad_sents,
@@ -251,22 +249,24 @@ class Blimp(evaluate.Metric):
                 sent_type="bad",
             )
 
-            # compute accuracy (mean of instances where good prob > bad prob)
-            accuracy = np.mean(good_probs > bad_probs)
+            # compute accuracy (mean of instances where good prob > bad prob) for this UID
+            sub_acc = (good_probs > bad_probs).float().mean().item()
+
+            phenom = dataset[0]["linguistics_term"]
 
-            results[category] = accuracy
-            phenom_results[phenom].append(accuracy)
+            results[category] = sub_acc
+            phenom_results[phenom].append(sub_acc)
 
         return {
             "by_uid": results,
-            "accuracy": np.mean(list(results.values())),
+            "accuracy": sum(results.values()) / len(results),  # overall accuracy
             "by_phenomenon": {
-                term: np.mean(acc) for term, acc in phenom_results.items()
+                term: sum(acc) / len(acc) for term, acc in phenom_results.items()
             },
         }
 
 
-def get_batch_probabilities(
+def _get_batch_probabilities(
     model,
     tokenizer,
     sentences: list[str],
@@ -276,7 +276,7 @@ def get_batch_probabilities(
     sent_type: str = "good",
 ):
     """Compute log probabilities for a batch of sentences"""
-    probs = []
+    probs = torch.zeros(len(sentences))
 
     for i in logging.tqdm(
         range(0, len(sentences), batch_size),
@@ -307,6 +307,6 @@ def get_batch_probabilities(
         # sum log probabilities
         sequence_log_probs = token_log_probs.sum(dim=1)
 
-        probs.append(sequence_log_probs.cpu().numpy())
+        probs[i : i + batch_size] = sequence_log_probs
 
-    return np.concatenate(probs)
+    return probs

requirements.txt

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