Replace # pytype: disable suppressions with # pyrefly: ignore

seqio/dataset_providers.py

@@ -185,7 +185,7 @@ def add(cls, name: str, provider_cls, provider_kwargs):
           "Expecting instance of %s, got %s"
           % (cls._PROVIDER_TYPE, provider_cls)
       )
-    provider = provider_cls(**provider_kwargs)  # pytype: disable=wrong-arg-types  # dynamic-method-lookup
+    provider = provider_cls(**provider_kwargs)
     cls.add_provider(name, provider)
     return provider
 
@@ -335,7 +335,7 @@ def list_shards(self, split: str) -> Sequence[str]:
 
   @abc.abstractmethod
   def get_dataset(  # pyrefly: ignore[bad-override]
-      self,  # pytype: disable=signature-mismatch  # overriding-default-value-checks
+      self,
       split: str = tfds.Split.TRAIN,  # pyrefly: ignore[missing-attribute]
       shuffle: bool = True,
       seed: Optional[int] = None,
@@ -358,7 +358,7 @@ def get_dataset(  # pyrefly: ignore[bad-override]
     """
     raise NotImplementedError
 
-  def num_input_examples(self, split: str) -> Optional[int]:  # pytype: disable=signature-mismatch  # overriding-return-type-checks
+  def num_input_examples(self, split: str) -> Optional[int]:
     if self._num_input_examples is None:
       return None
     return self._num_input_examples[split]
@@ -1240,7 +1240,7 @@ def _all_metric_fns(
   def metric_fns(self) -> Sequence[MetricFnCallable]:
     """List of all metric functions."""
     predict_fns, score_fns, predict_with_aux_fns = self._all_metric_fns
-    return predict_fns + score_fns + predict_with_aux_fns  # pytype: disable=unsupported-operands
+    return predict_fns + score_fns + predict_with_aux_fns
 
   @property
   def predict_metric_fns(self) -> Sequence[MetricFnCallable]:
@@ -1356,7 +1356,7 @@ def replace(self, **kwargs):
         task_kwargs[key] = getattr(self, key)
     return Task(**task_kwargs)
 
-  def num_input_examples(self, split: str) -> Optional[int]:  # pytype: disable=signature-mismatch  # overriding-return-type-checks
+  def num_input_examples(self, split: str) -> Optional[int]:
     return self.source.num_input_examples(split)
 
   def _preprocess_dataset(
@@ -1545,7 +1545,7 @@ def get_cached_stats(
     return self._stats[split]
 
   def get_dataset(
-      self,  # pytype: disable=signature-mismatch  # overriding-default-value-checks
+      self,
       sequence_length: Optional[Mapping[str, int]] = None,
       split: str = tfds.Split.TRAIN,  # pyrefly: ignore[missing-attribute]
       use_cached: bool = False,
@@ -1982,7 +1982,7 @@ def filter_features(ex):
   def _get_all_mixing_rates(self, tasks):
     return [self.get_rate(task) for task in tasks]
 
-  def get_dataset(  # pytype: disable=signature-mismatch  # overriding-parameter-type-checks
+  def get_dataset(
       self,
       sequence_length: Optional[Mapping[str, int]] = None,
       split: str = tfds.Split.TRAIN,  # pyrefly: ignore[missing-attribute]

seqio/dataset_providers_test.py

@@ -245,7 +245,7 @@ def _get_preps_with_cache_placeholder_buffer_size(self, buffer_size):
   def _mock_and_assert_cached_source(self, task_name):
     cached_task = dataset_providers.get_mixture_or_task(task_name)
     cached_task._get_cached_source = mock.MagicMock(
-        side_effect=cached_task._get_cached_source  # pytype: disable=attribute-error  # always-use-return-annotations
+        side_effect=cached_task._get_cached_source
     )
     _ = cached_task.get_dataset(None, "train", use_cached=True)
     cached_task._get_cached_source.assert_called_once_with("train")
@@ -961,7 +961,7 @@ def test_list_shards(self):
     def _get_formatted_shards_list(task_name, split):
       shards = dataset_providers.get_mixture_or_task(
           task_name
-      ).source.list_shards(  # pytype: disable=attribute-error
+      ).source.list_shards(
           split
       )  # always-use-return-annotations
       shards = [s.split("/")[-1] for s in shards]

seqio/evaluation.py

@@ -204,7 +204,7 @@ def _extract_model_output(cached_model_dataset, model_fn):
   def _permute(x, sorted_order):
     return [x[sorted_order[i]] for i in range(len(sorted_order))]
 
-  model_fn_result = model_fn(cached_model_dataset)  # pytype: disable=missing-parameter  # always-use-return-annotations
+  model_fn_result = model_fn(cached_model_dataset)
   if isinstance(model_fn_result, tuple):
     # Some of model functions return a tuple of two outputs per example.
     # e.g., ModelOutputType.PREDICTION_WITH_AUX,
@@ -528,7 +528,7 @@ def dataset_fn(task: Task) -> tf.data.Dataset:
       logging.warning(
           "'logger_cls' is empty so seqio.Evaluator will not log its results."
       )
-    self._loggers = tuple(cls(output_dir=log_dir) for cls in logger_cls)  # pytype:disable=not-instantiable
+    self._loggers = tuple(cls(output_dir=log_dir) for cls in logger_cls)
 
   def __del__(self):
     """Wait for metrics to be written before deletion."""
@@ -768,7 +768,7 @@ def _compute_clu_metrics(
 
     result: TaskMetricsType = {}
     # TODO(b/309107492): Fix that attribute-error instead of silencing it.
-    for k, v in itertools.chain(*[m.items() for m in task_metrics]):  # pytype: disable=attribute-error
+    for k, v in itertools.chain(*[m.items() for m in task_metrics]):
       if k in result:
         raise ValueError(f"Duplicate metric key '{k}' in Task '{task.name}'.")
       result[k] = v

seqio/evaluation_test.py

@@ -147,7 +147,7 @@ def _task_from_tensor_slices(name, tensor_slices, label_classes):
               tensor_slices
           ),
           splits="validation",
-      ),  # pytype: disable=wrong-arg-types
+      ),
       preprocessors=[
           utils.map_over_dataset(
               lambda ex: {
@@ -338,7 +338,7 @@ def _task_from_tensor_slices_rank2(name, tensor_slices, label_classes):
                   tensor_slices
               ),
               splits="validation",
-          ),  # pytype: disable=wrong-arg-types
+          ),
           preprocessors=[
               utils.map_over_dataset(
                   lambda ex: {
@@ -446,7 +446,7 @@ def mock_init(self):
       self._target_field_name = target_field_name
 
     with mock.patch.object(Evaluator, "__init__", new=mock_init):
-      evaluator = Evaluator()  # pytype: disable=missing-parameter
+      evaluator = Evaluator()  # pyrefly: ignore[missing-argument]
 
       # A dummy score function that always returns the same output.
       def predict_fn(
@@ -575,7 +575,7 @@ def mock_init(self):
       self._target_field_name = "targets"
 
     with mock.patch.object(Evaluator, "__init__", new=mock_init):
-      evaluator = Evaluator()  # pytype: disable=missing-parameter
+      evaluator = Evaluator()  # pyrefly: ignore[missing-argument]
 
       def predict_with_aux_fn(
           ds: tf.data.Dataset,
@@ -720,7 +720,7 @@ def mock_init(self):
       self._target_field_name = "targets"
 
     with mock.patch.object(Evaluator, "__init__", new=mock_init):
-      evaluator = Evaluator()  # pytype: disable=missing-parameter
+      evaluator = Evaluator()  # pyrefly: ignore[missing-argument]
 
       # A dummy prediction function that always returns the same output.
       # The first example is correct but the second is not.
@@ -765,7 +765,7 @@ def mock_init(self):
       self._target_field_name = "targets"
 
     with mock.patch.object(Evaluator, "__init__", new=mock_init):
-      evaluator = Evaluator()  # pytype: disable=missing-parameter
+      evaluator = Evaluator()  # pyrefly: ignore[missing-argument]
 
       # The output tokens will be docoded to ["e5", "e6", "e7"] and
       # postprocessed to [0, 1, 2].
@@ -844,7 +844,7 @@ def score_fn(
         self.assertEqual(ds, mock_ds1)
         return [(0, 1), (1, 2)]
 
-      evaluator = Evaluator()  # pytype: disable=missing-parameter
+      evaluator = Evaluator()  # pyrefly: ignore[missing-argument]
       all_metrics, _ = evaluator.evaluate(
           compute_metrics=True, predict_fn=predict_fn, score_fn=score_fn  # pyrefly: ignore[bad-argument-type]
       )
@@ -1112,7 +1112,7 @@ def mock_init(self):
       self._target_field_name = "targets"
 
     with mock.patch.object(Evaluator, "__init__", new=mock_init):
-      evaluator = Evaluator()  # pytype: disable=missing-parameter
+      evaluator = Evaluator()  # pyrefly: ignore[missing-argument]
       predict_fn = mock.Mock(return_value=[(0, 1)])
       evaluator.evaluate(
           compute_metrics=False,
@@ -1147,7 +1147,7 @@ def mock_init(self):
       self._target_field_name = "targets"
 
     with mock.patch.object(Evaluator, "__init__", new=mock_init):
-      evaluator = Evaluator()  # pytype: disable=missing-parameter
+      evaluator = Evaluator()  # pyrefly: ignore[missing-argument]
 
       # Dummy predict_fn where only the order is mixed.
       def mixing_order_predict_fn(
@@ -1280,7 +1280,7 @@ def mock_init(self):
       self._target_field_name = "targets"
 
     with mock.patch.object(Evaluator, "__init__", new=mock_init):
-      evaluator = Evaluator()  # pytype: disable=missing-parameter
+      evaluator = Evaluator()  # pyrefly: ignore[missing-argument]
 
       all_metrics, _ = evaluator.evaluate(
           compute_metrics=True,

seqio/experimental.py

@@ -390,7 +390,7 @@ def _get_maybe_sharded_dataset(
     eval_ds = _apply_preprocessors(eval_ds, self._eval_preprocessors)
     datasets['eval'] = eval_ds
 
-    return tf.data.Dataset.zip(datasets)  # pyrefly: ignore[bad-argument-type]
+    return tf.data.Dataset.zip(datasets)
 
 
 

seqio/feature_converters_test.py

@@ -278,7 +278,7 @@ def test_pass_through_and_packing(self):
                     dtype=tf.int32
                 ),
             },
-        )  # pytype: disable=not-instantiable
+        )
 
   def test_pass_through(self):
     with mock.patch.object(
@@ -291,7 +291,7 @@ def test_pass_through(self):
                   dtype=tf.int32
               ),
           },
-      )  # pytype: disable=not-instantiable
+      )
       self.assertDictEqual(
           converter._passthrough_features,
           {
@@ -309,7 +309,7 @@ def test_validate_dataset_missing_feature(self):
     with mock.patch.object(
         feature_converters.FeatureConverter, "__abstractmethods__", set()
     ):
-      converter = feature_converters.FeatureConverter()  # pytype: disable=not-instantiable
+      converter = feature_converters.FeatureConverter()  # pyrefly: ignore[bad-instantiation]
       expected_msg = (
           "Dataset is missing an expected feature during "
           "initial validation: 'inputs'"
@@ -342,7 +342,7 @@ def test_validate_dataset_incorrect_dtype(self):
       feature_converters.FeatureConverter.TASK_FEATURES = {
           k: FeatureSpec(v) for k, v in task_feature_dtypes.items()
       }
-      converter = feature_converters.FeatureConverter()  # pytype: disable=not-instantiable
+      converter = feature_converters.FeatureConverter()  # pyrefly: ignore[bad-instantiation]
       expected_msg = (
           "Dataset has incorrect type for feature 'inputs' during "
           "initial validation: Got int32, expected int64"
@@ -371,7 +371,7 @@ def test_validate_dataset_incorrect_rank(self):
     with mock.patch.object(
         feature_converters.FeatureConverter, "__abstractmethods__", set()
     ):
-      converter = feature_converters.FeatureConverter()  # pytype: disable=not-instantiable
+      converter = feature_converters.FeatureConverter()  # pyrefly: ignore[bad-instantiation]
       expected_msg = (
           "Dataset has incorrect rank for feature 'inputs' during "
           "initial validation: Got 2, expected 1"
@@ -415,7 +415,7 @@ def test_validate_dataset_rank_0(self):
           "inputs": FeatureSpec(dtype=tf.int64, rank=0, sequence_dim=0),
           "targets": FeatureSpec(dtype=tf.int64),
       }
-      converter = feature_converters.FeatureConverter()  # pytype: disable=not-instantiable
+      converter = feature_converters.FeatureConverter()  # pyrefly: ignore[bad-instantiation]
       converter._pack = False
       converter(ds, task_feature_lengths)
 
@@ -431,7 +431,7 @@ def test_validate_dataset_rank_2(self):
     with mock.patch.object(
         feature_converters.FeatureConverter, "__abstractmethods__", set()
     ):
-      converter = feature_converters.FeatureConverter()  # pytype: disable=not-instantiable
+      converter = feature_converters.FeatureConverter()  # pyrefly: ignore[bad-instantiation]
       converter._validate_dataset(
           ds,
           expected_features={
@@ -459,7 +459,7 @@ def test_validate_dataset_rank_2_with_pack(self):
         "_convert_features",
         return_value=ds,
     ):
-      converter = feature_converters.FeatureConverter(pack=True)  # pytype: disable=not-instantiable
+      converter = feature_converters.FeatureConverter(pack=True)  # pyrefly: ignore[bad-instantiation]
       feature_converters.FeatureConverter.TASK_FEATURES = {
           "inputs": FeatureSpec(tf.int64, rank=2),
           "targets": FeatureSpec(tf.int64),
@@ -488,7 +488,7 @@ def test_call_missing_input_lengths(self):
     with mock.patch.object(
         feature_converters.FeatureConverter, "__abstractmethods__", set()
     ):
-      converter = feature_converters.FeatureConverter()  # pytype: disable=not-instantiable
+      converter = feature_converters.FeatureConverter()  # pyrefly: ignore[bad-instantiation]
       feature_converters.FeatureConverter.TASK_FEATURES = {
           "inputs": FeatureSpec(tf.int64),
           "targets": FeatureSpec(tf.int64),
@@ -512,7 +512,7 @@ def test_validate_dataset_pretokenized_field(self):
     with mock.patch.object(
         feature_converters.FeatureConverter, "__abstractmethods__", set()
     ):
-      converter = feature_converters.FeatureConverter()  # pytype: disable=not-instantiable
+      converter = feature_converters.FeatureConverter()  # pyrefly: ignore[bad-instantiation]
       # _validate_dataset works even if ds has targets and targets_pretokenized
       ds = converter._validate_dataset(
           ds,

seqio/helpers.py

@@ -354,10 +354,10 @@ def mixture_with_missing_task_splits_removed(
   Returns:
     The new `Mixture` object.
   """
-  og_mix: dp.Mixture = dp.get_mixture_or_task(mixture_name)  # pytype: disable=annotation-type-mismatch  # always-use-return-annotations
+  og_mix: dp.Mixture = dp.get_mixture_or_task(mixture_name)
   new_tasks_and_rates = []
   for task_name, rate in og_mix._task_to_rate.items():
-    subtask: dp.Task = dp.get_mixture_or_task(task_name)  # pytype: disable=annotation-type-mismatch  # always-use-return-annotations
+    subtask: dp.Task = dp.get_mixture_or_task(task_name)
     if split in subtask.splits:
       new_tasks_and_rates.append((subtask.name, rate))
   new_mix = dp.Mixture(

seqio/loggers.py

@@ -254,7 +254,7 @@ def _get_summary_writer(self, task_name: str) -> tf.summary.SummaryWriter:
         )
     return self._summary_writers[task_name]
 
-  def __call__(  # pytype: disable=signature-mismatch  # overriding-parameter-type-checks
+  def __call__(
       self,
       task_name: str,
       step: Optional[int],

seqio/loggers_test.py

@@ -108,7 +108,7 @@ def _log_and_read(self, metrics):
   def test_log_scalar(self):
     task_metrics = {
         "rouge1": metrics_lib.Scalar(50),
-        "rouge2": metrics_lib.Scalar(np.float32(100)),  # pytype: disable=wrong-arg-types  # numpy-scalars
+        "rouge2": metrics_lib.Scalar(np.float32(100)),  # pyrefly: ignore[bad-argument-type]
     }
     logged_metrics, plugins = self._log_and_read(task_metrics)
     self.assertDictEqual(
@@ -421,7 +421,7 @@ def test_numpy_data(self):
     logger(
         task_name="test",
         step=42,
-        metrics={"accuracy": metrics_lib.Scalar(np.float32(100))},  # pytype: disable=wrong-arg-types  # numpy-scalars
+        metrics={"accuracy": metrics_lib.Scalar(np.float32(100))},  # pyrefly: ignore[bad-argument-type]
         dataset=task_dataset,
         inferences=inferences,
         targets=targets,

seqio/metrics.py

@@ -183,7 +183,7 @@ def from_model_output(  # pylint:disable=missing-function-docstring
           ])
       )
     return cls(
-        values={  # pyrefly: ignore[bad-argument-type]
+        values={
             "model_output": model_output,  # pyrefly: ignore[bad-assignment]
             "indices_2d": indices_2d,  # pyrefly: ignore[bad-assignment]
             "mask": mask,  # pyrefly: ignore[bad-assignment]

seqio/test_utils.py

@@ -959,7 +959,7 @@ def __init__(self, task_name, per_split_data):
     self._task = dataset_providers.get_mixture_or_task(task_name)
 
     self.per_split_data = per_split_data
-    self._saved_source = self._task._source  # pytype: disable=attribute-error  # always-use-return-annotations
+    self._saved_source = self._task._source
 
   def __enter__(self):
     def ds_fn(split, shuffle_files, seed=None):

seqio/utils.py

@@ -425,13 +425,11 @@ def load_shard(
       seed: Optional[int] = None,
   ):
     """Returns a dataset for a single shard of the TFDS TFRecord files."""
-    # pytype:disable=attribute-error
     ds = self.builder._tfrecords_reader.read_files(  # pylint:disable=protected-access
         [file_instruction],
         read_config=tfds.ReadConfig(shuffle_seed=seed),
         shuffle_files=shuffle_files,
     )
-    # pytype:enable=attribute-error
     return ds
 
   def size(self, split: str) -> Optional[int]: