Rename DatasetProviderBase to DatasetProvider

PiperOrigin-RevId: 479216470

Author: broken

seqio/dataset_providers.py

@@ -86,22 +86,27 @@ def splits(self) -> Sequence[str]:
   @abc.abstractmethod
   def get_dataset(
       self,
-      sequence_length: int,
-      split: str,
+      sequence_length: Optional[Mapping[str, int]] = None,
+      split: str = tfds.Split.TRAIN,
       use_cached: bool = False,
       shuffle: bool = True,
       seed: Optional[int] = None,
       shard_info: Optional[ShardInfo] = None,
-      num_epochs: int = 1,
+      num_epochs: Optional[int] = 1
   ) -> tf.data.Dataset:
     """Returns the requested tf.data.Dataset."""
     raise NotImplementedError
 
   @abc.abstractmethod
-  def num_input_examples(self, split: str) -> int:
+  def num_input_examples(self, split: str) -> Optional[int]:
     raise NotImplementedError
 
 
+class DatasetProvider(DatasetProviderBase, metaclass=abc.ABCMeta):
+  """Synonym for DatasetProviderBase."""
+  pass
+
+
 class DatasetProviderRegistry(object):
   """Base for registry of data providers.
 
@@ -110,8 +115,8 @@ class DatasetProviderRegistry(object):
   """
 
   # Class variables must be defined in subclasses.
-  _REGISTRY: MutableMapping[str, DatasetProviderBase]
-  _PROVIDER_TYPE: Type[DatasetProviderBase]
+  _REGISTRY: MutableMapping[str, DatasetProvider]
+  _PROVIDER_TYPE: Type[DatasetProvider]
 
   @classmethod
   def add_provider(cls, name: str, provider):
@@ -225,10 +230,10 @@ def get_dataset(
     ...
 
 
-class DataSource(DatasetProviderBase):
+class DataSource(DatasetProvider):
   """A `DatasetProvider` that provides raw data from an input source.
 
-  Inherits all abstract methods and properties of `DatasetProviderBase` except
+  Inherits all abstract methods and properties of `DatasetProvider` except
   those overidden below.
   """
 
@@ -264,7 +269,7 @@ def supports_arbitrary_sharding(self) -> bool:
 
   @property
   def output_features(self) -> Mapping[str, Feature]:
-    """Override unused property of `DatasetProviderBase`."""
+    """Override unused property of `DatasetProvider`."""
     raise NotImplementedError
 
   @abc.abstractmethod
@@ -275,10 +280,14 @@ def list_shards(self, split: str) -> Sequence[str]:
   @abc.abstractmethod
   def get_dataset(
       self,  # pytype: disable=signature-mismatch  # overriding-default-value-checks
-      split: str,
+      split: str = tfds.Split.TRAIN,
       shuffle: bool = True,
       seed: Optional[int] = None,
       shard_info: Optional[ShardInfo] = None,
+      *,  # remaining args are out of order from parent
+      sequence_length: Optional[Mapping[str, int]] = None,  # Unused
+      use_cached: bool = False,  # Unused
+      num_epochs: Optional[int] = 1  # Unused
   ) -> tf.data.Dataset:
     """Overrides base class to add shard identifier and remove use_cached.
 
@@ -287,6 +296,9 @@ def get_dataset(
       shuffle: bool, whether to shuffle the input source.
       seed: tf.int64 scalar tf.Tensor (or None) for shuffling input source.
       shard_info: optional specification for loading a shard of the split.
+      sequence_length: Unused
+      use_cached: Unused
+      num_epochs: Unused
     """
     raise NotImplementedError
 
@@ -376,10 +388,14 @@ def supports_arbitrary_sharding(self) -> bool:
 
   def get_dataset(
       self,
-      split: str,
+      split: str = tfds.Split.TRAIN,
       shuffle: bool = True,
       seed: Optional[int] = None,
       shard_info: Optional[ShardInfo] = None,
+      *,  # remaining args are out of order from parent
+      sequence_length: Optional[Mapping[str, int]] = None,  # Unused
+      use_cached: bool = False,  # Unused
+      num_epochs: Optional[int] = 1  # Unused
   ) -> tf.data.Dataset:
     if shard_info and shard_info.num_shards > 1:
       raise ValueError(
@@ -470,10 +486,14 @@ def supports_arbitrary_sharding(self) -> bool:
 
   def get_dataset(
       self,
-      split: str,
+      split: str = tfds.Split.TRAIN,
       shuffle: bool = True,
       seed: Optional[int] = None,
       shard_info: Optional[ShardInfo] = None,
+      *,  # remaining args are out of order from parent
+      sequence_length: Optional[Mapping[str, int]] = None,  # Unused
+      use_cached: bool = False,  # Unused
+      num_epochs: Optional[int] = 1  # Unused
   ) -> tf.data.Dataset:
     return self.tfds_dataset.load(
         split, shuffle_files=shuffle, seed=seed, shard_info=shard_info
@@ -539,10 +559,14 @@ def supports_arbitrary_sharding(self) -> bool:
 
   def get_dataset(
       self,
-      split: str,
+      split: str = tfds.Split.TRAIN,
       shuffle: bool = True,
       seed: Optional[int] = None,
       shard_info: Optional[ShardInfo] = None,
+      *,  # remaining args are out of order from parent
+      sequence_length: Optional[Mapping[str, int]] = None,  # Unused
+      use_cached: bool = False,  # Unused
+      num_epochs: Optional[int] = 1  # Unused
   ) -> tf.data.Dataset:
     files = self.list_shards(split)
 
@@ -862,7 +886,7 @@ def __call__(self, dataset):
 MetricFnCallable = metrics_lib.MetricFnCallable
 
 
-class Task(DatasetProviderBase):
+class Task(DatasetProvider):
   """A class to manage a dataset and its related metrics."""
 
   def __init__(
@@ -1260,15 +1284,15 @@ def get_cached_stats(
 
   def get_dataset(
       self,  # pytype: disable=signature-mismatch  # overriding-default-value-checks
-      sequence_length: Optional[Mapping[str, int]],
+      sequence_length: Optional[Mapping[str, int]] = None,
       split: str = tfds.Split.TRAIN,
       use_cached: bool = False,
       shuffle: bool = True,
-      shuffle_buffer_size: Optional[int] = None,
+      shuffle_buffer_size: Optional[int] = None,  # Unique to Task
       seed: Optional[int] = None,
       shard_info: Optional[ShardInfo] = None,
       num_epochs: Optional[int] = 1,
-      trim_output_features: bool = True,
+      trim_output_features: bool = True,  # Unique to Task
   ) -> tf.data.Dataset:
     """Returns a tf.data.Dataset from cache or generated on the fly.
 
@@ -1465,7 +1489,7 @@ def get(cls, name) -> Task:
 MixtureRate = Union[int, float, Callable[[Union[Task, "Mixture"]], float]]
 
 
-class Mixture(DatasetProviderBase):
+class Mixture(DatasetProvider):
   """Class for mixing multiple tasks."""
 
   def __init__(
@@ -1607,14 +1631,14 @@ def _check_compatible_features(self) -> None:
 
   def get_dataset(  # pytype: disable=signature-mismatch  # overriding-parameter-type-checks
       self,
-      sequence_length: Optional[Mapping[str, int]],
+      sequence_length: Optional[Mapping[str, int]] = None,
       split: str = tfds.Split.TRAIN,
       use_cached: bool = False,
       shuffle: bool = True,
       seed: Optional[int] = None,
       shard_info: Optional[ShardInfo] = None,
-      num_epochs: Optional[int] = None,
-      copy_pretokenized: bool = False,
+      num_epochs: Optional[int] = None,  # Unique default for Mixture
+      copy_pretokenized: bool = False,  # Unique (and all below) to Mixture
       compute_stats_empirically: bool = False,
       log_mixing_proportions: bool = True,
       passthrough_features: Optional[Sequence[str]] = None,
@@ -2030,7 +2054,7 @@ def get_dataset(
 
   mixture_or_task = (
       get_mixture_or_task(mixture_or_task_name)
-      if not isinstance(mixture_or_task_name, DatasetProviderBase)
+      if not isinstance(mixture_or_task_name, DatasetProvider)
       else mixture_or_task_name
   )
   is_grain_task = False

seqio/experimental.py

@@ -22,7 +22,7 @@
 from seqio import preprocessors as seqio_preprocessors
 from seqio import utils
 import tensorflow as tf
-
+import tensorflow_datasets as tfds
 
 CacheDatasetPlaceholder = dataset_providers.CacheDatasetPlaceholder
 Mixture = dataset_providers.Mixture
@@ -232,10 +232,14 @@ def list_shards(self, split: str) -> Sequence[str]:
 
   def get_dataset(
       self,
-      split: str,
+      split: str = tfds.Split.TRAIN,
       shuffle: bool = True,
       seed: Optional[int] = None,
       shard_info: Optional[ShardInfo] = None,
+      *,  # remaining args are out of order from parent
+      sequence_length: Optional[Mapping[str, int]] = None,  # Unused
+      use_cached: bool = False,  # Unused
+      num_epochs: Optional[int] = 1  # Unused
   ) -> tf.data.Dataset:
     shard_info: ShardInfo = shard_info or ShardInfo(0, 1)
     if self._train_split not in self._original_source.splits:

seqio/helpers.py

@@ -23,6 +23,7 @@
 from seqio import dataset_providers as dp
 from seqio import vocabularies as vc
 import tensorflow.compat.v2 as tf
+import tensorflow_datasets as tfds
 
 
 def mixture_or_task_with_new_vocab(
@@ -201,10 +202,14 @@ def list_shards(self, split: str) -> Sequence[str]:
 
   def get_dataset(
       self,
-      split: str,
+      split: str = tfds.Split.TRAIN,
       shuffle: bool = True,
       seed: Optional[int] = None,
       shard_info: Optional[dp.ShardInfo] = None,
+      *,  # remaining args are out of order from parent
+      sequence_length: Optional[Mapping[str, int]] = None,  # Unused
+      use_cached: bool = False,  # Unused
+      num_epochs: Optional[int] = 1  # Unused
   ) -> tf.data.Dataset:
     """See base class for documentation."""
     if split not in self.split_sizes: