Heuristics to retrieve the license from Hugging Face.

PiperOrigin-RevId: 630328078

Author: marcenacp

tensorflow_datasets/core/dataset_builders/huggingface_dataset_builder.py

@@ -273,7 +273,7 @@ def _info(self) -> dataset_info_lib.DatasetInfo:
         description=self._hf_info.description,
         features=huggingface_utils.convert_hf_features(self._hf_features()),
         citation=self._hf_info.citation,
-        license=self._hf_info.license,
+        license=self._get_license(),
         supervised_keys=_extract_supervised_keys(self._hf_info),
     )
 
@@ -407,6 +407,24 @@ def _write_shards(
       )
     return shard_infos_by_split
 
+  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.
+    if 'license' in dataset_info.card_data:
+      return dataset_info.card_data['license']
+    # Third heuristic: check the tags.
+    for tag in dataset_info.tags:
+      if tag.startswith('license:'):
+        return tag[len('license:') :]
+    return None
+
 
 def builder(
     name: str, config: Optional[str] = None, **builder_kwargs

tensorflow_datasets/core/dataset_builders/huggingface_dataset_builder_test.py

@@ -30,6 +30,7 @@ def _info(self):
     return hf_datasets.DatasetInfo(
         description='description',
         citation='citation',
+        license='test-license',
         features=None,
         version='1.0.0',
     )
@@ -89,6 +90,9 @@ def mock_huggingface_dataset_builder(
       'foo/bar', 'config', other_arg='this is another arg'
   )
   login_to_hf.assert_called_once_with('SECRET_TOKEN')
+  assert builder.info.description == 'description'
+  assert builder.info.citation == 'citation'
+  assert builder.info.redistribution_info.license == 'test-license'
   yield builder