minor code cleanup

syntaxgym.py

@@ -197,14 +197,10 @@ class SyntaxGym(evaluate.EvaluationModule):
         surps_shifted = surprisals[:, :-1, :]
         expected_ids = input_ids[:, 1:]
 
-        # TODO: check this logic
-        tt = expected_ids.unsqueeze(2)
         # reindexed surprisals: B * (T - 1)
         surprisals = torch.gather(surps_shifted, 2, expected_ids.unsqueeze(2)) \
             .squeeze(2)
 
-        # surprisals is now B * (T - 1)
-
         #### aggregate
         condition_names = item["conditions"]["condition_name"]
         region_totals = {condition_name: defaultdict(float)

test.py

@@ -14,6 +14,7 @@ def syntaxgym_dataset():
 
 @pytest.fixture(scope="session")
 def syntaxgym_metric():
+    # TODO work out reference
     return evaluate.load("./syntaxgym.py")
 
 
@@ -488,17 +489,14 @@ GPT2_SUBORDINATION_SRC_REFERENCE = \
   ('sub_no-matrix', 5): 4.819862633503057}]
 
 
-def test_gpt_subordination_region_totals():
+def test_gpt_subordination_region_totals(syntaxgym_metric):
     """
     Check region-level surprisals against the original syntaxgym-core
     implementation, using the same underlying `gpt2` model.
     """
-    reference = ...  # TODO
 
-    # TODO work out references
     dataset = datasets.load_dataset("cpllab/syntaxgym", "subordination_src-src")
-    metric = evaluate.load("./syntaxgym.py")
-    result = metric.compute(suite=dataset["test"], model_id="gpt2")
+    result = syntaxgym_metric.compute(suite=dataset["test"], model_id="gpt2")
 
     from pprint import pprint
     pprint(result["region_totals"][0])