Fix more pyre errors for llm_attr and tests [2/n]

Summary: Get rid of some more pyre errors by adding better typing

Differential Revision: D63365945

Author: csauper

captum/attr/_core/llm_attr.py

@@ -343,7 +343,7 @@ def _forward_func(
             caching=use_cached_outputs,
         )
 
-        log_prob_list = []
+        log_prob_list: List[Tensor] = []
         outputs = None
         for target_token in target_tokens:
             if use_cached_outputs:
@@ -382,17 +382,15 @@ def _forward_func(
                 (model_inp, torch.tensor([[target_token]]).to(self.device)), dim=1
             )
 
-        # pyre-ignore[9] pyre/mypy thinks sum returns int here, but it will return
-        # Tensor
         total_log_prob: Tensor = sum(log_prob_list)  # type: ignore
         # 1st element is the total prob, rest are the target tokens
         # add a leading dim for batch even we only support single instance for now
         if self.include_per_token_attr:
             target_log_probs = torch.stack(
-                [total_log_prob, *log_prob_list], dim=0  # type: ignore
+                [total_log_prob, *log_prob_list], dim=0
             ).unsqueeze(0)
         else:
-            target_log_probs = total_log_prob  # type: ignore
+            target_log_probs = total_log_prob
         target_probs = torch.exp(target_log_probs)
 
         if _inspect_forward:
@@ -544,8 +542,7 @@ def attribute(
             _convert_ids_to_pretty_tokens(target_tokens, self.tokenizer),
         )
 
-    # pyre-fixme[24] Generic type `Callable` expects 2 type parameters.
-    def attribute_future(self) -> Callable:
+    def attribute_future(self) -> Callable[[], LLMAttributionResult]:
         r"""
         This method is not implemented for LLMAttribution.
         """
@@ -745,8 +742,7 @@ def attribute(
             _convert_ids_to_pretty_tokens(target_tokens, self.tokenizer),
         )
 
-    # pyre-fixme[24] Generic type `Callable` expects 2 type parameters.
-    def attribute_future(self) -> Callable:
+    def attribute_future(self) -> Callable[[], LLMAttributionResult]:
         r"""
         This method is not implemented for LLMGradientAttribution.
         """

tests/attr/test_llm_attr.py

@@ -8,6 +8,7 @@
     cast,
     Dict,
     List,
+    Literal,
     NamedTuple,
     Optional,
     overload,
@@ -18,7 +19,6 @@
 
 import torch
 from captum._utils.models.linear_model import SkLearnLasso
-from captum._utils.typing import Literal
 from captum.attr._core.feature_ablation import FeatureAblation
 from captum.attr._core.kernel_shap import KernelShap
 from captum.attr._core.layer.layer_gradient_shap import LayerGradientShap
@@ -44,9 +44,6 @@ class DummyTokenizer:
     @overload
     def encode(self, text: str, return_tensors: None = None) -> List[int]: ...
     @overload
-    # pyre-fixme[43]: Incompatible overload. The implementation of
-    # `DummyTokenizer.encode` does not accept all possible arguments of overload.
-    # pyre-ignore[11]: Annotation `pt` is not defined as a type
     def encode(self, text: str, return_tensors: Literal["pt"]) -> Tensor: ...
 
     def encode(

tests/attr/test_llm_attr_gpu.py

@@ -3,11 +3,21 @@
 # pyre-strict
 
 import copy
-from typing import Any, cast, Dict, List, NamedTuple, Optional, overload, Type, Union
+from typing import (
+    Any,
+    cast,
+    Dict,
+    List,
+    Literal,
+    NamedTuple,
+    Optional,
+    overload,
+    Type,
+    Union,
+)
 
 import torch
 
-from captum._utils.typing import Literal
 from captum.attr._core.feature_ablation import FeatureAblation
 from captum.attr._core.kernel_shap import KernelShap
 from captum.attr._core.layer.layer_integrated_gradients import LayerIntegratedGradients
@@ -32,9 +42,6 @@ class DummyTokenizer:
     @overload
     def encode(self, text: str, return_tensors: None = None) -> List[int]: ...
     @overload
-    # pyre-fixme[43]: Incompatible overload. The implementation of
-    # `DummyTokenizer.encode` does not accept all possible arguments of overload.
-    # pyre-ignore[11]: Annotation `pt` is not defined as a type
     def encode(self, text: str, return_tensors: Literal["pt"]) -> Tensor: ...
 
     def encode(