From 016a26933cffc4117ac71edc18ffab13146ad117 Mon Sep 17 00:00:00 2001
From: Oleh Prypin <oprypin@google.com>
Date: Wed, 7 Oct 2026 03:23:07 -0700
Subject: [PATCH] Replace `# pytype: disable` suppressions with `# pyrefly:
 ignore`

PiperOrigin-RevId: 994989816
---
 .../core/community/dataset_sources.py         |  2 +-
 .../core/community/register_package.py        |  4 +-
 .../conll/conll_dataset_builder.py            |  2 +-
 .../conll/conllu_dataset_builder.py           |  2 +-
 .../rlds/datasets/locomotion/locomotion.py    | 10 +--
 .../robosuite_panda_pick_place_can.py         | 45 +++++-----
 .../robomimic/dataset_utils.py                |  6 +-
 .../deployment/export_community_datasets.py   |  2 +-
 .../documentation/document_datasets.py        | 14 ++-
 .../cherry_blossoms/cherry_blossoms.py        |  2 +-
 .../structured/howell/howell.py               |  2 +-
 .../structured/wine_quality/wine_quality.py   |  2 +-
 .../testing/huggingface/dummy_dataset.py      |  2 +-
 .../text/c4_wsrs/c4_wsrs_utils.py             |  2 +-
 .../dolphin_number_word.py                    |  2 +-
 tensorflow_datasets/text/gem/gem.py           | 90 +++++++++----------
 tensorflow_datasets/text/mrqa/mrqa.py         | 36 ++++----
 .../text/unifiedqa/unifiedqa.py               | 28 +++---
 tensorflow_datasets/translate/wmt.py          |  4 +-
 tensorflow_datasets/video/tao/tao.py          |  2 +-
 tensorflow_datasets/video/ucf101.py           |  2 +-
 .../video/youtube_vis/youtube_vis.py          |  8 +-
 .../vision_language/laion400m/laion400m.py    |  2 +-
 .../vision_language/wit_kaggle/wit_kaggle.py  |  2 -
 24 files changed, 126 insertions(+), 147 deletions(-)

diff --git a/tensorflow_datasets/core/community/dataset_sources.py b/tensorflow_datasets/core/community/dataset_sources.py
index 32dc01935be..56db8bdabdb 100644
--- a/tensorflow_datasets/core/community/dataset_sources.py
+++ b/tensorflow_datasets/core/community/dataset_sources.py
@@ -55,7 +55,7 @@ def to_json(self) -> utils.JsonValue:
     if len(self.filenames) == 1:
       return os.fspath(self.root_path / self.filenames[0])
     else:
-      return {  # pyrefly: ignore[bad-return]
+      return {
           'root_path': os.fspath(self.root_path),
           'filenames': self.filenames,  # pyrefly: ignore[bad-assignment]
       }
diff --git a/tensorflow_datasets/core/community/register_package.py b/tensorflow_datasets/core/community/register_package.py
index 76758e00ea2..f32856dfe5e 100644
--- a/tensorflow_datasets/core/community/register_package.py
+++ b/tensorflow_datasets/core/community/register_package.py
@@ -187,7 +187,7 @@ def refresh(self) -> None:
       return
 
     # If read was sucessful, update the cache with the new dataset list
-    self._cached_path.write_text(content)  # pytype: disable=attribute-error  # attribute-variable-annotations
+    self._cached_path.write_text(content)
     self._refresh_from_content(content)
 
 
@@ -260,7 +260,7 @@ def builder(
       **builder_kwargs: Any,
   ) -> dataset_builder.DatasetBuilder:
     """Returns the dataset builder."""
-    return self.builder_cls(name)(**builder_kwargs)  # pytype: disable=not-instantiable
+    return self.builder_cls(name)(**builder_kwargs)
 
 
 def list_ds_packages_for_namespace(
diff --git a/tensorflow_datasets/core/dataset_builders/conll/conll_dataset_builder.py b/tensorflow_datasets/core/dataset_builders/conll/conll_dataset_builder.py
index a29d3c65860..b6ee0dcf42e 100644
--- a/tensorflow_datasets/core/dataset_builders/conll/conll_dataset_builder.py
+++ b/tensorflow_datasets/core/dataset_builders/conll/conll_dataset_builder.py
@@ -97,7 +97,7 @@ class ConllDatasetBuilder(
   @property
   def builder_config(self) -> ConllBuilderConfig:
     """`tfds.core.BuilderConfig` for this builder."""
-    return self._builder_config  # pytype: disable=bad-return-type  # always-use-return-annotations
+    return self._builder_config  # pyrefly: ignore[bad-return]
 
   def create_dataset_info(
       self,
diff --git a/tensorflow_datasets/core/dataset_builders/conll/conllu_dataset_builder.py b/tensorflow_datasets/core/dataset_builders/conll/conllu_dataset_builder.py
index 1f61e2b7913..b5cc94487e8 100644
--- a/tensorflow_datasets/core/dataset_builders/conll/conllu_dataset_builder.py
+++ b/tensorflow_datasets/core/dataset_builders/conll/conllu_dataset_builder.py
@@ -158,7 +158,7 @@ class ConllUDatasetBuilder(
   @property
   def builder_config(self) -> ConllUBuilderConfig:
     """`tfds.core.BuilderConfig` for this builder."""
-    return self._builder_config  # pytype: disable=bad-return-type  # always-use-return-annotations
+    return self._builder_config  # pyrefly: ignore[bad-return]
 
   def create_dataset_info(
       self,
diff --git a/tensorflow_datasets/rlds/datasets/locomotion/locomotion.py b/tensorflow_datasets/rlds/datasets/locomotion/locomotion.py
index ace700b4cf0..dce56d665a0 100644
--- a/tensorflow_datasets/rlds/datasets/locomotion/locomotion.py
+++ b/tensorflow_datasets/rlds/datasets/locomotion/locomotion.py
@@ -70,7 +70,7 @@ class Locomotion(tfds.core.GeneratorBasedBuilder):
           description=(
               'Dataset generated by a SAC agent trained for 1M steps for Ant.'
           ),
-          supervised_keys=None,  # pytype: disable=wrong-arg-types  # gen-stub-imports
+          supervised_keys=None,
       ),
       rlds_base.DatasetConfig(
           name='hopper_sac_1M_single_policy_stochastic',
@@ -85,7 +85,7 @@ class Locomotion(tfds.core.GeneratorBasedBuilder):
               'Dataset generated by a SAC agent trained for 1M steps for'
               ' Hopper.'
           ),
-          supervised_keys=None,  # pytype: disable=wrong-arg-types  # gen-stub-imports
+          supervised_keys=None,
       ),
       rlds_base.DatasetConfig(
           name='halfcheetah_sac_1M_single_policy_stochastic',
@@ -100,7 +100,7 @@ class Locomotion(tfds.core.GeneratorBasedBuilder):
               'Dataset generated by a SAC agent trained for 1M steps for'
               ' HalfCheetah.'
           ),
-          supervised_keys=None,  # pytype: disable=wrong-arg-types  # gen-stub-imports
+          supervised_keys=None,
       ),
       rlds_base.DatasetConfig(
           name='walker2d_sac_1M_single_policy_stochastic',
@@ -115,7 +115,7 @@ class Locomotion(tfds.core.GeneratorBasedBuilder):
               'Dataset generated by a SAC agent trained for 1M steps for'
               ' Walker2d.'
           ),
-          supervised_keys=None,  # pytype: disable=wrong-arg-types  # gen-stub-imports
+          supervised_keys=None,
       ),
       rlds_base.DatasetConfig(
           name='humanoid_sac_15M_single_policy_stochastic',
@@ -130,7 +130,7 @@ class Locomotion(tfds.core.GeneratorBasedBuilder):
               'Dataset generated by a SAC agent trained for 15M steps for'
               ' Humanoid.'
           ),
-          supervised_keys=None,  # pytype: disable=wrong-arg-types  # gen-stub-imports
+          supervised_keys=None,
       ),
   ]
 
diff --git a/tensorflow_datasets/rlds/datasets/robosuite_panda_pick_place_can/robosuite_panda_pick_place_can.py b/tensorflow_datasets/rlds/datasets/robosuite_panda_pick_place_can/robosuite_panda_pick_place_can.py
index 180ed45f138..b2ed981d15f 100644
--- a/tensorflow_datasets/rlds/datasets/robosuite_panda_pick_place_can/robosuite_panda_pick_place_can.py
+++ b/tensorflow_datasets/rlds/datasets/robosuite_panda_pick_place_can/robosuite_panda_pick_place_can.py
@@ -72,7 +72,6 @@ class RobosuitePandaPickPlaceCan(tfds.core.GeneratorBasedBuilder):
       'synthetic_stochastic_sac_afe13968': 'https://storage.googleapis.com/rlds_external_data_release/rlds_robosuite_panda_pick_place_can_synthetic_stochastic_sac_afe13968.tar.gz',
   }
 
-  # pytype: disable=wrong-keyword-args
   BUILDER_CONFIGS = [
       rlds_base.DatasetConfig(
           name='human_dc29b40a',
@@ -114,22 +113,22 @@ class RobosuitePandaPickPlaceCan(tfds.core.GeneratorBasedBuilder):
               ),
           },
           action_info=tfds.features.Tensor(shape=(7,), dtype=np.float64),
-          reward_info=np.float64,
-          discount_info=np.float64,
+          reward_info=np.float64,  # pyrefly: ignore[bad-argument-type]
+          discount_info=np.float64,  # pyrefly: ignore[bad-argument-type]
           episode_metadata_info={
-              'agent_id': np.str_,
-              'episode_index': np.int32,
-              'episode_id': np.str_,
+              'agent_id': np.str_,  # pyrefly: ignore[bad-assignment]
+              'episode_index': np.int32,  # pyrefly: ignore[bad-assignment]
+              'episode_id': np.str_,  # pyrefly: ignore[bad-assignment]
           },
           step_metadata_info={
-              'tag:placed': np.bool_,
+              'tag:placed': np.bool_,  # pyrefly: ignore[bad-assignment]
               'image': tfds.features.Image(),
           },
           citation=_CITATION,
           homepage=_HOMEPAGE,
           overall_description=_DESCRIPTION,
           description='Human generated dataset (50 episodes).',
-          supervised_keys=None,  # pytype: disable=wrong-arg-types  # gen-stub-imports
+          supervised_keys=None,
       ),
       rlds_base.DatasetConfig(
           name='human_images_dc29b40a',
@@ -183,15 +182,15 @@ class RobosuitePandaPickPlaceCan(tfds.core.GeneratorBasedBuilder):
               ),
           },
           action_info=tfds.features.Tensor(shape=(7,), dtype=np.float64),
-          reward_info=np.float64,
-          discount_info=np.float64,
+          reward_info=np.float64,  # pyrefly: ignore[bad-argument-type]
+          discount_info=np.float64,  # pyrefly: ignore[bad-argument-type]
           episode_metadata_info={
-              'agent_id': np.str_,
-              'episode_index': np.int32,
-              'episode_id': np.str_,
+              'agent_id': np.str_,  # pyrefly: ignore[bad-assignment]
+              'episode_index': np.int32,  # pyrefly: ignore[bad-assignment]
+              'episode_id': np.str_,  # pyrefly: ignore[bad-assignment]
           },
           step_metadata_info={
-              'tag:placed': np.bool_,
+              'tag:placed': np.bool_,  # pyrefly: ignore[bad-assignment]
               'image': tfds.features.Image(),
           },
           citation=_CITATION,
@@ -201,7 +200,7 @@ class RobosuitePandaPickPlaceCan(tfds.core.GeneratorBasedBuilder):
               ' angles in the observation.'
               ' Note that it may take some time to generate.'
           ),
-          supervised_keys=None,  # pytype: disable=wrong-arg-types  # gen-stub-imports
+          supervised_keys=None,
       ),
       rlds_base.DatasetConfig(
           name='synthetic_stochastic_sac_afe13968',
@@ -243,15 +242,15 @@ class RobosuitePandaPickPlaceCan(tfds.core.GeneratorBasedBuilder):
               ),
           },
           action_info=tfds.features.Tensor(shape=(7,), dtype=np.float32),
-          reward_info=np.float64,
-          discount_info=np.float64,
+          reward_info=np.float64,  # pyrefly: ignore[bad-argument-type]
+          discount_info=np.float64,  # pyrefly: ignore[bad-argument-type]
           episode_metadata_info={
-              'agent_id': np.str_,
-              'episode_index': np.int32,
-              'episode_id': np.str_,
+              'agent_id': np.str_,  # pyrefly: ignore[bad-assignment]
+              'episode_index': np.int32,  # pyrefly: ignore[bad-assignment]
+              'episode_id': np.str_,  # pyrefly: ignore[bad-assignment]
           },
           step_metadata_info={
-              'tag:placed': np.bool_,
+              'tag:placed': np.bool_,  # pyrefly: ignore[bad-assignment]
               'image': tfds.features.Image(),
           },
           citation=_CITATION,
@@ -261,12 +260,10 @@ class RobosuitePandaPickPlaceCan(tfds.core.GeneratorBasedBuilder):
               'Synthetic dataset generated by a stochastic agent trained with'
               ' SAC (200 episodes).'
           ),
-          supervised_keys=None,  # pytype: disable=wrong-arg-types  # gen-stub-imports
+          supervised_keys=None,
       ),
   ]
 
-  # pytype: enable=wrong-keyword-args
-
 
   def _info(self) -> tfds.core.DatasetInfo:
     """Returns the dataset metadata."""
diff --git a/tensorflow_datasets/robomimic/dataset_utils.py b/tensorflow_datasets/robomimic/dataset_utils.py
index 743b4f66f4b..1a1914d31d5 100644
--- a/tensorflow_datasets/robomimic/dataset_utils.py
+++ b/tensorflow_datasets/robomimic/dataset_utils.py
@@ -112,7 +112,6 @@ def make_builder_configs(dataset: DataSource):
   for task, details in TASKS.items():
     if dataset in details['datasets']:  # pyrefly: ignore[not-iterable]
       for observation_type in [ObservationType.IMAGE, ObservationType.LOW_DIM]:
-        # pytype: disable=wrong-keyword-args
         configs.append(
             # name is inherited from base dataclass unbeknownst to the linter
             BuilderConfig(  # pylint: disable=unexpected-keyword-arg
@@ -120,10 +119,9 @@ def make_builder_configs(dataset: DataSource):
                 task=task,
                 dataset=dataset,
                 filename=observation_type,
-                horizon=details['horizon'],
+                horizon=details['horizon'],  # pyrefly: ignore[bad-argument-type]
             )
         )
-        # pytype: enable=wrong-keyword-args
   return configs
 
 
@@ -301,7 +299,7 @@ def _get_features(self) -> tfds.features.FeaturesDict:
     features = tfds.features.FeaturesDict({
         'horizon': np.int32,
         'episode_id': np.str_,
-        'steps': tfds.features.Dataset({  # pyrefly: ignore[bad-argument-type]
+        'steps': tfds.features.Dataset({
             'action': tensor_feature(action_size),  # pyrefly: ignore[bad-argument-type]
             'observation': observation,  # pyrefly: ignore[bad-assignment]
             'reward': np.float64,
diff --git a/tensorflow_datasets/scripts/deployment/export_community_datasets.py b/tensorflow_datasets/scripts/deployment/export_community_datasets.py
index 55002ea87d1..750d4bcf81c 100644
--- a/tensorflow_datasets/scripts/deployment/export_community_datasets.py
+++ b/tensorflow_datasets/scripts/deployment/export_community_datasets.py
@@ -36,7 +36,7 @@
 
 
 def main(_):
-  export_community_datasets(in_path=_IN_PATH, out_path=_OUT_PATH)  # pytype: disable=wrong-arg-types
+  export_community_datasets(in_path=_IN_PATH, out_path=_OUT_PATH)
 
 
 def export_community_datasets(
diff --git a/tensorflow_datasets/scripts/documentation/document_datasets.py b/tensorflow_datasets/scripts/documentation/document_datasets.py
index d7cb512e2aa..d6d04184d5b 100644
--- a/tensorflow_datasets/scripts/documentation/document_datasets.py
+++ b/tensorflow_datasets/scripts/documentation/document_datasets.py
@@ -235,7 +235,7 @@ def get_config_builder(config) -> tfds.core.DatasetBuilder:
     return BuilderToDocument(
         sections=sections,
         namespace=None,
-        builder=builder_cls(),  # pytype: disable=not-instantiable
+        builder=builder_cls(),
         config_builders=[],
     )
 
@@ -365,30 +365,28 @@ def iter_documentation_builders(
   print('Retrieving the list of builders...')
   datasets = datasets or _all_tfds_datasets()
 
-  # pytype: disable=attribute-error
-  if doc_util_paths.fig_base_path:
+  if doc_util_paths.fig_base_path:  # pyrefly: ignore[missing-attribute]
     visu_doc_util = doc_utils.VisualizationDocUtil(
         base_path=doc_util_paths.fig_base_path,
-        base_url=doc_util_paths.fig_base_url,
+        base_url=doc_util_paths.fig_base_url,  # pyrefly: ignore[missing-attribute]
     )
   else:
     visu_doc_util = None
 
-  if doc_util_paths.df_base_path:
+  if doc_util_paths.df_base_path:  # pyrefly: ignore[missing-attribute]
     df_doc_util = doc_utils.DataframeDocUtil(
         base_path=doc_util_paths.df_base_path,
-        base_url=doc_util_paths.df_base_url,
+        base_url=doc_util_paths.df_base_url,  # pyrefly: ignore[missing-attribute]
     )
   else:
     df_doc_util = None
 
   if doc_util_paths.fig_base_path:
     nightly_doc_util = doc_utils.NightlyDocUtil(
-        path=doc_util_paths.nightly_path,
+        path=doc_util_paths.nightly_path,  # pyrefly: ignore[missing-attribute]
     )
   else:
     nightly_doc_util = None
-  # pytype: enable=attribute-error
 
   document_single_builder_fn = functools.partial(
       _document_single_builder,
diff --git a/tensorflow_datasets/structured/cherry_blossoms/cherry_blossoms.py b/tensorflow_datasets/structured/cherry_blossoms/cherry_blossoms.py
index 80f634cf60f..2457adce82a 100644
--- a/tensorflow_datasets/structured/cherry_blossoms/cherry_blossoms.py
+++ b/tensorflow_datasets/structured/cherry_blossoms/cherry_blossoms.py
@@ -90,7 +90,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
 
   def _generate_examples(self, path):
     """Yields examples."""
-    with path.open() as f:  # pytype: disable=attribute-error  # gen-stub-imports
+    with path.open() as f:
       reader = csv.DictReader(f, delimiter=';')
       for index, row in enumerate(reader):
         example = {k: v.replace('NA', 'nan') for k, v in row.items()}
diff --git a/tensorflow_datasets/structured/howell/howell.py b/tensorflow_datasets/structured/howell/howell.py
index 09a63c1541c..338bb717836 100644
--- a/tensorflow_datasets/structured/howell/howell.py
+++ b/tensorflow_datasets/structured/howell/howell.py
@@ -89,7 +89,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
 
   def _generate_examples(self, file_path):
     """Yields examples."""
-    with file_path.open() as f:  # pytype: disable=attribute-error  # gen-stub-imports
+    with file_path.open() as f:
       reader = csv.DictReader(f, delimiter=";")
       for index, row in enumerate(reader):
         example = dict(row.items())
diff --git a/tensorflow_datasets/structured/wine_quality/wine_quality.py b/tensorflow_datasets/structured/wine_quality/wine_quality.py
index b9847ab1657..f6a17ebc546 100644
--- a/tensorflow_datasets/structured/wine_quality/wine_quality.py
+++ b/tensorflow_datasets/structured/wine_quality/wine_quality.py
@@ -119,7 +119,7 @@ def _info(self):
     return tfds.core.DatasetInfo(
         builder=self,
         description=_DESCRIPTION,
-        features=tfds.features.FeaturesDict({  # pyrefly: ignore[bad-argument-type]
+        features=tfds.features.FeaturesDict({
             "quality": np.int32,
             "features": features_dict,  # pyrefly: ignore[bad-assignment]
         }),
diff --git a/tensorflow_datasets/testing/huggingface/dummy_dataset.py b/tensorflow_datasets/testing/huggingface/dummy_dataset.py
index 200ffee40b3..1a885b1cf22 100644
--- a/tensorflow_datasets/testing/huggingface/dummy_dataset.py
+++ b/tensorflow_datasets/testing/huggingface/dummy_dataset.py
@@ -17,7 +17,7 @@
 
 from pathlib import Path  # < Direct import of path should be patched  # pylint: disable=g-importing-member
 
-import datasets  # pytype: disable=import-error
+import datasets  # pyrefly: ignore[missing-import]
 import tensorflow_datasets.public_api as tfds
 
 
diff --git a/tensorflow_datasets/text/c4_wsrs/c4_wsrs_utils.py b/tensorflow_datasets/text/c4_wsrs/c4_wsrs_utils.py
index 9a91a188538..85e131f3656 100644
--- a/tensorflow_datasets/text/c4_wsrs/c4_wsrs_utils.py
+++ b/tensorflow_datasets/text/c4_wsrs/c4_wsrs_utils.py
@@ -265,7 +265,7 @@ def extract_snippets(
         continue
       count_by_pair[('', '')] += 1
 
-    yield key, (snippet, rarest_pair, index_to_pair)  # pytype: disable=bad-return-type  # container-simplification
+    yield key, (snippet, rarest_pair, index_to_pair)  # pyrefly: ignore[invalid-yield]
 
 
 def _reverse_substitute_snippet(
diff --git a/tensorflow_datasets/text/dolphin_number_word/dolphin_number_word.py b/tensorflow_datasets/text/dolphin_number_word/dolphin_number_word.py
index 68149ed511d..dde15cdac15 100644
--- a/tensorflow_datasets/text/dolphin_number_word/dolphin_number_word.py
+++ b/tensorflow_datasets/text/dolphin_number_word/dolphin_number_word.py
@@ -90,7 +90,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
   def _generate_examples(self, json_paths, split):
     """Parses available JSON files in a given split, returns TF examples."""
     # Currently assume all files in archive are json-formatted.
-    for file in json_paths.glob('*' + split + '*'):  # pytype: disable=attribute-error  # gen-stub-imports
+    for file in json_paths.glob('*' + split + '*'):
       json_string = file.read_text()
       examples = json.loads(json_string)
       for e in examples:
diff --git a/tensorflow_datasets/text/gem/gem.py b/tensorflow_datasets/text/gem/gem.py
index 2a5098873ae..bbee89a5ae5 100644
--- a/tensorflow_datasets/text/gem/gem.py
+++ b/tensorflow_datasets/text/gem/gem.py
@@ -786,7 +786,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
                 name=challenge_split,
                 gen_kwargs={
                     "filepath": os.path.join(
-                        files["challenge_set"],  # pyrefly: ignore[bad-index]
+                        files["challenge_set"],
                         self.builder_config.name,
                         filename,
                     ),
@@ -800,7 +800,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
               name=tfds.Split.TRAIN,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
                   "filepath": os.path.join(
-                      files["data"], "commongen.train.jsonl"  # pyrefly: ignore[bad-index]
+                      files["data"], "commongen.train.jsonl"
                   ),
                   "set_name": "train",
               },
@@ -809,7 +809,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
               name=tfds.Split.VALIDATION,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
                   "filepath": os.path.join(
-                      files["data"], "commongen.dev.jsonl"  # pyrefly: ignore[bad-index]
+                      files["data"], "commongen.dev.jsonl"
                   ),
                   "set_name": "validation",
               },
@@ -818,7 +818,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
               name=tfds.Split.TEST,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
                   "filepath": os.path.join(
-                      files["data"], "commongen.test_noref.jsonl"  # pyrefly: ignore[bad-index]
+                      files["data"], "commongen.test_noref.jsonl"
                   ),
                   "set_name": "test",
               },
@@ -847,7 +847,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
                 name=challenge_split,
                 gen_kwargs={
                     "filepath": os.path.join(
-                        files["challenge_set"],  # pyrefly: ignore[bad-index]
+                        files["challenge_set"],
                         self.builder_config.name,
                         filename,
                     ),
@@ -860,21 +860,21 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
           tfds.core.SplitGenerator(
               name=tfds.Split.TRAIN,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
-                  "filepath": files["train"],  # pyrefly: ignore[bad-index]
+                  "filepath": files["train"],
                   "set_name": "train",
               },
           ),
           tfds.core.SplitGenerator(
               name=tfds.Split.VALIDATION,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
-                  "filepath": files["validation"],  # pyrefly: ignore[bad-index]
+                  "filepath": files["validation"],
                   "set_name": "validation",
               },
           ),
           tfds.core.SplitGenerator(
               name=tfds.Split.TEST,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
-                  "filepath": files["test"],  # pyrefly: ignore[bad-index]
+                  "filepath": files["test"],
                   "set_name": "test",
               },
           ),
@@ -884,21 +884,21 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
           tfds.core.SplitGenerator(
               name=tfds.Split.TRAIN,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
-                  "filepath": files["train"],  # pyrefly: ignore[bad-index]
+                  "filepath": files["train"],
                   "set_name": "train",
               },
           ),
           tfds.core.SplitGenerator(
               name=tfds.Split.VALIDATION,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
-                  "filepath": files["validation"],  # pyrefly: ignore[bad-index]
+                  "filepath": files["validation"],
                   "set_name": "validation",
               },
           ),
           tfds.core.SplitGenerator(
               name=tfds.Split.TEST,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
-                  "filepath": files["test"],  # pyrefly: ignore[bad-index]
+                  "filepath": files["test"],
                   "set_name": "test",
               },
           ),
@@ -922,7 +922,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
                 name=challenge_split,
                 gen_kwargs={
                     "filepath": os.path.join(
-                        files["challenge_set"],  # pyrefly: ignore[bad-index]
+                        files["challenge_set"],
                         self.builder_config.name,
                         filename,
                     ),
@@ -934,21 +934,21 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
           tfds.core.SplitGenerator(
               name=tfds.Split.TRAIN,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
-                  "filepath": files["train"],  # pyrefly: ignore[bad-index]
+                  "filepath": files["train"],
                   "set_name": "train",
               },
           ),
           tfds.core.SplitGenerator(
               name=tfds.Split.VALIDATION,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
-                  "filepath": files["validation"],  # pyrefly: ignore[bad-index]
+                  "filepath": files["validation"],
                   "set_name": "validation",
               },
           ),
           tfds.core.SplitGenerator(
               name=tfds.Split.TEST,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
-                  "filepath": files["test"],  # pyrefly: ignore[bad-index]
+                  "filepath": files["test"],
                   "set_name": "test",
               },
           ),
@@ -974,7 +974,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
                 name=challenge_split,
                 gen_kwargs={
                     "filepath": os.path.join(  # pyrefly: ignore[no-matching-overload]
-                        files["challenge_set"],  # pyrefly: ignore[bad-index]
+                        files["challenge_set"],
                         self.builder_config.name,
                         filename,
                     ),
@@ -987,31 +987,31 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
               name=tfds.Split.TRAIN,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
                   "filepath": os.path.join(
-                      files["train"], lang + "_train.jsonl"  # pyrefly: ignore[bad-index]
+                      files["train"], lang + "_train.jsonl"
                   ),
                   "set_name": "train",
                   "lang": lang,
-                  "filepaths": files["bad_ids"],  # pyrefly: ignore[bad-index]
+                  "filepaths": files["bad_ids"],
               },
           ),
           tfds.core.SplitGenerator(
               name=tfds.Split.VALIDATION,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
                   "filepath": os.path.join(
-                      files["validation"], lang + "_val.jsonl"  # pyrefly: ignore[bad-index]
+                      files["validation"], lang + "_val.jsonl"
                   ),
                   "set_name": "validation",
                   "lang": lang,
-                  "filepaths": files["bad_ids"],  # pyrefly: ignore[bad-index]
+                  "filepaths": files["bad_ids"],
               },
           ),
           tfds.core.SplitGenerator(
               name=tfds.Split.TEST,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
-                  "filepath": os.path.join(files["test"], lang + "_test.jsonl"),  # pyrefly: ignore[bad-index]
+                  "filepath": os.path.join(files["test"], lang + "_test.jsonl"),
                   "set_name": "test",
                   "lang": lang,
-                  "filepaths": files["bad_ids"],  # pyrefly: ignore[bad-index]
+                  "filepaths": files["bad_ids"],
               },
           ),
       ] + challenge_splits
@@ -1053,7 +1053,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
                 name=challenge_split,
                 gen_kwargs={
                     "filepath": os.path.join(
-                        files["challenge_set"],  # pyrefly: ignore[bad-index]
+                        files["challenge_set"],
                         self.builder_config.name,
                         filename,
                     ),
@@ -1071,7 +1071,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
             tfds.core.SplitGenerator(
                 name=tfds_spl,
                 gen_kwargs={
-                    "filepath": os.path.join(files["data"], "gem_sgd.json"),  # pyrefly: ignore[bad-index]
+                    "filepath": os.path.join(files["data"], "gem_sgd.json"),
                     "set_name": spl,
                 },
             )
@@ -1096,7 +1096,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
                 name=challenge_split,
                 gen_kwargs={
                     "filepath": os.path.join(
-                        files["challenge_set"],  # pyrefly: ignore[bad-index]
+                        files["challenge_set"],
                         self.builder_config.name,
                         filename,
                     ),
@@ -1110,7 +1110,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
               name=tfds.Split.TRAIN,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
                   "filepath": os.path.join(
-                      files["data"], "totto_data/totto_train_data.jsonl"  # pyrefly: ignore[bad-index]
+                      files["data"], "totto_data/totto_train_data.jsonl"
                   ),
                   "set_name": "train",
               },
@@ -1119,7 +1119,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
               name=tfds.Split.VALIDATION,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
                   "filepath": os.path.join(
-                      files["data"], "totto_data/totto_dev_data.jsonl"  # pyrefly: ignore[bad-index]
+                      files["data"], "totto_data/totto_dev_data.jsonl"
                   ),
                   "set_name": "validation",
               },
@@ -1128,7 +1128,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
               name=tfds.Split.TEST,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
                   "filepath": os.path.join(
-                      files["data"],  # pyrefly: ignore[bad-index]
+                      files["data"],
                       "totto_data/unlabeled_totto_test_data.jsonl",
                   ),
                   "set_name": "test",
@@ -1164,7 +1164,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
                 name=challenge_split,
                 gen_kwargs={
                     "filepath": os.path.join(  # pyrefly: ignore[no-matching-overload]
-                        files["challenge_set"],  # pyrefly: ignore[bad-index]
+                        files["challenge_set"],
                         self.builder_config.name,
                         filename,
                     ),
@@ -1175,18 +1175,18 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
       return [
           tfds.core.SplitGenerator(
               name=tfds.Split.TRAIN,  # pyrefly: ignore[missing-attribute]
-              gen_kwargs={"filepath": files["train"], "set_name": "train"},  # pyrefly: ignore[bad-index]
+              gen_kwargs={"filepath": files["train"], "set_name": "train"},
           ),
           tfds.core.SplitGenerator(
               name=tfds.Split.VALIDATION,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
-                  "filepath": files["validation"],  # pyrefly: ignore[bad-index]
+                  "filepath": files["validation"],
                   "set_name": "validation",
               },
           ),
           tfds.core.SplitGenerator(
               name=tfds.Split.TEST,  # pyrefly: ignore[missing-attribute]
-              gen_kwargs={"filepath": files["test"], "set_name": "test"},  # pyrefly: ignore[bad-index]
+              gen_kwargs={"filepath": files["test"], "set_name": "test"},
           ),
       ] + challenge_splits
     elif self.builder_config.name == "wiki_auto_asset_turk":
@@ -1239,7 +1239,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
                 name=challenge_split,
                 gen_kwargs={
                     "filepath": os.path.join(
-                        files["challenge_set"],  # pyrefly: ignore[bad-index]
+                        files["challenge_set"],
                         self.builder_config.name,
                         filename,
                     ),
@@ -1251,14 +1251,14 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
           tfds.core.SplitGenerator(
               name=tfds.Split.TRAIN,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
-                  "filepath": files["train"],  # pyrefly: ignore[bad-index]
+                  "filepath": files["train"],
                   "set_name": "train",
               },
           ),
           tfds.core.SplitGenerator(
               name=tfds.Split.VALIDATION,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
-                  "filepath": files["validation"],  # pyrefly: ignore[bad-index]
+                  "filepath": files["validation"],
                   "set_name": "validation",
               },
           ),
