Replace # pytype: disable suppressions with # pyrefly: ignore

lit_nlp/api/dataset.py

@@ -127,7 +127,7 @@ def description(self) -> str:
     Returns:
       (string) A human-readable description for display in the UI.
     """
-    return self._description or inspect.getdoc(self) or ''  # pytype: disable=bad-return-type
+    return self._description or inspect.getdoc(self) or ''
 
   @classmethod
   def init_spec(cls) -> Optional[lit_types.Spec]:

lit_nlp/app.py

@@ -336,7 +336,7 @@ def _annotate_new_data(self,
       example['data'] = dict(annotated_dataset.examples[i], _id=new_id)
       example['id'] = new_id
 
-    return data['inputs']  # pytype: disable=bad-return-type  # always-use-return-annotations
+    return data['inputs']
 
   def _post_new_data(
       self,
@@ -595,11 +595,12 @@ def _get_generated(
 
     dataset = self._datasets[dataset_name]
     # Nested list, containing generated examples from each input.
-    all_generated: list[list[Input]] = genny.run(  # pytype: disable=annotation-type-mismatch  # always-use-return-annotations
+    all_generated: list[list[Input]] = genny.run(
         [ex['data'] for ex in data['inputs']],
         self._models[model],
         dataset,
-        config=config)
+        config=config,
+    )
 
     # Annotate datapoints
     def annotate_generated(datapoints):

lit_nlp/components/ablation_flip.py

@@ -318,7 +318,7 @@ def generate(self,
         continue
 
       # Create counterfactual and obtain model prediction.
-      cf = self._create_cf(example, input_spec, ablation_idxs)  # pytype: disable=wrong-arg-types  # enable-nested-classes
+      cf = self._create_cf(example, input_spec, ablation_idxs)  # pyrefly: ignore[bad-argument-type]
       cf_output = list(model.predict([cf]))[0]
 
       # Check if counterfactual results in a prediction flip.

lit_nlp/components/classification_results.py

@@ -90,13 +90,14 @@ def get_classifications(
 class ClassificationInterpreter(lit_components.Interpreter):
   """Calculates and returns classification results, using thresholds."""
 
-  def run(  # pytype: disable=signature-mismatch  # overriding-parameter-type-checks
+  def run(  # pyrefly: ignore[bad-override]
       self,
       inputs: list[JsonDict],
       model: lit_model.Model,
       dataset: lit_dataset.IndexedDataset,
       model_outputs: Optional[list[JsonDict]] = None,
-      config: Optional[JsonDict] = None):
+      config: Optional[JsonDict] = None,
+  ):
 
     # Find the prediction field key in the model output to use for calculations.
     output_spec = model.output_spec()

lit_nlp/components/core.py

@@ -35,22 +35,20 @@
 from lit_nlp.components import word_replacer
 
 # pylint: disable=g-import-not-at-top
-# pytype: disable=import-error
 try:
-  from lit_nlp.components import shap_explainer
+  from lit_nlp.components import shap_explainer  # pyrefly: ignore[missing-module-attribute]
 
   _SHAP_AVAILABLE = True
 except (ModuleNotFoundError, ImportError):
   _SHAP_AVAILABLE = False
 
 try:
-  from lit_nlp.components import umap
+  from lit_nlp.components import umap  # pyrefly: ignore[missing-module-attribute]
 
   _UMAP_AVAILABLE = True
 except (ModuleNotFoundError, ImportError):
   _UMAP_AVAILABLE = False
 # pylint: enable=g-import-not-at-top
-# pytype: enable=import-error
 
 
 def default_generators() -> dict[str, lit_components.Generator]:

lit_nlp/components/hotflip.py

@@ -322,7 +322,7 @@ def generate(self,
     # nested (and easier to track loop state).
     for token_field, v in tokens_and_gradients.items():
       tokens, grads = v
-      text_field = input_spec[token_field].parent  # pytype: disable=attribute-error
+      text_field = input_spec[token_field].parent  # pyrefly: ignore[missing-attribute]
       logging.info("Identifying Hotflips for input field: %s", str(text_field))
       direction = -1
       if is_regression:

lit_nlp/components/lime_explainer.py

@@ -177,11 +177,11 @@ def run(
     config_defaults = {k: v.default for k, v in self.config_spec().items()}
     config = dict(config_defaults, **(config or {}))  # update and return
 
-    kernel_width = int(config[KERNEL_WIDTH_KEY])  # pyrefly: ignore[bad-argument-type]
-    num_samples = int(config[NUM_SAMPLES_KEY])  # pyrefly: ignore[bad-argument-type]
+    kernel_width = int(config[KERNEL_WIDTH_KEY])
+    num_samples = int(config[NUM_SAMPLES_KEY])
     mask_string = (config[MASK_KEY])
     # pylint: disable=g-explicit-bool-comparison
-    seed = int(config[SEED_KEY]) if config[SEED_KEY] != '' else None  # pyrefly: ignore[bad-argument-type]
+    seed = int(config[SEED_KEY]) if config[SEED_KEY] != '' else None
     # pylint: enable=g-explicit-bool-comparison
 
     # Find keys of input (text) segments to explain.
@@ -201,7 +201,7 @@ def run(
       return None
 
     if (field := config[TARGET_HEAD_KEY]) and (
-        cls_idx := int(config[CLASS_KEY])  # pyrefly: ignore[bad-argument-type]
+        cls_idx := int(config[CLASS_KEY])
     ) != -1:
       # TODO(b/205996131): remove this case
       pred_key = field
@@ -249,7 +249,7 @@ def run(
             class_to_explain=class_to_explain,
             num_samples=num_samples,
             tokenizer=str.split,
-            mask_token=mask_string,  # pyrefly: ignore[bad-argument-type]
+            mask_token=mask_string,
             kernel=functools.partial(
                 lime.exponential_kernel, kernel_width=kernel_width
             ),

lit_nlp/components/minimal_targeted_counterfactuals.py

@@ -178,12 +178,13 @@ def generate(self,
     original_pred = list(model.predict([example]))[0]
 
     # Find dataset examples that are flips.
-    filtered_examples = self._filter_ds_examples(  # pytype: disable=wrong-arg-types  # enable-nested-classes
+    filtered_examples = self._filter_ds_examples(
         dataset=dataset,  # pyrefly: ignore[bad-argument-type]
         model=model,
         reference_output=original_pred,
         pred_key=pred_key,
-        regression_thresh=regression_thresh)
+        regression_thresh=regression_thresh,
+    )
 
     supported_field_names = self._find_all_fields_to_consider(
         ds_spec=dataset.spec(),
@@ -516,7 +517,7 @@ def _calculate_stats(self, dataset: lit_dataset.Dataset) -> dict[str, float]:
     return field_stats
 
   def _calculate_std_dev(self, values: list[float]) -> float:
-    return np.std(values)  # pyrefly: ignore[bad-return]
+    return np.std(values)
 
   def _calculate_categorical_prob(self, values: list[float]) -> float:
     """Returns probability of two values from the list having the same value."""

lit_nlp/components/projection.py

@@ -151,7 +151,7 @@ def _train_instance(
   ) -> ProjectionInterpreter:
     # Ignore pytype warning about abstract methods, since this should always
     # be a subclass of ProjectorModel which has these implemented.
-    projector = self._model_factory(**config.get("proj_kw", {}))  # pytype: disable=not-instantiable
+    projector = self._model_factory(**config.get("proj_kw", {}))
     train_inputs = dataset.examples
 
     # If using input values, then treat inputs as outputs instead of running

lit_nlp/components/regression_results.py

@@ -31,13 +31,14 @@
 class RegressionInterpreter(lit_components.Interpreter):
   """Calculates and returns regression results from model outputs."""
 
-  def run(  # pytype: disable=signature-mismatch  # overriding-parameter-type-checks
+  def run(  # pyrefly: ignore[bad-override]
       self,
       inputs: list[JsonDict],
       model: lit_model.Model,
       dataset: lit_dataset.IndexedDataset,
       model_outputs: Optional[list[JsonDict]] = None,
-      config: Optional[JsonDict] = None):
+      config: Optional[JsonDict] = None,
+  ):
 
     # Find the prediction field key in the model output to use for calculations.
     output_spec = model.output_spec()

lit_nlp/components/remote_model.py

@@ -91,7 +91,7 @@ def predict_minibatch(self, inputs: list[JsonDict]) -> list[JsonDict]:
     preds = query_lit_server(
         self._url,
         'get_preds',
-        params={  # pyrefly: ignore[bad-argument-type]
+        params={
             'model': self._name,
             'response_simple_json': False,  # pyrefly: ignore[bad-assignment]
         },

lit_nlp/components/salience_clustering.py

@@ -153,8 +153,9 @@ def _compute_fixed_length_representation(
         token_weights = {}
 
         for token, score in zip(token_salience.tokens, token_salience.salience):
-          token_weights[token] = max([token_weights.get(token, score), score],  # pyrefly: ignore[no-matching-overload]
-                                     key=abs)
+          token_weights[token] = max(
+              [token_weights.get(token, score), score], key=abs
+          )
 
         representation = self._convert_to_bow_vector(token_weights,
                                                      vocab_lookup, top_k)

lit_nlp/components/shap_explainer.py

@@ -152,7 +152,7 @@ def run(
 
     example_data = inputs or dataset.examples
     examples: pd.DataFrame = pd.DataFrame(example_data)[input_feats]
-    sample_size = int(config.get(SAMPLE_KEY, 0))  # pyrefly: ignore[bad-argument-type]
+    sample_size = int(config.get(SAMPLE_KEY, 0))
     if sample_size and len(examples) > sample_size:
       inputs_to_use: pd.DataFrame = examples.sample(sample_size)
     else:

lit_nlp/examples/glue/models.py

@@ -160,7 +160,7 @@ def _preprocess(self, inputs: Iterable[JsonDict]) -> dict[str, tf.Tensor]:
         tokens_b,
         max_length=self.config.max_seq_length,
     )
-    return encoded_input  # pytype: disable=bad-return-type
+    return encoded_input  # pyrefly: ignore[bad-return]
 
   def _make_dataset(self, inputs: Iterable[JsonDict]) -> tf.data.Dataset:
     """Make a tf.data.Dataset from inputs in LIT format."""
@@ -180,7 +180,7 @@ def _make_dataset(self, inputs: Iterable[JsonDict]) -> tf.data.Dataset:
       )
     # encoded_input is actually a transformers.BatchEncoding
     # object, which tf.data.Dataset doesn't like. Convert to a regular dict.
-    return tf.data.Dataset.from_tensor_slices((dict(encoded_input), labels))  # pyrefly: ignore[bad-argument-type]
+    return tf.data.Dataset.from_tensor_slices((dict(encoded_input), labels))
 
   def train(
       self,
@@ -303,7 +303,7 @@ def _postprocess(self, output: dict[str, Any]):
         # Return the label corresponding to the class index used for gradients.
         output[self.config.label_name] = self.config.labels[  # pyrefly: ignore[unsupported-operation]
             output[self.config.label_name]
-        ]  # pytype: disable=container-type-mismatch
+        ]
 
       # Remove "input_emb_grad" since it's not in the output spec.
       del output["input_emb_grad"]
@@ -321,7 +321,7 @@ def _postprocess(self, output: dict[str, Any]):
       output[key] = output[key][:, :ntok, :ntok].transpose((0, 2, 1))
       # Make a copy of this array to avoid memory leaks, since NumPy otherwise
       # keeps a pointer around that prevents the source array from being GCed.
-      output[key] = output[key].copy()  # pytype: disable=attribute-error
+      output[key] = output[key].copy()
 
     return output
 
@@ -513,7 +513,7 @@ def predict_minibatch(self, inputs: Iterable[JsonDict]):
         )
         # <float32>[batch_size, 1]
         denom = tf.reduce_sum(token_mask, axis=1)
-        for i, layer_output in enumerate(out.hidden_states):  # pyrefly: ignore[bad-argument-type, not-iterable]
+        for i, layer_output in enumerate(out.hidden_states):  # pyrefly: ignore[not-iterable]
           # layer_output is <float32>[batch_size, num_tokens, emb_dim]
           # average over tokens to get <float32>[batch_size, emb_dim]
           batched_outputs[f"layer_{i}/avg_emb"] = (
@@ -529,7 +529,7 @@ def predict_minibatch(self, inputs: Iterable[JsonDict]):
               f"{self.model.config.num_hidden_layers}, got "
               f"{len(out.attentions)}."
           )
-        for i, layer_attention in enumerate(out.attentions):  # pyrefly: ignore[bad-argument-type, not-iterable]
+        for i, layer_attention in enumerate(out.attentions):  # pyrefly: ignore[not-iterable]
           batched_outputs[f"layer_{i+1}/attention"] = layer_attention
 
       if self.is_regression:

lit_nlp/examples/prompt_debugging/keras_lms.py

@@ -17,19 +17,17 @@
 
 
 # pylint: disable=g-import-not-at-top
-# pytype: disable=import-error
 # NOTE: The Keras backend must be set before loading the Keras library. You can
 # set the backend using the KERAS_BACKEND environment variable or your
 # ~/.keras/keras.json configuration file. For more information, see:
 # https://keras.io/getting_started/#configuring-your-backend
 if keras.backend.backend() == "tensorflow":
   import tensorflow as tf
 elif keras.backend.backend() == "torch":
-  import torch
+  import torch  # pyrefly: ignore[missing-import]
 else:
   # TODO(b/333373960): Update imports once a JAX salience is supported.
   raise ValueError(f"Unsupported backend: {keras.backend.backend()}")
-# pytype: enable=import-error
 # pylint: enable=g-import-not-at-top
 
 

lit_nlp/examples/prompt_debugging/models.py

@@ -30,13 +30,13 @@ def _initialize_modeling_environment(
 
     # NOTE: Imported here and not at the top of the file to avoid
     # initialization issues with the environment variables above.
-    import keras  # pylint: disable=g-import-not-at-top # pytype: disable=import-error
+    import keras  # pylint: disable=g-import-not-at-top  # pyrefly: ignore[missing-import]
 
     keras.config.set_floatx(precision)
   elif dl_runtime == "torch":
     # NOTE: Keras sets precision for all backends with set_floatx(), but for
     # HuggingFace Transformers with PyTorch we need to set it explicitly.
-    import torch  # pylint: disable=g-import-not-at-top # pytype: disable=import-error
+    import torch  # pylint: disable=g-import-not-at-top  # pyrefly: ignore[missing-import]
 
     torch.set_default_dtype(
         torch.bfloat16 if precision == "bfloat16" else torch.float32
@@ -85,7 +85,7 @@ def get_models(
     logging.info("Loading model '%s' from '%s'", model_name, path)
 
     if dl_framework == "kerasnlp":
-      from lit_nlp.examples.prompt_debugging import keras_lms  # pylint: disable=g-import-not-at-top # pytype: disable=import-error
+      from lit_nlp.examples.prompt_debugging import keras_lms  # pylint: disable=g-import-not-at-top  # pyrefly: ignore[missing-module-attribute]
 
       models |= keras_lms.initialize_model_group_for_salience(
           model_name,
@@ -94,7 +94,7 @@ def get_models(
           batch_size=batch_size,
       )
     else:
-      from lit_nlp.examples.prompt_debugging import transformers_lms  # pylint: disable=g-import-not-at-top # pytype: disable=import-error
+      from lit_nlp.examples.prompt_debugging import transformers_lms  # pylint: disable=g-import-not-at-top  # pyrefly: ignore[missing-module-attribute]
 
       models |= transformers_lms.initialize_model_group_for_salience(
           model_name,
@@ -140,7 +140,7 @@ def get_model_loaders(
   }
 
   if dl_framework == "kerasnlp":
-    from lit_nlp.examples.prompt_debugging import keras_lms  # pylint: disable=g-import-not-at-top # pytype: disable=import-error
+    from lit_nlp.examples.prompt_debugging import keras_lms  # pylint: disable=g-import-not-at-top  # pyrefly: ignore[missing-module-attribute]
 
     keras_init_spec: lit_types.Spec = {
         **common_init_spec,
@@ -156,7 +156,7 @@ def get_model_loaders(
         )
     }
   else:
-    from lit_nlp.examples.prompt_debugging import transformers_lms  # pylint: disable=g-import-not-at-top # pytype: disable=import-error
+    from lit_nlp.examples.prompt_debugging import transformers_lms  # pylint: disable=g-import-not-at-top  # pyrefly: ignore[missing-module-attribute]
 
     transformers_init_spec: lit_types.Spec = {
         **common_init_spec,

lit_nlp/examples/prompt_debugging/transformers_lms.py

@@ -22,17 +22,15 @@
 import transformers
 
 # pylint: disable=g-import-not-at-top
-# pytype: disable=import-error
 try:
   import tensorflow as tf
 except (ModuleNotFoundError, ImportError):
   logging.warning("TensorFlow is not available.")
 
 try:
-  import torch
+  import torch  # pyrefly: ignore[missing-import]
 except (ModuleNotFoundError, ImportError):
   logging.warning("PyTorch is not available.")
-# pytype: enable=import-error
 # pylint: enable=g-import-not-at-top
 
 
@@ -345,7 +343,7 @@ def _left_pad_target_masks(self, seq_length, target_masks):
         pad_val=0,
         pad_left=self.pad_left,
     )
-    padded_target_masks = np.stack(  # pyrefly: ignore[no-matching-overload]
+    padded_target_masks = np.stack(
         [pad_fn(mask) for mask in modified_masks],
         axis=0,
     )
@@ -395,7 +393,7 @@ def _pred_tf(self, encoded_inputs, target_masks):
           from_logits=True, reduction="none"
       )
       # <tf.float>[batch_size, num_tokens]
-      per_token_loss = loss_fn(target_ids, out.logits)  # pyrefly: ignore[not-callable]
+      per_token_loss = loss_fn(target_ids, out.logits)
       masked_loss = per_token_loss * tf.cast(loss_mask, per_token_loss.dtype)
 
     grads = tape.gradient(

lit_nlp/lib/caching.py

@@ -72,7 +72,7 @@ def load(self) -> dict[CacheKey, Any]:
     except IOError:
       logging.error("Failed loading cache at %s.", self._cache_path)
       data = {}
-    return data  # pytype: disable=name-error  # py310-upgrade
+    return data
 
 
 class PredsCache(object):
@@ -307,7 +307,7 @@ def predict(self,
                    len(miss_idxs), len(cached_results))
     else:
       # If all results were already cached, return them.
-      return cached_results  # pytype: disable=bad-return-type
+      return cached_results
 
     with self._cache.get_pred_lock(input_keys):
       model_preds = list(self.wrapped.predict(progress_indicator(misses)))  # pyrefly: ignore[not-callable]
@@ -326,7 +326,7 @@ def predict(self,
       # Remove the prediction lock from the cache as the request is complete
       self._cache.delete_pred_lock(input_keys)
 
-    return cached_results  # pytype: disable=bad-return-type
+    return cached_results
 
   def _get_results_from_cache(self, input_keys: list[CacheKey]):
     with self._cache.lock: