Fix typing issues in metrics/ and remove check_untyped_defs. (stanford-crfm#1942)

Author: brianwgoldman

src/helm/benchmark/metrics/summac/model_summac.py

@@ -1,11 +1,12 @@
-# mypy: check_untyped_defs = False
 ###############################################
 # Source: https://github.com/tingofurro/summac
 ###############################################
 
+from typing import Dict, List
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import nltk
 import numpy as np
+import numpy.typing as npt
 import torch
 import os
 import json
@@ -145,6 +146,7 @@ def build_image(self, original, generated):
 
         if self.model is None:
             self.load_nli()
+            assert self.model
 
         dataset = [
             {"premise": original_chunks[i], "hypothesis": generated_chunks[j], "doc_i": i, "gen_i": j}
@@ -303,7 +305,7 @@ def compute_histogram(self, original=None, generated=None, image=None):
 
         full_histogram = []
         for i_gen in range(N_gen):
-            histos = []
+            histos: List[npt.NDArray] = []
 
             for i_depth in range(N_depth):
                 if (
@@ -317,32 +319,31 @@ def compute_histogram(self, original=None, generated=None, image=None):
                     histos.append(histo)
 
             if self.norm_histo:
-                histos = [[N_ori, N_gen]] + histos
+                histos = [np.array([N_ori, N_gen])] + histos
             histogram_row = np.concatenate(histos)
             full_histogram.append(histogram_row)
 
         n_rows_missing = self.n_rows - len(full_histogram)
         full_histogram += [[0.0] * self.full_size] * n_rows_missing
         full_histogram = full_histogram[: self.n_rows]
-        full_histogram = np.array(full_histogram)
-        return image, full_histogram
+        return image, np.array(full_histogram)
 
     def forward(self, originals, generateds, images=None):
         if images is not None:
             # In case they've been pre-computed.
-            histograms = []
+            histogram_list = []
             for image in images:
                 _, histogram = self.compute_histogram(image=image)
-                histograms.append(histogram)
+                histogram_list.append(histogram)
         else:
-            images, histograms = [], []
+            images, histogram_list = [], []
             for original, generated in zip(originals, generateds):
                 image, histogram = self.compute_histogram(original=original, generated=generated)
                 images.append(image)
-                histograms.append(histogram)
+                histogram_list.append(histogram)
 
-        N = len(histograms)
-        histograms = torch.FloatTensor(histograms).to(self.device)
+        N = len(histogram_list)
+        histograms = torch.FloatTensor(histogram_list).to(self.device)
 
         non_zeros = (torch.sum(histograms, dim=-1) != 0.0).long()
         seq_lengths = non_zeros.sum(dim=-1).tolist()
@@ -379,8 +380,8 @@ def forward(self, originals, generateds, images=None):
                     )
             else:
                 features.append(torch.FloatTensor([0.0, 0.0, 0.0]).unsqueeze(0))  # .cuda()
-        features = torch.cat(features)
-        logits = self.layer_final(features)
+        features_tensor = torch.cat(features)
+        logits = self.layer_final(features_tensor)
         histograms_out = [histogram.cpu().numpy() for histogram in histograms]
         return logits, histograms_out, images
 
@@ -451,7 +452,7 @@ def score_one(self, original, generated):
         return {"score": final_score, "image": image}
 
     def score(self, sources, generateds, **kwargs):
-        output = {"scores": [], "images": []}
+        output: Dict[str, List] = {"scores": [], "images": []}
         for source, gen in zip(sources, generateds):
             score = self.score_one(source, gen)
             output["scores"].append(score["score"])

src/helm/benchmark/metrics/test_bias_metrics.py

@@ -1,4 +1,3 @@
-# mypy: check_untyped_defs = False
 from dataclasses import dataclass
 from typing import Callable, List, Optional
 
@@ -12,7 +11,7 @@ class TestCase:
     rel_tol: float = 0.01
 
 
-def check_test_cases(test_cases: List[TestCase], bias_func: Callable[[List[str]], float]):
+def check_test_cases(test_cases: List[TestCase], bias_func: Callable[[List[str]], Optional[float]]):
     for test_case in test_cases:
         bias_score = bias_func(test_case.texts)
         error_msg = f"Expected: {test_case.bias_score}, Actual:{bias_score}"