@@ -1268,14 +1268,14 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
                   "filepath": "",
                   "set_name": "test_asset",
                   "filepaths": [
-                      files["test_asset_" + str(i)] for i in range(10)  # pyrefly: ignore[bad-index]
+                      files["test_asset_" + str(i)] for i in range(10)
                   ],
               },
           ),
           tfds.core.SplitGenerator(
               name="test_turk",
               gen_kwargs={
-                  "filepath": files["test_turk"],  # pyrefly: ignore[bad-index]
+                  "filepath": files["test_turk"],
                   "set_name": "test_turk",
               },
           ),
@@ -1283,7 +1283,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
     elif self.builder_config.name.startswith("wiki_lingua"):  # pyrefly: ignore[missing-attribute]
       lang_name = self.builder_config.name.split("_")[-2]  # pyrefly: ignore[missing-attribute]
       lang = self.builder_config.name.split("_")[-1]  # pyrefly: ignore[missing-attribute]
-      base_dir = os.path.join(files["data"], lang_name)  # pyrefly: ignore[bad-index]
+      base_dir = os.path.join(files["data"], lang_name)
       return [
           tfds.core.SplitGenerator(
               name=tfds.Split.TRAIN,  # pyrefly: ignore[missing-attribute]
@@ -1339,7 +1339,7 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
                 name=challenge_split,
                 gen_kwargs={
                     "filepath": os.path.join(
-                        files["challenge_set"],  # pyrefly: ignore[bad-index]
+                        files["challenge_set"],
                         self.builder_config.name,
                         filename,
                     ),
@@ -1351,25 +1351,25 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
           tfds.core.SplitGenerator(
               name=tfds.Split.TRAIN,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
-                  "filepath": files["splits"],  # pyrefly: ignore[bad-index]
+                  "filepath": files["splits"],
                   "set_name": "train",
-                  "filepaths": os.path.join(files["data"], "bbc-summary-data"),  # pyrefly: ignore[bad-index]
+                  "filepaths": os.path.join(files["data"], "bbc-summary-data"),
               },
           ),
           tfds.core.SplitGenerator(
               name=tfds.Split.VALIDATION,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
-                  "filepath": files["splits"],  # pyrefly: ignore[bad-index]
+                  "filepath": files["splits"],
                   "set_name": "validation",
-                  "filepaths": os.path.join(files["data"], "bbc-summary-data"),  # pyrefly: ignore[bad-index]
+                  "filepaths": os.path.join(files["data"], "bbc-summary-data"),
               },
           ),
           tfds.core.SplitGenerator(
               name=tfds.Split.TEST,  # pyrefly: ignore[missing-attribute]
               gen_kwargs={
-                  "filepath": files["splits"],  # pyrefly: ignore[bad-index]
+                  "filepath": files["splits"],
                   "set_name": "test",
-                  "filepaths": os.path.join(files["data"], "bbc-summary-data"),  # pyrefly: ignore[bad-index]
+                  "filepaths": os.path.join(files["data"], "bbc-summary-data"),
               },
           ),
       ] + challenge_splits
