Expose the subplits num_examples and instructions to the public API

PiperOrigin-RevId: 292380707

Author: Conchylicultor

tensorflow_datasets/core/__init__.py

@@ -39,6 +39,7 @@
 from tensorflow_datasets.core.splits import SplitDict
 from tensorflow_datasets.core.splits import SplitGenerator
 from tensorflow_datasets.core.splits import SplitInfo
+from tensorflow_datasets.core.splits import SubSplitInfo
 
 from tensorflow_datasets.core.tfrecords_reader import ReadInstruction
 

tensorflow_datasets/core/dataset_builder.py

@@ -870,7 +870,7 @@ def _download_and_prepare(self, dl_manager, **prepare_split_kwargs):
       tf.io.gfile.makedirs(self._data_dir)
 
     # Generating data for all splits
-    split_dict = splits_lib.SplitDict()
+    split_dict = splits_lib.SplitDict(dataset_name=self.name)
     split_generators_kwargs = self._make_split_generators_kwargs(
         prepare_split_kwargs)
     for split_generator in self._split_generators(

tensorflow_datasets/core/dataset_info.py

@@ -151,7 +151,7 @@ def __init__(self,
             "the top-level. Got {}".format(features))
       features._set_top_level()  # pylint: disable=protected-access
     self._features = features
-    self._splits = splits_lib.SplitDict()
+    self._splits = splits_lib.SplitDict(self._builder.name)
     if supervised_keys is not None:
       assert isinstance(supervised_keys, tuple)
       assert len(supervised_keys) == 2
@@ -203,6 +203,10 @@ def homepage(self):
   def citation(self):
     return self.as_proto.citation
 
+  @property
+  def data_dir(self):
+    return self._builder.data_dir
+
   @property
   def size_in_bytes(self):
     size_in_bytes = sum(split.num_bytes for split in self.splits.values())
@@ -362,7 +366,8 @@ def read_from_directory(self, dataset_info_dir):
     parsed_proto = read_from_json(json_filename)
 
     # Update splits
-    self._set_splits(splits_lib.SplitDict.from_proto(parsed_proto.splits))
+    split_dict = splits_lib.SplitDict.from_proto(self.name, parsed_proto.splits)
+    self._set_splits(split_dict)
 
     # Restore the feature metadata (vocabulary, labels names,...)
     if self.features:

tensorflow_datasets/core/naming.py

@@ -67,11 +67,21 @@ def filepattern_for_dataset_split(dataset_name, split, data_dir,
   return "%s*" % filepath
 
 
-def filepaths_for_dataset_split(dataset_name, split, num_shards, data_dir,
-                                filetype_suffix=None):
+def filenames_for_dataset_split(
+    dataset_name, split, num_shards, filetype_suffix=None):
   prefix = filename_prefix_for_split(dataset_name, split)
   if filetype_suffix:
     prefix += ".%s" % filetype_suffix
-  filenames = sharded_filenames(prefix, num_shards)
+  return sharded_filenames(prefix, num_shards)
+
+
+def filepaths_for_dataset_split(dataset_name, split, num_shards, data_dir,
+                                filetype_suffix=None):
+  filenames = filenames_for_dataset_split(
+      dataset_name=dataset_name,
+      split=split,
+      num_shards=num_shards,
+      filetype_suffix=filetype_suffix,
+  )
   filepaths = [os.path.join(data_dir, fname) for fname in filenames]
   return filepaths

tensorflow_datasets/core/naming_test.py

@@ -64,6 +64,15 @@ def test_filename_prefix_for_split(self, prefix, expected):
     split = splits.Split.TRAIN
     self.assertEqual(expected, naming.filename_prefix_for_split(prefix, split))
 
+  def test_filenames_for_dataset_split(self):
+    self.assertEqual([
+        "foo-train-00000-of-00002",
+        "foo-train-00001-of-00002",
+    ], naming.filenames_for_dataset_split(
+        dataset_name="foo",
+        split=splits.Split.TRAIN,
+        num_shards=2))
+
   def test_filepaths_for_dataset_split(self):
     self.assertEqual([
         "/tmp/bar/foo-train-00000-of-00002",

tensorflow_datasets/core/splits.py

@@ -27,6 +27,7 @@
 from six.moves import range  # pylint: disable=redefined-builtin
 
 from tensorflow_datasets.core import proto
+from tensorflow_datasets.core import tfrecords_reader
 from tensorflow_datasets.core import utils
 
 
@@ -45,6 +46,37 @@ def __repr__(self):
     return "<tfds.core.SplitInfo num_examples=%s>" % str(num_examples)
 
 
+class SubSplitInfo(object):
+  """Wrapper around a sub split info.
+
+  This class expose info on the subsplit:
+
+  ```
+  ds, info = tfds.load(..., split='train[75%:]', with_info=True)
+  info.splits['train[75%:]'].num_examples
+  ```
+
+  """
+
+  def __init__(self, file_instructions):
+    """Constructor.
+
+    Args:
+      file_instructions: _FileInstructionOutput
+    """
+    self._file_instructions = file_instructions
+
+  @property
+  def num_examples(self):
+    """Returns the number of example in the subsplit."""
+    return self._file_instructions.num_examples
+
+  @property
+  def file_instructions(self):
+    """Returns the list of dict(filename, take, skip)."""
+    return self._file_instructions.file_instructions
+
+
 @six.add_metaclass(abc.ABCMeta)
 class SplitBase(object):
   # pylint: disable=line-too-long
@@ -527,14 +559,22 @@ def compute_mask_offsets(shard_id2num_examples):
 class SplitDict(utils.NonMutableDict):
   """Split info object."""
 
-  def __init__(self):
+  def __init__(self, dataset_name):
     super(SplitDict, self).__init__(error_msg="Split {key} already present")
+    self._dataset_name = dataset_name
 
   def __getitem__(self, key):
-    if str(key) not in self:
-      raise KeyError("Invalid split %s. Available splits are: %s" % (
-          key, sorted(list(self.keys()))))
-    return super(SplitDict, self).__getitem__(str(key))
+    # 1st case: The key exists: `info.splits['train']`
+    if str(key) in self:
+      return super(SplitDict, self).__getitem__(str(key))
+    # 2nd case: Uses instructions: `info.splits['train[50%]']`
+    else:
+      instructions = tfrecords_reader.make_file_instructions(
+          name=self._dataset_name,
+          split_infos=self.values(),
+          instruction=key,
+      )
+      return SubSplitInfo(instructions)
 
   def __setitem__(self, key, value):
     raise ValueError("Cannot add elem. Use .add() instead.")
@@ -547,9 +587,9 @@ def add(self, split_info):
     super(SplitDict, self).__setitem__(split_info.name, split_info)
 
   @classmethod
-  def from_proto(cls, repeated_split_infos):
+  def from_proto(cls, dataset_name, repeated_split_infos):
     """Returns a new SplitDict initialized from the `repeated_split_infos`."""
-    split_dict = cls()
+    split_dict = cls(dataset_name)
     for split_info_proto in repeated_split_infos:
       split_info = SplitInfo()
       split_info.CopyFrom(split_info_proto)
@@ -567,7 +607,7 @@ def total_num_examples(self):
     return sum(s.num_examples for s in self.values())
 
   def copy(self):
-    return SplitDict.from_proto(self.to_proto())
+    return SplitDict.from_proto(self._dataset_name, self.to_proto())
 
 
 def check_splits_equals(splits1, splits2):

tensorflow_datasets/core/splits_test.py

@@ -90,7 +90,7 @@ class SplitsUnitTest(testing.TestCase):
   @classmethod
   def setUpClass(cls):
     super(SplitsUnitTest, cls).setUpClass()
-    cls._splits = tfds.core.SplitDict()
+    cls._splits = tfds.core.SplitDict("ds_name")
     cls._splits.add(tfds.core.SplitInfo(name="train", num_shards=10))
     cls._splits.add(tfds.core.SplitInfo(name="test", num_shards=2))
     cls._splits.add(tfds.core.SplitInfo(name="custom", num_shards=2))
@@ -270,7 +270,8 @@ def test_split_equality(self):
     self.assertNotEqual(train, train.subsplit(tfds.percent[:50]))
     self.assertNotEqual(train.subsplit(tfds.percent[:50]), train)
 
-    self.assertFalse(tfds.Split.TRAIN != "train")
+    # Explictly want to test the `!=` operator.
+    self.assertFalse(tfds.Split.TRAIN != "train")  # pylint: disable=g-generic-assert
 
   def _info(self, split):
     read_instruction = split.get_read_instruction(self._splits)
@@ -328,6 +329,7 @@ class SplitsOffsetIntegrationTest(testing.TestCase):
 
   @classmethod
   def setUpClass(cls):
+    super(SplitsOffsetIntegrationTest, cls).setUpClass()
     cls._builder = DummyDataset(
         data_dir=testing.make_tmp_dir(),
         range_train=range(0, 666),
@@ -375,6 +377,7 @@ class SplitsIntegrationTest(testing.TestCase):
 
   @classmethod
   def setUpClass(cls):
+    super(SplitsIntegrationTest, cls).setUpClass()
     cls._builder = DummyDataset(data_dir=testing.make_tmp_dir())
     cls._builder.download_and_prepare()
 
@@ -506,7 +509,7 @@ class SplitsDictTest(testing.TestCase):
 
   @property
   def split_dict(self):
-    sd = splits.SplitDict()
+    sd = splits.SplitDict("ds_name")
     sd.add(tfds.core.SplitInfo(name="train", num_shards=10))
     sd.add(tfds.core.SplitInfo(name="test", num_shards=1))
     return sd
@@ -519,10 +522,10 @@ def test_get(self):
 
   def test_from_proto(self):
     sd = splits.SplitDict.from_proto(
-        [proto.SplitInfo(name="validation", num_shards=5)])
-    self.assertTrue("validation" in sd)
-    self.assertFalse("train" in sd)
-    self.assertFalse("test" in sd)
+        "ds_name", [proto.SplitInfo(name="validation", num_shards=5)])
+    self.assertIn("validation", sd)
+    self.assertNotIn("train", sd)
+    self.assertNotIn("test", sd)
 
   def test_to_proto(self):
     sd = self.split_dict
@@ -535,26 +538,26 @@ def test_to_proto(self):
     self.assertEqual(10, sdp[1].num_shards)
 
   def test_bool(self):
-    sd = splits.SplitDict()
+    sd = splits.SplitDict("ds_name")
     self.assertFalse(sd)  # Empty split is False
     sd.add(tfds.core.SplitInfo(name="train", num_shards=10))
     self.assertTrue(sd)  # Non-empty split is True
 
   def test_check_splits_equals(self):
-    s1 = splits.SplitDict()
+    s1 = splits.SplitDict("ds_name")
     s1.add(tfds.core.SplitInfo(name="train", num_shards=10))
     s1.add(tfds.core.SplitInfo(name="test", num_shards=3))
 
-    s2 = splits.SplitDict()
+    s2 = splits.SplitDict("ds_name")
     s2.add(tfds.core.SplitInfo(name="train", num_shards=10))
     s2.add(tfds.core.SplitInfo(name="test", num_shards=3))
 
-    s3 = splits.SplitDict()
+    s3 = splits.SplitDict("ds_name")
     s3.add(tfds.core.SplitInfo(name="train", num_shards=10))
     s3.add(tfds.core.SplitInfo(name="test", num_shards=3))
     s3.add(tfds.core.SplitInfo(name="valid", num_shards=0))
 
-    s4 = splits.SplitDict()
+    s4 = splits.SplitDict("ds_name")
     s4.add(tfds.core.SplitInfo(name="train", num_shards=11))
     s4.add(tfds.core.SplitInfo(name="test", num_shards=3))
 
@@ -564,10 +567,10 @@ def test_check_splits_equals(self):
     self.assertFalse(splits.check_splits_equals(s1, s4))  # Nb of shards !=
 
   def test_split_overwrite(self):
-    s1 = splits.SplitDict()
+    s1 = splits.SplitDict("ds_name")
     s1.add(tfds.core.SplitInfo(name="train", shard_lengths=[15]))
 
-    s2 = splits.SplitDict()
+    s2 = splits.SplitDict("ds_name")
     s2.add(tfds.core.SplitInfo(name="train", shard_lengths=[15]))
 
     self.assertTrue(splits.check_splits_equals(s1, s2))
@@ -579,5 +582,45 @@ def test_split_overwrite(self):
     self.assertFalse(splits.check_splits_equals(s1, s2))
 
 
+class SplitsSubsplitTest(testing.TestCase):
+
+  @classmethod
+  def setUpClass(cls):
+    super(SplitsSubsplitTest, cls).setUpClass()
+    cls._builder = testing.DummyDatasetSharedGenerator(
+        data_dir=testing.make_tmp_dir())
+    cls._builder.download_and_prepare()
+
+  def test_sub_split_num_examples(self):
+    s = self._builder.info.splits
+    self.assertEqual(s["train[75%:]"].num_examples, 5)
+    self.assertEqual(s["train[:75%]"].num_examples, 15)
+    self.assertEqual(
+        s["train"].num_examples,
+        s["train[75%:]"].num_examples + s["train[:75%]"].num_examples,
+    )
+
+    self.assertEqual(s["test[75%:]"].num_examples, 2)
+    self.assertEqual(s["test[:75%]"].num_examples, 8)
+    self.assertEqual(
+        s["test"].num_examples,
+        s["test[75%:]"].num_examples + s["test[:75%]"].num_examples,
+    )
+
+  def test_sub_split_file_instructions(self):
+    fi = self._builder.info.splits["train[75%:]"].file_instructions
+    self.assertEqual(fi, [{
+        "filename":
+            "dummy_dataset_shared_generator-train.tfrecord-00000-of-00001",
+        "skip": 15,
+        "take": -1,
+    }])
+
+  def test_sub_split_wrong_key(self):
+    with self.assertRaisesWithPredicateMatch(
+        ValueError, "Unknown split \"unknown\""):
+      _ = self._builder.info.splits["unknown"]
+
+
 if __name__ == "__main__":
   testing.test_main()