From f7ea12f4908c8bad4a17f23f7174edbd04f1aba1 Mon Sep 17 00:00:00 2001 From: Oleh Prypin Date: Wed, 7 Oct 2026 03:16:20 -0700 Subject: [PATCH] Replace `# pytype: disable` suppressions with `# pyrefly: ignore` PiperOrigin-RevId: 994986656 --- lit_nlp/api/dataset.py | 2 +- lit_nlp/app.py | 7 ++++--- lit_nlp/components/ablation_flip.py | 2 +- lit_nlp/components/classification_results.py | 5 +++-- lit_nlp/components/core.py | 6 ++---- lit_nlp/components/hotflip.py | 2 +- lit_nlp/components/lime_explainer.py | 10 +++++----- .../components/minimal_targeted_counterfactuals.py | 7 ++++--- lit_nlp/components/projection.py | 2 +- lit_nlp/components/regression_results.py | 5 +++-- lit_nlp/components/remote_model.py | 2 +- lit_nlp/components/salience_clustering.py | 5 +++-- lit_nlp/components/shap_explainer.py | 2 +- lit_nlp/examples/glue/models.py | 12 ++++++------ lit_nlp/examples/prompt_debugging/keras_lms.py | 4 +--- lit_nlp/examples/prompt_debugging/models.py | 12 ++++++------ .../examples/prompt_debugging/transformers_lms.py | 8 +++----- lit_nlp/lib/caching.py | 6 +++--- lit_nlp/lib/file_cache.py | 4 ++-- lit_nlp/lib/serialize.py | 4 ++-- 20 files changed, 53 insertions(+), 54 deletions(-) diff --git a/lit_nlp/api/dataset.py b/lit_nlp/api/dataset.py index 9d3b75f7..faba9d87 100644 --- a/lit_nlp/api/dataset.py +++ b/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]: diff --git a/lit_nlp/app.py b/lit_nlp/app.py index 2a667dba..7b6923f0 100644 --- a/lit_nlp/app.py +++ b/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): diff --git a/lit_nlp/components/ablation_flip.py b/lit_nlp/components/ablation_flip.py index 97647efa..8ea7c8c3 100644 --- a/lit_nlp/components/ablation_flip.py +++ b/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. diff --git a/lit_nlp/components/classification_results.py b/lit_nlp/components/classification_results.py index 5761e73b..797b2e80 100644 --- a/lit_nlp/components/classification_results.py +++ b/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() diff --git a/lit_nlp/components/core.py b/lit_nlp/components/core.py index 7415094f..5bbc14cb 100644 --- a/lit_nlp/components/core.py +++ b/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]: diff --git a/lit_nlp/components/hotflip.py b/lit_nlp/components/hotflip.py index 8dc3b086..ea7483a1 100644 --- a/lit_nlp/components/hotflip.py +++ b/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: diff --git a/lit_nlp/components/lime_explainer.py b/lit_nlp/components/lime_explainer.py index 284d55da..a15dd2ab 100644 --- a/lit_nlp/components/lime_explainer.py +++ b/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 ), diff --git a/lit_nlp/components/minimal_targeted_counterfactuals.py b/lit_nlp/components/minimal_targeted_counterfactuals.py index c7f5be35..a9d13628 100644 --- a/lit_nlp/components/minimal_targeted_counterfactuals.py +++ b/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.""" diff --git a/lit_nlp/components/projection.py b/lit_nlp/components/projection.py index 633c3aca..92c70f58 100644 --- a/lit_nlp/components/projection.py +++ b/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 diff --git a/lit_nlp/components/regression_results.py b/lit_nlp/components/regression_results.py index b6beeaf9..57a95d4c 100644 --- a/lit_nlp/components/regression_results.py +++ b/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() diff --git a/lit_nlp/components/remote_model.py b/lit_nlp/components/remote_model.py index 714389a2..f999023e 100644 --- a/lit_nlp/components/remote_model.py +++ b/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] }, diff --git a/lit_nlp/components/salience_clustering.py b/lit_nlp/components/salience_clustering.py index dbb488f8..aa1d429f 100644 --- a/lit_nlp/components/salience_clustering.py +++ b/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) diff --git a/lit_nlp/components/shap_explainer.py b/lit_nlp/components/shap_explainer.py index da667001..031c0518 100644 --- a/lit_nlp/components/shap_explainer.py +++ b/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: diff --git a/lit_nlp/examples/glue/models.py b/lit_nlp/examples/glue/models.py index e075747f..7b147978 100644 --- a/lit_nlp/examples/glue/models.py +++ b/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]): ) # [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 [batch_size, num_tokens, emb_dim] # average over tokens to get [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: diff --git a/lit_nlp/examples/prompt_debugging/keras_lms.py b/lit_nlp/examples/prompt_debugging/keras_lms.py index d384511a..133ca9a3 100644 --- a/lit_nlp/examples/prompt_debugging/keras_lms.py +++ b/lit_nlp/examples/prompt_debugging/keras_lms.py @@ -17,7 +17,6 @@ # 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: @@ -25,11 +24,10 @@ 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 diff --git a/lit_nlp/examples/prompt_debugging/models.py b/lit_nlp/examples/prompt_debugging/models.py index 178b5409..fcfecc06 100644 --- a/lit_nlp/examples/prompt_debugging/models.py +++ b/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, diff --git a/lit_nlp/examples/prompt_debugging/transformers_lms.py b/lit_nlp/examples/prompt_debugging/transformers_lms.py index 4749f751..384e004b 100644 --- a/lit_nlp/examples/prompt_debugging/transformers_lms.py +++ b/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" ) # [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( diff --git a/lit_nlp/lib/caching.py b/lit_nlp/lib/caching.py index ed81bf63..f4b9ac5d 100644 --- a/lit_nlp/lib/caching.py +++ b/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: diff --git a/lit_nlp/lib/file_cache.py b/lit_nlp/lib/file_cache.py index 5205a41d..b3008813 100644 --- a/lit_nlp/lib/file_cache.py +++ b/lit_nlp/lib/file_cache.py @@ -209,11 +209,11 @@ def _get_from_cache( ) with temp_file_manager() as temp_file: logging.info('%s not found in cache.', url) - _fetch_content( # pytype: disable=wrong-arg-types + _fetch_content( url_to_download, temp_file, # pyrefly: ignore[bad-argument-type] headers=headers, - progress_indicator=progress_indicator + progress_indicator=progress_indicator, ) logging.info('%s downloaded to %s', url, temp_file.name) diff --git a/lit_nlp/lib/serialize.py b/lit_nlp/lib/serialize.py index ab51b824..762c040c 100644 --- a/lit_nlp/lib/serialize.py +++ b/lit_nlp/lib/serialize.py @@ -39,7 +39,7 @@ def _obj_to_json(o: object): # .tolist() on a NumPy array. return cast(np.number, o).tolist() # to regular Python scalar elif isinstance(o, types.LitType): - return o.to_json() # pytype: disable=attribute-error # enable-nested-classes + return o.to_json() elif isinstance(o, dtypes.DataTuple): return o.to_json() elif isinstance(o, tuple): @@ -64,7 +64,7 @@ def _obj_to_json_simple(o: object): # .tolist() on a NumPy array. return cast(np.number, o).tolist() # to regular Python scalar elif isinstance(o, types.LitType): - return o.to_json() # pytype: disable=attribute-error # enable-nested-classes + return o.to_json() elif isinstance(o, dtypes.DataTuple): return o.to_json() elif isinstance(o, dtypes.EnumSerializableAsValues):