diff --git a/tensorflow_datasets/text/mrqa/mrqa.py b/tensorflow_datasets/text/mrqa/mrqa.py
index 624dd275b4c..ff8cbcc3a8b 100644
--- a/tensorflow_datasets/text/mrqa/mrqa.py
+++ b/tensorflow_datasets/text/mrqa/mrqa.py
@@ -606,31 +606,25 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
     split_generators = {}
 
     if 'train' in self.builder_config.data_urls:  # pyrefly: ignore[missing-attribute]
-      split_generators.update(
-          {
-              tfds.Split.TRAIN: self._generate_examples(  # pyrefly: ignore[missing-attribute]
-                  path=data_dir['train'], split='train'  # pyrefly: ignore[bad-index]
-              ),
-          }
-      )
+      split_generators.update({
+          tfds.Split.TRAIN: self._generate_examples(  # pyrefly: ignore[missing-attribute]
+              path=data_dir['train'], split='train'
+          ),
+      })
 
     if 'validation' in self.builder_config.data_urls:  # pyrefly: ignore[not-iterable]
-      split_generators.update(
-          {
-              tfds.Split.VALIDATION: self._generate_examples(  # pyrefly: ignore[missing-attribute]
-                  path=data_dir['validation'], split='validation'  # pyrefly: ignore[bad-index]
-              ),
-          }
-      )
+      split_generators.update({
+          tfds.Split.VALIDATION: self._generate_examples(  # pyrefly: ignore[missing-attribute]
+              path=data_dir['validation'], split='validation'
+          ),
+      })
 
     if 'test' in self.builder_config.data_urls:  # pyrefly: ignore[not-iterable]
