Internal change

PiperOrigin-RevId: 788353351

Author: The TensorFlow Datasets Authors

tensorflow_datasets/core/load.py

@@ -183,6 +183,7 @@ def get_dataset_repr() -> str:
           f' explicitly set. Wrong arguments for {get_dataset_repr()}'
       )
     builder_kwargs['data_dir'] = gcs_utils.gcs_path('datasets')
+
   if name.namespace:
     if name.namespace == 'huggingface':
       return huggingface_dataset_builder.builder(
@@ -827,12 +828,8 @@ def data_source(
   """  # fmt:skip
   builder_kwargs = _set_file_format_for_data_source(data_dir, builder_kwargs)
   _validate_file_format_for_data_source(builder_kwargs)
-  dbuilder = _fetch_builder(
-      name,
-      data_dir,
-      builder_kwargs,
-      try_gcs,
-  )
+
+  dbuilder = _fetch_builder(name, data_dir, builder_kwargs, try_gcs)
   _download_and_prepare_builder(dbuilder, download, download_and_prepare_kwargs)
   return dbuilder.as_data_source(
       split=split, decoders=decoders, deserialize_method=deserialize_method

tensorflow_datasets/core/utils/file_utils.py

@@ -180,6 +180,14 @@ def get_dataset_dir(
     version: version_lib.Version | str | None = None,
 ) -> epath.Path:
   """Returns the data directory for the given dataset."""
+  # If the data dir is a TF Hub path, use the TF Hub dataset directory.
+  if tfhub_path.is_tfhub_filepath(str(data_dir)):
+    return tfhub_file_utils.get_tfhub_dataset_dir(
+        dataset_name=builder_name,
+        config_name=config_name,
+        version=version,
+        tfhub_env=tfhub_file_utils.TfHubEnv.PROD,
+    )
   dataset_dir = epath.Path(data_dir) / builder_name
   if config_name:
     dataset_dir /= config_name