add support for config-level freeform tags in TFDS.

PiperOrigin-RevId: 523952029

Author: pierrot0

CHANGELOG.md

@@ -10,6 +10,21 @@ and this project adheres to
 
 ### Added
 
+-   [Experimental] A list of freeform text tags can now be attached to a
+    `BuilderConfig`. For example:
+    ```py
+    BUILDER_CONFIGS = [
+        tfds.core.BuilderConfig(name="foo", tags=["foo", "live"]),
+        tfds.core.BuilderConfig(name="bar", tags=["bar", "old"]),
+    ]
+    ```
+    The tags are recorded with the dataset metadata and can later be retrieved
+    using the info object:
+    ```py
+    builder.info.config_tags  # ["foo", "live"]
+    ```
+    This feature is experimental and there are no guidelines on tags format.
+
 ### Changed
 
 ### Deprecated

tensorflow_datasets/core/dataset_builder.py

@@ -83,6 +83,16 @@ class BuilderConfig:
 
   DatasetBuilder subclasses with data configuration options should subclass
   `BuilderConfig` and add their own properties.
+
+  Attributes:
+    name: The name of the config.
+    version: The version of the config.
+    release_notes: A dictionary associating versions to changes.
+    supported_versions: A list of versions which this Builder Config supports.
+    description: a human description of the config.
+    tags: [Experimental] a list of freeform tags applying to the config. This is
+      not used by TFDS, but can be retrieved later from a ConfigBuilder
+      instance.
   """
 
   # TODO(py3.10): Should update dataclass to be:
@@ -96,6 +106,7 @@ class BuilderConfig:
       default_factory=list
   )
   description: Optional[str] = None
+  tags: List[str] = dataclasses.field(default_factory=list)
 
   @classmethod
   def from_dataset_info(
@@ -108,6 +119,7 @@ def from_dataset_info(
         description=info_proto.config_description,
         version=info_proto.version,
         release_notes=info_proto.release_notes or {},
+        tags=info_proto.config_tags or [],
     )
 
 

tensorflow_datasets/core/dataset_builder_test.py

@@ -57,13 +57,15 @@ class DummyDatasetWithConfigs(dataset_builder.GeneratorBasedBuilder):
           name="plus1",
           version=utils.Version("0.0.1"),
           description="Add 1 to the records",
+          tags=["foo:bar"],
           increment=1,
       ),
       DummyBuilderConfig(
           name="plus2",
           version=utils.Version("0.0.2"),
           supported_versions=[utils.Version("0.0.1")],
           description="Add 2 to the records",
+          tags=["foo:baz"],
           increment=2,
       ),
   ]

tensorflow_datasets/core/dataset_info.py

@@ -39,7 +39,7 @@
 import posixpath
 import tempfile
 import time
-from typing import Any, Dict, Iterable, Optional, Tuple, Union
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
 
 from absl import logging
 from etils import epath
@@ -106,23 +106,27 @@ class DatasetIdentity:
   module_name: str
   config_name: Optional[str] = None
   config_description: Optional[str] = None
+  config_tags: Optional[List[str]] = None
   release_notes: Optional[Dict[str, str]] = None
 
   @classmethod
   def from_builder(cls, builder) -> "DatasetIdentity":
     if builder.builder_config:
       config_name = builder.builder_config.name
       config_description = builder.builder_config.description
+      config_tags = builder.builder_config.tags
     else:
       config_name = None
       config_description = None
+      config_tags = None
     return cls(
         name=builder.name,
         version=utils.Version(builder.version),
         data_dir=builder.data_dir,
         module_name=str(builder.__module__),
         config_name=config_name,
         config_description=config_description,
+        config_tags=config_tags,
         release_notes=builder.release_notes,
     )
 
@@ -139,6 +143,7 @@ def from_proto(
         module_name=info_proto.module_name,
         config_name=info_proto.config_name,
         config_description=info_proto.config_description,
+        config_tags=info_proto.config_tags or [],
         release_notes={k: v for k, v in info_proto.release_notes.items()},
     )
 
@@ -228,6 +233,7 @@ def __init__(
         disable_shuffling=disable_shuffling,
         config_name=self._identity.config_name,
         config_description=self._identity.config_description,
+        config_tags=self._identity.config_tags,
         citation=utils.dedent(citation),
         module_name=self._identity.module_name,
         redistribution_info=dataset_info_pb2.RedistributionInfo(
@@ -320,6 +326,10 @@ def config_name(self) -> str:
   def config_description(self) -> str:
     return self._identity.config_description
 
+  @property
+  def config_tags(self) -> List[str]:
+    return self._identity.config_tags
+
   @property
   def full_name(self):
     """Full canonical name: (<dataset_name>/<config_name>/<version>)."""
@@ -650,6 +660,8 @@ def read_from_directory(self, dataset_info_dir: epath.PathLike) -> None:
       # Otherwise, we restore the dataset_info.json value
       if field.type == field.TYPE_MESSAGE:
         field_value.MergeFrom(field_value_restored)
+      elif field.label == field.LABEL_REPEATED:
+        field_value.extend(field_value_restored)
       else:
         setattr(self._info_proto, field_name, field_value_restored)
 
@@ -754,6 +766,11 @@ def __repr__(self):
     else:
       config_description = SKIP
 
+    if self._info_proto.config_tags:
+      config_tags = ", ".join(self.config_tags)
+    else:
+      config_tags = SKIP
+
     file_format_str = (
         self.file_format.value
         if self.file_format
@@ -765,6 +782,7 @@ def __repr__(self):
         ("full_name", repr(self.full_name)),
         ("description", _indent(f'"""\n{self.description}\n"""')),
         ("config_description", config_description),
+        ("config_tags", config_tags),
         ("homepage", repr(self.homepage)),
         ("data_path", repr(self.data_dir)),
         ("file_format", file_format_str),

tensorflow_datasets/core/dataset_info_test.py

@@ -198,6 +198,7 @@ def test_restore_after_modification(self):
         citation="some citation",
         license="some license",
     )
+    info.as_proto.config_tags.extend(["foo", "bar"])
     info.download_size = 456
     filepath_template = "{DATASET}-{SPLIT}.{FILEFORMAT}-{SHARD_X_OF_Y}"
     info.as_proto.splits.add(

tensorflow_datasets/core/proto/dataset_info.proto

@@ -156,6 +156,9 @@ message DatasetInfo {
   // Config description of the dataset
   string config_description = 14;
 
+  // Config tags of the dataset
+  repeated string config_tags = 21;
+
   // The structure and characteristics of the features of this dataset.
   Feature features = 19;
 
@@ -203,5 +206,5 @@ message DatasetInfo {
   // The data that was used to generate this dataset.
   repeated DataSourceAccess data_source_accesses = 20;
 
-  // Next available: 21
+  // Next available: 22
 }

tensorflow_datasets/core/proto/dataset_info_generated_pb2.py

@@ -63,12 +63,13 @@
     b' \x01(\t\x12\x0f\n\x07version\x18\x03'
     b' \x01(\t\x12\x10\n\x08\x64\x61ta_dir\x18\x04'
     b' \x01(\t\x12\x14\n\x0c\x64s_namespace\x18\x05'
-    b' \x01(\t"\x9f\x07\n\x0b\x44\x61tasetInfo\x12\x0c\n\x04name\x18\x01'
+    b' \x01(\t"\xb4\x07\n\x0b\x44\x61tasetInfo\x12\x0c\n\x04name\x18\x01'
     b' \x01(\t\x12\x13\n\x0b\x64\x65scription\x18\x02'
     b' \x01(\t\x12\x0f\n\x07version\x18\t \x01(\t\x12I\n\rrelease_notes\x18\x12'
     b' \x03(\x0b\x32\x32.tensorflow_datasets.DatasetInfo.ReleaseNotesEntry\x12\x13\n\x0b\x63onfig_name\x18\r'
     b' \x01(\t\x12\x1a\n\x12\x63onfig_description\x18\x0e'
-    b' \x01(\t\x12.\n\x08\x66\x65\x61tures\x18\x13'
+    b' \x01(\t\x12\x13\n\x0b\x63onfig_tags\x18\x15'
+    b' \x03(\t\x12.\n\x08\x66\x65\x61tures\x18\x13'
     b' \x01(\x0b\x32\x1c.tensorflow_datasets.Feature\x12\x10\n\x08\x63itation\x18\x03'
     b' \x01(\t\x12\x19\n\rsize_in_bytes\x18\x04'
     b' \x01(\x03\x42\x02\x18\x01\x12\x15\n\rdownload_size\x18\x0c'
@@ -138,9 +139,9 @@
   _TFDSDATASETREFERENCE._serialized_start = 1203
   _TFDSDATASETREFERENCE._serialized_end = 1312
   _DATASETINFO._serialized_start = 1315
-  _DATASETINFO._serialized_end = 2242
-  _DATASETINFO_RELEASENOTESENTRY._serialized_start = 2133
-  _DATASETINFO_RELEASENOTESENTRY._serialized_end = 2184
-  _DATASETINFO_DOWNLOADCHECKSUMSENTRY._serialized_start = 2186
-  _DATASETINFO_DOWNLOADCHECKSUMSENTRY._serialized_end = 2242
+  _DATASETINFO._serialized_end = 2263
+  _DATASETINFO_RELEASENOTESENTRY._serialized_start = 2154
+  _DATASETINFO_RELEASENOTESENTRY._serialized_end = 2205
+  _DATASETINFO_DOWNLOADCHECKSUMSENTRY._serialized_start = 2207
+  _DATASETINFO_DOWNLOADCHECKSUMSENTRY._serialized_end = 2263
 # @@protoc_insertion_point(module_scope)