-      split_generators.update(
-          {
-              tfds.Split.TEST: self._generate_examples(  # pyrefly: ignore[missing-attribute]
-                  path=data_dir['test'], split='test'  # pyrefly: ignore[bad-index]
-              ),
-          }
-      )
+      split_generators.update({
+          tfds.Split.TEST: self._generate_examples(  # pyrefly: ignore[missing-attribute]
+              path=data_dir['test'], split='test'
+          ),
+      })
     return split_generators
 
   def _generate_examples(self, path, split):
diff --git a/tensorflow_datasets/text/unifiedqa/unifiedqa.py b/tensorflow_datasets/text/unifiedqa/unifiedqa.py
index b4e02035bbe..c33dde170c6 100644
--- a/tensorflow_datasets/text/unifiedqa/unifiedqa.py
+++ b/tensorflow_datasets/text/unifiedqa/unifiedqa.py
@@ -1970,27 +1970,21 @@ def _split_generators(self, dl_manager: tfds.download.DownloadManager):
     split_generators = {}
 
     if 'train' in self.builder_config.data_urls:  # pyrefly: ignore[missing-attribute]
-      split_generators.update(
-          {
-              tfds.Split.TRAIN: self._generate_examples(path=data_dir['train']),  # pyrefly: ignore[bad-index, missing-attribute]
-          }
-      )
+      split_generators.update({
+          tfds.Split.TRAIN: self._generate_examples(path=data_dir['train']),  # pyrefly: ignore[missing-attribute]
+      })
 
     if 'validation' in self.builder_config.data_urls:  # pyrefly: ignore[not-iterable]
-      split_generators.update(
-          {
-              tfds.Split.VALIDATION: self._generate_examples(  # pyrefly: ignore[missing-attribute]
-                  path=data_dir['validation']  # pyrefly: ignore[bad-index]
-              ),
-          }
-      )
+      split_generators.update({
+          tfds.Split.VALIDATION: self._generate_examples(  # pyrefly: ignore[missing-attribute]
+              path=data_dir['validation']
+          ),
+      })
 
     if 'test' in self.builder_config.data_urls:  # pyrefly: ignore[not-iterable]
-      split_generators.update(
-          {
-              tfds.Split.TEST: self._generate_examples(path=data_dir['test']),  # pyrefly: ignore[bad-index, missing-attribute]
-          }
-      )
+      split_generators.update({
+          tfds.Split.TEST: self._generate_examples(path=data_dir['test']),  # pyrefly: ignore[missing-attribute]
+      })
     return split_generators
 
   def _generate_examples(self, path):
diff --git a/tensorflow_datasets/translate/wmt.py b/tensorflow_datasets/translate/wmt.py
index 402e501b972..84d6759cd7a 100644
--- a/tensorflow_datasets/translate/wmt.py
+++ b/tensorflow_datasets/translate/wmt.py
@@ -1056,7 +1056,7 @@ def _get_tuv_seg(tuv):
 
   with epath.Path(path).open("rb") as f:
     utf_f = codecs.getreader("utf-8")(f)
