fix LLM attribution related mypy issues (#1200)

Summary:

as title

Reviewed By: vivekmig

Differential Revision: D50715178

Author: aobo-y

captum/attr/_core/llm_attr.py

@@ -153,7 +153,7 @@ def _format_model_input(self, model_input: Union[str, Tensor]):
         raw text and text token tensors
         """
         # return tensor(1, n_tokens)
-        if type(model_input) is str:
+        if isinstance(model_input, str):
             return self.tokenizer.encode(model_input, return_tensors="pt").to(
                 self.device
             )

captum/attr/_utils/interpretable_input.py

@@ -210,14 +210,24 @@ def __init__(
         if baselines is None:
             # default baseline is to remove the element
             baselines = [""] * len(values)
-        elif dict_keys:
-            assert isinstance(baselines, dict), (
-                "if values is dict, the baselines must also be a dict, "
-                f"received: {type(baselines)}"
-            )
+        elif not callable(baselines):
+            if dict_keys:
+                assert isinstance(baselines, dict), (
+                    "if values is a dict, the baselines must also be a dict "
+                    "or a callable which return a dict, "
+                    f"received: {type(baselines)}"
+                )
 
-            # convert dict to list
-            baselines = [baselines[k] for k in self.dict_keys]
+                # convert dict to list
+                baselines = [baselines[k] for k in dict_keys]
+            else:
+                assert isinstance(baselines, list), (
+                    "if values is a list, the baselines must also be a list "
+                    "or a callable which return a list, "
+                    f"received: {type(baselines)}"
+                )
+
+        self.baselines = baselines
 
         if mask is None:
             n_itp_features = n_features
@@ -247,14 +257,13 @@ def __init__(
         if isinstance(template, str):
             template = template.format
         else:
-            assert isinstance(template, Callable), (
+            assert callable(template), (
                 "the template must be either a string or a callable, "
                 f"received: {type(template)}"
             )
             template = template
         self.format_fn = template
 
-        self.baselines = baselines
         self.mask = mask
 
     def to_tensor(self) -> torch.Tensor:
@@ -265,13 +274,23 @@ def to_model_input(self, perturbed_tensor: Optional[Tensor] = None) -> str:
         values = list(self.values)  # clone
 
         if perturbed_tensor is not None:
-            baselines = self.baselines
-            if isinstance(baselines, Callable):
+            if callable(self.baselines):
                 # a placeholder for advanced baselines
                 # TODO: support callable baselines
                 baselines = self.baselines()
                 if self.dict_keys:
+                    assert isinstance(baselines, dict), (
+                        "if values is a dict and the baselines is a callable"
+                        f"it must return a dict, received: {type(baselines)}"
+                    )
                     baselines = [baselines[k] for k in self.dict_keys]
+                else:
+                    assert isinstance(baselines, list), (
+                        "if values is a list and the baselines is a callable"
+                        f"it must return a list, received: {type(baselines)}"
+                    )
+            else:
+                baselines = self.baselines
 
             for i in range(len(values)):
                 itp_idx = i
@@ -284,8 +303,8 @@ def to_model_input(self, perturbed_tensor: Optional[Tensor] = None) -> str:
                     values[i] = baselines[i]
 
         if self.dict_keys:
-            values = dict(zip(self.dict_keys, values))
-            input_str = self.format_fn(**values)
+            dict_values = dict(zip(self.dict_keys, values))
+            input_str = self.format_fn(**dict_values)
         else:
             input_str = self.format_fn(*values)
 

tests/attr/test_interpretable_input.py

@@ -16,8 +16,8 @@ class TestTextTemplateInput(BaseTest):
             ),
         ]
     )
-    def test_input(self, template, inputs) -> None:
-        tt_input = TextTemplateInput(template, inputs)
+    def test_input(self, template, values) -> None:
+        tt_input = TextTemplateInput(template, values)
 
         expected_tensor = torch.tensor([[1.0] * 4])
         assertTensorAlmostEqual(self, tt_input.to_tensor(), expected_tensor)
@@ -37,11 +37,11 @@ def test_input(self, template, inputs) -> None:
             ),
         ]
     )
-    def test_input_with_baselines(self, template, inputs, baselines) -> None:
+    def test_input_with_baselines(self, template, values, baselines) -> None:
         perturbed_tensor = torch.tensor([[1.0, 0.0, 1.0, 0.0]])
 
         # single instance baselines
-        tt_input = TextTemplateInput(template, inputs, baselines=baselines)
+        tt_input = TextTemplateInput(template, values, baselines=baselines)
         self.assertEqual(tt_input.to_model_input(perturbed_tensor), "a b x d e z")
 
     @parameterized.expand(
@@ -54,8 +54,8 @@ def test_input_with_baselines(self, template, inputs, baselines) -> None:
             ),
         ]
     )
-    def test_input_with_mask(self, template, inputs, mask) -> None:
-        tt_input = TextTemplateInput(template, inputs, mask=mask)
+    def test_input_with_mask(self, template, values, mask) -> None:
+        tt_input = TextTemplateInput(template, values, mask=mask)
 
         expected_tensor = torch.tensor([[1.0] * 2])
         assertTensorAlmostEqual(self, tt_input.to_tensor(), expected_tensor)
@@ -75,8 +75,8 @@ def test_input_with_mask(self, template, inputs, mask) -> None:
             ),
         ]
     )
-    def test_format_attr(self, template, inputs, mask) -> None:
-        tt_input = TextTemplateInput(template, inputs, mask=mask)
+    def test_format_attr(self, template, values, mask) -> None:
+        tt_input = TextTemplateInput(template, values, mask=mask)
 
         attr = torch.tensor([[0.1, 0.2]])