Retrieve description/citation metadata from 1) Hugging Face Hub (allegedly source of truth among the Hugging Face API), then 2) Hugging Face Datasets.

PiperOrigin-RevId: 643005946

Author: marcenacp

tensorflow_datasets/core/dataset_builders/huggingface_dataset_builder.py

@@ -42,14 +42,14 @@
 from tensorflow_datasets.core import example_serializer
 from tensorflow_datasets.core import features as feature_lib
 from tensorflow_datasets.core import file_adapters
-from tensorflow_datasets.core import lazy_imports_lib
 from tensorflow_datasets.core import split_builder as split_builder_lib
 from tensorflow_datasets.core import splits as splits_lib
 from tensorflow_datasets.core.utils import huggingface_utils
 from tensorflow_datasets.core.utils import shard_utils
 from tensorflow_datasets.core.utils import tqdm_utils
 from tensorflow_datasets.core.utils import version as version_lib
 from tensorflow_datasets.core.utils.lazy_imports_utils import datasets as hf_datasets
+from tensorflow_datasets.core.utils.lazy_imports_utils import huggingface_hub
 
 
 def _extract_supervised_keys(hf_info):
@@ -224,21 +224,21 @@ def __init__(
         or '1.0.0'
     )
     self.VERSION = version_lib.Version(version)  # pylint: disable=invalid-name
-    if self._hf_config:
-      self._converted_builder_config = dataset_builder.BuilderConfig(
-          name=tfds_config,
-          version=self.VERSION,
-          description=self._hf_info.description,
-      )
-    else:
-      self._converted_builder_config = None
     self.name = huggingface_utils.convert_hf_name(hf_repo_id)
     self._hf_hub_token = hf_hub_token
     self._hf_num_proc = hf_num_proc
     self._tfds_num_proc = tfds_num_proc
     self._verification_mode = (
         'no_checks' if ignore_verifications else 'all_checks'
     )
+    if self._hf_config:
+      self._converted_builder_config = dataset_builder.BuilderConfig(
+          name=tfds_config,
+          version=self.VERSION,
+          description=self._get_text_field('description'),
+      )
+    else:
+      self._converted_builder_config = None
     super().__init__(
         file_format=file_format, config=tfds_config, data_dir=data_dir
     )
@@ -266,8 +266,16 @@ def _hf_download_and_prepare(self):
 
   @property
   def _hf_info(self) -> hf_datasets.DatasetInfo:
+    """Retrieves the dataset info from the HuggingFace Datasets."""
     return self._hf_builder.info
 
+  @functools.cached_property
+  def _hf_hub_info(self) -> huggingface_hub.hf_api.DatasetInfo:
+    """Retrieves the dataset info from the HuggingFace Hub and caches it."""
+    return huggingface_hub.dataset_info(
+        self._hf_repo_id, token=self._hf_hub_token
+    )
+
   def _hf_features(self) -> hf_datasets.Features:
     if not self._hf_info.features:
       # We need to download and prepare the data to know its features.
@@ -278,9 +286,9 @@ def _hf_features(self) -> hf_datasets.Features:
   def _info(self) -> dataset_info_lib.DatasetInfo:
     return dataset_info_lib.DatasetInfo(
         builder=self,
-        description=self._hf_info.description,
+        description=self._get_text_field('description'),
         features=huggingface_utils.convert_hf_features(self._hf_features()),
-        citation=self._hf_info.citation,
+        citation=self._get_text_field('citation'),
         license=self._get_license(),
         supervised_keys=_extract_supervised_keys(self._hf_info),
     )