-    for line_id, (_, elem) in enumerate(ElementTree.iterparse(utf_f)):  # pytype: disable=wrong-arg-types
+    for line_id, (_, elem) in enumerate(ElementTree.iterparse(utf_f)):
       if elem.tag == "tu":
         yield line_id, {
             _get_tuv_lang(tuv): _get_tuv_seg(tuv)
@@ -1106,7 +1106,7 @@ def _parse_czeng(*paths, **kwargs):
     re_block = re.compile(r"^[^-]+-b(\d+)-\d\d[tde]")
     with epath.Path(filter_path).open() as f:
       bad_blocks = set(
-          re.search(r"qw{([\s\d]*)}", f.read()).groups()[0].split()  # pytype: disable=attribute-error
+          re.search(r"qw{([\s\d]*)}", f.read()).groups()[0].split()  # pyrefly: ignore[missing-attribute]
       )
     logging.info(
         "Loaded %d bad blocks to filter from CzEng v1.6 to make v1.7.",
diff --git a/tensorflow_datasets/video/tao/tao.py b/tensorflow_datasets/video/tao/tao.py
index a56e5aeb235..bc33b7bb4d3 100644
--- a/tensorflow_datasets/video/tao/tao.py
+++ b/tensorflow_datasets/video/tao/tao.py
@@ -296,7 +296,7 @@ def _info(self) -> tfds.core.DatasetInfo:
         3,
     )
     all_features = {
-        'video': tfds.features.Video(video_shape),  # pytype: disable=wrong-arg-types  # gen-stub-imports
+        'video': tfds.features.Video(video_shape),
         'metadata': {
             'height': np.int32,
             'width': np.int32,
diff --git a/tensorflow_datasets/video/ucf101.py b/tensorflow_datasets/video/ucf101.py
index 5b7ddfba375..9a547ba6302 100644
--- a/tensorflow_datasets/video/ucf101.py
+++ b/tensorflow_datasets/video/ucf101.py
@@ -146,7 +146,7 @@ def _info(self):
             video_shape,  # pyrefly: ignore[bad-argument-type]
             ffmpeg_extra_args=ffmpeg_extra_args,
             encoding_format='jpeg',
-        ),  # pytype: disable=wrong-arg-types  # gen-stub-imports
+        ),
         'label': tfds.features.ClassLabel(names_file=labels_names_file),
     })
     return tfds.core.DatasetInfo(
diff --git a/tensorflow_datasets/video/youtube_vis/youtube_vis.py b/tensorflow_datasets/video/youtube_vis/youtube_vis.py
index 9e6c688c492..27c7fbd21cf 100644
--- a/tensorflow_datasets/video/youtube_vis/youtube_vis.py
+++ b/tensorflow_datasets/video/youtube_vis/youtube_vis.py
@@ -71,7 +71,7 @@ def _convert_bbox(
       ymin=box[1] / height,
       xmax=(box[0] + box[2]) / width,
       ymax=(box[1] + box[3]) / height,
-  )  # pytype: disable=bad-return-type  # gen-stub-imports
+  )
 
 
 def _decode_segmentation(
@@ -95,7 +95,7 @@ def _decode_segmentation(
         (desired_width, desired_height),
         interpolation=cv2.INTER_NEAREST,
     )
-  segmentation = np.expand_dims(segmentation, axis=-1)  # pyrefly: ignore[bad-assignment, no-matching-overload]
+  segmentation = np.expand_dims(segmentation, axis=-1)  # pyrefly: ignore[no-matching-overload]
   assert len(segmentation.shape) == 3  # pyrefly: ignore[missing-attribute]
   return segmentation
 
@@ -432,7 +432,7 @@ def _info(self) -> tfds.core.DatasetInfo:
     )
     seg_shape = (None, self.builder_config.height, self.builder_config.width, 1)  # pyrefly: ignore[missing-attribute]
     all_features = {
-        'video': tfds.features.Video(video_shape),  # pytype: disable=wrong-arg-types  # gen-stub-imports
+        'video': tfds.features.Video(video_shape),
         'metadata': {
             'height': np.int32,
             'width': np.int32,
@@ -441,7 +441,7 @@ def _info(self) -> tfds.core.DatasetInfo:
         },
         'tracks': tfds.features.Sequence({
             'bboxes': tfds.features.Sequence(tfds.features.BBoxFeature()),
-            'segmentations': tfds.features.Video(seg_shape, use_colormap=True),  # pytype: disable=wrong-arg-types  # gen-stub-imports
+            'segmentations': tfds.features.Video(seg_shape, use_colormap=True),
             'category': tfds.features.ClassLabel(names_file=names_file),
             'is_crowd': np.bool_,
             'areas': tfds.features.Sequence(np.float32),
diff --git a/tensorflow_datasets/vision_language/laion400m/laion400m.py b/tensorflow_datasets/vision_language/laion400m/laion400m.py
index 28d1cb6d27f..ad3f888d124 100644
--- a/tensorflow_datasets/vision_language/laion400m/laion400m.py
+++ b/tensorflow_datasets/vision_language/laion400m/laion400m.py
@@ -236,7 +236,7 @@ def _download_data(
 
       file_name_to_dl_path = dl_manager.download(file_name_to_url)
 
-      return file_name_to_dl_path  # pyrefly: ignore[bad-return]
+      return file_name_to_dl_path
 
   def _split_generators(self, dl_manager: tfds.download.DownloadManager):
     """Returns SplitGenerators."""
diff --git a/tensorflow_datasets/vision_language/wit_kaggle/wit_kaggle.py b/tensorflow_datasets/vision_language/wit_kaggle/wit_kaggle.py
index d43b35b82af..9269ed30c39 100644
--- a/tensorflow_datasets/vision_language/wit_kaggle/wit_kaggle.py
+++ b/tensorflow_datasets/vision_language/wit_kaggle/wit_kaggle.py
@@ -265,7 +265,6 @@ class WitKaggle(tfds.core.GeneratorBasedBuilder):
       """,
   }
 
-  # pytype: disable=wrong-keyword-args
   # In the Wikipedia - Image/Caption Matching competition, train samples are
   # associated with a rich set of metadata, while test samples only have a
   # sample_id and four image-related fields.
@@ -315,7 +314,6 @@ class WitKaggle(tfds.core.GeneratorBasedBuilder):
           }),
       ),
   ]
-  # pytype: enable=wrong-keyword-args
 
   MANUAL_DOWNLOAD_INSTRUCTIONS = """
   Depending on the config called, manual_dir should contain some of the