@@ -417,24 +425,32 @@ def _write_shards(
 
   def _get_license(self) -> str | None:
     """Implements heuristics to get the license from HuggingFace."""
-    # First heuristic: check the DatasetInfo from Hugging Face datasets.
-    if self._hf_info.license:
-      return self._hf_info.license
-    huggingface_hub = lazy_imports_lib.lazy_imports.huggingface_hub
-    # Retrieve the dataset info from the HuggingFace Hub.
-    repo_id, token = self._hf_repo_id, self._hf_hub_token
-    dataset_info = huggingface_hub.dataset_info(repo_id, token=token)
-    # Second heuristic: check the card data.
+    # Heuristic #1: check the DatasetInfo from Hugging Face Hub/Datasets.
+    if info_license := self._get_text_field('license'):
+      return info_license
+    dataset_info = self._hf_hub_info
+    # Heuristic #2: check the card data.
     if dataset_info.card_data:
       if card_data_license := dataset_info.card_data.get('license'):
         return card_data_license
-    # Third heuristic: check the tags.
+    # Heuristic #3: check the tags.
     if dataset_info.tags:
       for tag in dataset_info.tags:
         if tag.startswith('license:'):
           return tag.removeprefix('license:')
     return None
 
+  def _get_text_field(self, field: str) -> str | None:
+    """Get the field from either HF Hub or HF Datasets."""
+    # The information retrieved from the Hub has priority over the one in the
+    # builder, because the Hub which is allegedly the new source of truth.
+    for dataset_info in [self._hf_hub_info, self._hf_info]:
+      # `description` and `citation` are not official fields in the Hugging Face
+      # Hub API but they're still exposed in its __dict__.
+      if value := getattr(dataset_info, field, None):
+        return value
+    return None
+
 
 def builder(
     name: str, config: Optional[str] = None, **builder_kwargs
@@ -449,5 +465,4 @@ def login_to_hf(hf_hub_token: Optional[str] = None):
   """Logs in to Hugging Face Hub with the token as arg or env variable."""
   hf_hub_token = hf_hub_token or os.environ.get('HUGGING_FACE_HUB_TOKEN')
   if hf_hub_token is not None:
-    huggingface_hub = lazy_imports_lib.lazy_imports.huggingface_hub
     huggingface_hub.login(token=hf_hub_token)

tensorflow_datasets/core/dataset_builders/huggingface_dataset_builder_test.py

@@ -20,6 +20,7 @@
 import pytest
 from tensorflow_datasets.core import lazy_imports_lib
 from tensorflow_datasets.core.dataset_builders import huggingface_dataset_builder
+from tensorflow_datasets.core.utils.lazy_imports_utils import huggingface_hub
 
 PIL_Image = lazy_imports_lib.lazy_imports.PIL_Image
 
@@ -72,6 +73,22 @@ def mock_login_to_hf():
     yield login_to_hf
 
 
+@pytest.fixture(autouse=True)
+def mock_hub_dataset_info():
+  fake_dataset_info = huggingface_hub.hf_api.DatasetInfo(
+      id='foo/bar',
+      citation='citation from the hub',
+      private=False,
+      downloads=123,
+      likes=456,
+      tags=[],
+  )
+  with mock.patch.object(
+      huggingface_hub, 'dataset_info', return_value=fake_dataset_info
+  ) as dataset_info:
+    yield dataset_info
+
+
 @pytest.fixture(name='builder')
 def mock_huggingface_dataset_builder(
     tmp_path, load_dataset_builder, login_to_hf
@@ -91,7 +108,7 @@ def mock_huggingface_dataset_builder(
   )
   login_to_hf.assert_called_once_with('SECRET_TOKEN')
   assert builder.info.description == 'description'
-  assert builder.info.citation == 'citation'
+  assert builder.info.citation == 'citation from the hub'
   assert builder.info.redistribution_info.license == 'test-license'
   yield builder
 

tensorflow_datasets/core/lazy_imports_lib.py

@@ -229,11 +229,6 @@ def zarr(cls):
   def conllu(cls):
     return _try_import("conllu")
 
-  @utils.classproperty
-  @classmethod
-  def huggingface_hub(cls):
-    return _try_import("huggingface_hub")
-
 
 lazy_imports = LazyImporter  # pylint: disable=invalid-name
 

tensorflow_datasets/core/utils/lazy_imports_utils.py

@@ -129,6 +129,7 @@ def datasets_error_callback(module_name: Exception):
 
 with epy.lazy_imports(error_callback=datasets_error_callback):
   import datasets  # pytype: disable=import-error
+  import huggingface_hub  # pytype: disable=import-error
 
 with epy.lazy_imports(error_callback=array_record_error_callback):
   from array_record.python import array_record_data_source