text: # coding=utf-8
  # Copyright 2021 The Uncertainty Baselines Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  """Abstract base classes which defines interfaces for datasets."""
  
  import logging
  from typing import Callable, Dict, Optional, Sequence, Type, TypeVar, Union
  
  from robustness_metrics.common import ops
  from robustness_metrics.common import types
  from robustness_metrics.datasets import tfds as robustness_metrics_base
  import tensorflow.compat.v2 as tf
  import tensorflow_datasets as tfds
  
  
  # For datasets like UCI, the tf.data.Dataset returned by _read_examples will
  # have elements that are Sequence[tf.Tensor], for TFDS datasets they will be
  # Dict[Text, tf.Tensor] (types.Features), for Criteo they are a tf.Tensor.
  PreProcessFn = Callable[
      [Union[int, tf.Tensor, Sequence[tf.Tensor], types.Features]],
      types.Features]
  
  # Same as PreProcessFn except also takes an integer first argument.
  _EnumeratedPreProcessFn = Callable[
      [int, Union[int, tf.Tensor, Sequence[tf.Tensor], types.Features]],
      types.Features]
  
  
  def _absolute_split_len(absolute_split, dataset_splits):
    if absolute_split.from_ is None:
      start = 0
    else:
      start = absolute_split.from_
    if absolute_split.to is None:
      end = dataset_splits[absolute_split.splitname].num_examples
    else:
      end = absolute_split.to
    return end - start
  
  
  def get_validation_percent_split(
      dataset_builder,
      validation_percent,
      split,
      test_split=tfds.Split.TEST):
    """Calculate a validation set from a provided validation_percent in [0, 1]."""
    if validation_percent < 0.0 or validation_percent >= 1.0:
      raise ValueError(
          'validation_percent must be in [0, 1), received {}.'.format(
              validation_percent))
    if validation_percent == 0.:
      train_split = tfds.Split.TRAIN
      # We cannot use None here because that will return all the splits if passed
      # to builder.as_dataset().
      validation_split = tfds.Split.VALIDATION
    else:
      num_train_examples = dataset_builder.info.splits['train'].num_examples
      num_validation_examples = int(num_train_examples * validation_percent)
      train_split = tfds.core.ReadInstruction(
          'train', to=-num_validation_examples, unit='abs')
      validation_split = tfds.core.ReadInstruction(
          'train', from_=-num_validation_examples, unit='abs')
  
    if split in ['train', tfds.Split.TRAIN]:
      new_split = train_split
    elif split in ['validation', tfds.Split.VALIDATION]:
      new_split = validation_split
    elif split in ['test', tfds.Split.TEST]:
      new_split = test_split
    elif isinstance(split, str):
      # For Python 3 this should be save to check for the Text type, see
      # https://github.com/google/pytype/blob/a4d56ef763eb4a29b5db03a2013c3373f9f46146/pytype/pytd/builtins/2and3/typing.pytd#L31.
      raise ValueError(
          'Invalid string name for split, must be one of ["train", "validation"'
          ', "test"], received {}.'.format(split))
    else:
      new_split = split
    return new_split
  
  
  class BaseDataset(robustness_metrics_base.TFDSDataset):
    """Abstract base dataset class.
  
    Requires subclasses to override _read_examples, _create_process_example_fn.
    """
  
    def __init__(
        self,
        name: str,
        dataset_builder: tfds.core.DatasetBuilder,
        split: Union[float, str, tfds.Split],
        seed: Optional[Union[int, tf.Tensor]] = None,
        is_training: Optional[bool] = None,
        shuffle_buffer_size: Optional[int] = None,
        num_parallel_parser_calls: int = tf.data.experimental.AUTOTUNE,
        drop_remainder: bool = True,
        fingerprint_key: Optional[str] = None,
        download_data: bool = False,
        decoders: Optional[Dict[str, tfds.decode.Decoder]] = None,
        cache: bool = False):
      """Create a tf.data.Dataset builder.
  
      Args:
        name: the name of this dataset.
        dataset_builder: the TFDS dataset builder used to read examples given a
          split.
        split: a dataset split, either a custom tfds.Split or one of the
          tfds.Split enums [TRAIN, VALIDAITON, TEST] or their lowercase string
          names. For Criteo it can also be a float to represent the level of data
          augmentation. For speech commands it can be a tuple of a string and
          float for shifted data splits.
        seed: the seed used as a source of randomness.
        is_training: whether or not `split` is the training split. This is
          necessary because tf.data subsplits can sometimes be derived from the
          training split, such as when using part of the training split as a
          validation set, and this complicates calculating `is_training`
          in these cases. Only required when the passed split is not one of
          ['train', 'validation', 'test', tfds.Split.TRAIN, tfds.Split.VALIDATION,
          tfds.Split.TEST], otherwise it is set automatically.
        shuffle_buffer_size: the number of example to use in the shuffle buffer
          for tf.data.Dataset.shuffle().
        num_parallel_parser_calls: the number of parallel threads to use while
          preprocessing in tf.data.Dataset.map().
        drop_remainder: whether or not to drop the last batch of data if the
          number of points is not exactly equal to the batch size. This option
          needs to be True for running on TPUs.
        fingerprint_key: The name of the feature holding a string that will be
          used to create an element id using a fingerprinting function. If None,
          then `ds.enumerate()` is added before the `ds.map(preprocessing_fn)` is
          called and an `id` field is added to the example Dict.
        download_data: Whether or not to download data before loading.
        decoders: Optional TFDS decoders to provide to
          `dataset_builder.as_dataset`, the same as passed to `tfds.load`.
        cache: Whether or not to cache the dataset after it is returned from
          dataset_builder.as_dataset(...) (before preprocessing is applied).
      """
      self.name = name
      self._split = split
  
      # Stateless random ops require a (2,) shaped seed.
      if seed is None:
        self._seed = tf.random.uniform((2,), maxval=int(1e10), dtype=tf.int32)
      elif isinstance(seed, int):
        self._seed = (seed, seed + 1)
      elif isinstance(seed, tf.Tensor) and tf.shape(seed).shape == 0:
        self._seed = tf.stack([seed, seed+1])
      else:
        self._seed = seed
  
      self._num_parallel_parser_calls = num_parallel_parser_calls
      self._drop_remainder = drop_remainder
      self._download_data = download_data
      self._decoders = decoders
      self._cache = cache
  
      known_splits = [
          'train', 'validation', 'test', tfds.Split.TRAIN, tfds.Split.VALIDATION,
          tfds.Split.TEST
      ]
      if is_training is None:
        if split in known_splits:
          is_training = split in ['train', tfds.Split.TRAIN]
        else:
          raise ValueError(
              'Received ambiguous split {}, must set is_training for splits other '
              'than "train", "validation", "test".'.format(split))
  
      self._is_training = is_training
      # TODO(znado): properly parse the number of train/validation/test examples
      # from the provided split, see `make_file_instructions(...)` in
      # tensorflow_datasets/core/tfrecords_reader.py.
      if 'train' in dataset_builder.info.splits and shuffle_buffer_size is None:
        num_train_examples = dataset_builder.info.splits['train'].num_examples
        self._shuffle_buffer_size = num_train_examples
      else:
        self._shuffle_buffer_size = shuffle_buffer_size
      super(BaseDataset, self).__init__(
          dataset_builder=dataset_builder,
          fingerprint_key=fingerprint_key,
          split=self._split,
          label_key='label')
  
      self._enumerate_id_key = '_enumerate_added_per_step_id'
  
      self._add_fingerprint_key = False
      if self._fingerprint_key is None:
        self._fingerprint_key = '_enumerate_added_example_id'
        self._add_fingerprint_key = True
  
    # This method can be overridden to add custom info via info.metadata.
    @property
    def tfds_info(self) -> tfds.core.DatasetInfo:
      return self._dataset_builder.info
  
    @property
    def split(self):
      return self._split
  
    @property
    def num_examples(self):
      if isinstance(self._split, tfds.core.ReadInstruction):
        absolute_split = self._split.to_absolute(
            {
                name: self.tfds_info.splits[name].num_examples
                for name in self.tfds_info.splits.keys()
            })[0]
        return _absolute_split_len(absolute_split, self.tfds_info.splits)
      return self.tfds_info.splits[self._split].num_examples
  
    def _create_process_example_fn(self) -> Optional[PreProcessFn]:
      """Create a function to perform dataset pre-processing, if needed.
  
      Returns:
        None if no processing is necessary. Otherwise, a function which takes as
        inputs a single element of the dataset (passed from dataset.map()), and
        returns a dict with keys 'features' and 'labels' and their corresponding
        Tensor values.
      """
      raise NotImplementedError(
          'Must override dataset _create_process_example_fn!')
  
    def _create_process_batch_fn(self, batch_size: int) -> Optional[PreProcessFn]:
      """Create a function to perform pre-processing on batches, such as mixup.
  
      Args:
        batch_size: the size of the batch to be processed.
  
      Returns:
        None if no processing is necessary. Otherwise, a function which takes as
        inputs a batch of elements of the dataset (passed from dataset.batch()),
        and returns a dict with keys 'features' and 'labels' and their
        corresponding Tensor values.
      """
      return None
  
    def _create_element_id(self, features: types.Features) -> types.Features:
      """Hash element id for a unique id per data element (NOT per-step)."""
      if 'element_id' in features:
        raise ValueError(
            '`element_id` should not be already present in the feature set.')
      fingerprint_feature = features[self._fingerprint_key]
      features['element_id'] = ops.fingerprint_int64(fingerprint_feature)
      return features
  
    def _add_enumerate_id(self, enumerate_key: str) -> _EnumeratedPreProcessFn:
      def _add_example_id(enumerate_id, example):
        """Turn an id added by ds.enumerate() as a field in the example dict."""
        if isinstance(example, tf.Tensor):
          example = {'features': example}
        example[enumerate_key] = enumerate_id
        return example
      return _add_example_id
  
    def _load(self,
              *,
              preprocess_fn: Optional[PreProcessFn] = None,
              batch_size: int = -1) -> tf.data.Dataset:
      """Transforms the dataset from builder.as_dataset() to batch, repeat, etc.
  
      Note that we do not handle replication/sharding here, because the
      DistributionStrategy experimental_distribute_dataset() will shard the input
      files for us.
  
      Args:
        preprocess_fn: an optional preprocessing function, if not provided then a
          subclass must define _create_process_example_fn() which will be used to
          preprocess the data.
        batch_size: the batch size to use.
  
      Returns:
        A tf.data.Dataset of elements that are a dict with keys 'features' and
        'labels' and their corresponding Tensor values.
      """
      if batch_size <= 0:
        raise ValueError(
            'Must provide a positive batch size, received {}.'.format(batch_size))
  
      self._seed, self._shuffle_seed = tf.random.experimental.stateless_split(
          self._seed, num=2)
  
      if self._download_data:
        self._dataset_builder.download_and_prepare()
      dataset = self._dataset_builder.as_dataset(
          split=self._split, decoders=self._decoders)
  
      # Possibly cache the original dataset before preprocessing is applied.
      if self._cache:
        dataset = dataset.cache()
  
      # This must be done *before* repeating the dataset so that each example has
      # a unique and stable fingerprint key.
      if self._add_fingerprint_key:
        dataset = dataset.enumerate()
        add_fingerprint_key_fn = self._add_enumerate_id(self._fingerprint_key)
        dataset = dataset.map(
            add_fingerprint_key_fn,
            num_parallel_calls=self._num_parallel_parser_calls)
  
      # Shuffle and repeat only for the training split.
      if self._is_training:
        dataset = dataset.shuffle(
            self._shuffle_buffer_size,
            seed=tf.cast(self._shuffle_seed[0], tf.int64),
            reshuffle_each_iteration=True)
        dataset = dataset.repeat()
  
      # Enumerate the dataset to generate a unique per-example, per-step id, that
      # is then added to the feature dict as `self._enumerate_id_key`.
      # Note that this is distinct from just a per-example id that is used by
      # Robustness Metrics to identify examples, because we want an id that is
      # different for each step so that we can fold it into a source of randomness
      # for deterministic random preprocessing.
      # This must be done *after* repeating the dataset so that each example has a
      # different key per-step.
      dataset = dataset.enumerate()
      add_per_step_id_key_fn = self._add_enumerate_id(self._enumerate_id_key)
  
      if preprocess_fn is None:
        preprocess_fn = self._create_process_example_fn()
  
      # `self._create_element_id` must come before `preprocess_fn` so that we
      # guarantee the field with key `self._fingerprint_key` is still present
      # (many preprocess_fn's may not return it).
      preprocess_fn = ops.compose(
          add_per_step_id_key_fn, self._create_element_id, preprocess_fn)
      dataset = dataset.map(
          preprocess_fn,
          num_parallel_calls=self._num_parallel_parser_calls)
  
      # Note that unless the default value of `drop_remainder=True` is overriden
      # in `__init__`, we always drop the last batch when the batch size does not
      # evenly divide the number of examples.
      # TODO(znado): add padding to last partial eval batch.
      dataset = dataset.batch(batch_size, drop_remainder=self._drop_remainder)
  
      process_batch_fn = self._create_process_batch_fn(batch_size)  # pylint: disable=assignment-from-none
      if process_batch_fn:
        dataset = dataset.map(
            process_batch_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE)
  
      if not self._is_training:
        dataset = dataset.cache()
  
      dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
  
      # The AutoSharding policy in DistributionStrategy defaults to AUTO, which
      # will fallback to DATA if it can, which is safe to do but possibly
      # wasteful compared to `distribute_datasets_from_function`.
      options = tf.data.Options()
      # Optimize dataset performance.
      # Keep commented out, unclear if will always improve performance.
      # options.experimental_optimization.parallel_batch = True
      options.experimental_optimization.map_fusion = True
      options.experimental_optimization.map_parallelization = True
      options.experimental_threading.private_threadpool_size = 48
      options.experimental_threading.max_intra_op_parallelism = 1
      dataset = dataset.with_options(options)
      return dataset
  
    def load(
        self,
        *,
        preprocess_fn: Optional[PreProcessFn] = None,
        batch_size: int = -1,
        strategy: Optional[tf.distribute.Strategy] = None) -> tf.data.Dataset:
      """Function definition to support multi-host dataset sharding.
  
      This is preferred over strategy.experimental_distribute_dataset(...) because
      not all datasets will have enough input files to have >=1 per host, which
      will result in an error unless the auto_shard_policy is set to
      tf.data.experimental.AutoShardPolicy.OFF. However, if auto sharding is OFF,
      then each host will process the same set of files, in the same order, which
      will be the same as using a single host. To correctly distribute across
      multiple hosts, we must either shard the input files across hosts, or
      shuffle the data in a different order on each host. In order to get a
      per-host id to fold into the shuffle RNG, we use
      strategy.distribute_datasets_from_function to get a
      tf.distribute.InputContext. This is preferrred over AutoShardPolicy.DATA
      because DATA will prepare and throw out (n - 1)/n elements of data, where n
      is the number of devices (DATA also relies on the data files being in the
      same ordering across hosts, which may be a fair assumption). On an ImageNet
      test run on a TPUv2-32, we saw a 37% slowdown using DATA instead of
      `distribute_datasets_from_function`. See this documentation for more info
      https://www.tensorflow.org/tutorials/distribute/input#sharding.
  
  
      Args:
        preprocess_fn: see `load()`.
        batch_size: the *global* batch size to use. This should equal
          `per_replica_batch_size * num_replica_in_sync`.
        strategy: the DistributionStrategy used to shard the dataset. Note that
          this is only required if TensorFlow for training, otherwise it can be
          ignored.
      Returns:
        A sharded dataset, with its seed combined with the per-host id.
      """
      if strategy:
        def _load_distributed(ctx: tf.distribute.InputContext):
          self._seed = tf.random.experimental.stateless_fold_in(
              self._seed, ctx.input_pipeline_id)
          per_replica_batch_size = ctx.get_per_replica_batch_size(batch_size)
          return self._load(
              preprocess_fn=preprocess_fn, batch_size=per_replica_batch_size)
  
        return strategy.distribute_datasets_from_function(_load_distributed)
      else:
        return self._load(preprocess_fn=preprocess_fn, batch_size=batch_size)
  
  
  _BaseDatasetClass = Type[TypeVar('B', bound=BaseDataset)]
  
  
  def make_ood_dataset(ood_dataset_cls: _BaseDatasetClass) -> _BaseDatasetClass:
    """Generate a BaseDataset with in/out distribution labels."""
  
    class _OodBaseDataset(ood_dataset_cls):
      """Combine two datasets to form one with in/out of distribution labels."""
  
      def __init__(
          self,
          in_distribution_dataset: BaseDataset,
          shuffle_datasets: bool = False,
          **kwargs):
        super(_OodBaseDataset, self).__init__(**kwargs)
        # This should be the builder for whatever split will be considered
        # in-distribution (usually the test split).
        self._in_distribution_dataset = in_distribution_dataset
        self._shuffle_datasets = shuffle_datasets
  
      def load(self,
               *,
               preprocess_fn=None,
               batch_size: int = -1) -> tf.data.Dataset:
        # Set up the in-distribution dataset using the provided dataset builder.
        if preprocess_fn:
          dataset_preprocess_fn = preprocess_fn
        else:
          dataset_preprocess_fn = (
              self._in_distribution_dataset._create_process_example_fn())  # pylint: disable=protected-access
        dataset_preprocess_fn = ops.compose(
            dataset_preprocess_fn,
            _create_ood_label_fn(True))
        dataset = self._in_distribution_dataset.load(
            preprocess_fn=dataset_preprocess_fn,
            batch_size=batch_size)
        dataset = dataset.map(
            _remove_fingerprint_id_key(self._in_distribution_dataset))
  
        # Set up the OOD dataset using this class.
        if preprocess_fn:
          ood_dataset_preprocess_fn = preprocess_fn
        else:
          ood_dataset_preprocess_fn = (
              super(_OodBaseDataset, self)._create_process_example_fn())
        ood_dataset_preprocess_fn = ops.compose(
            ood_dataset_preprocess_fn,
            _create_ood_label_fn(False))
        ood_dataset = super(_OodBaseDataset, self).load(
            preprocess_fn=ood_dataset_preprocess_fn,
            batch_size=batch_size)
        ood_dataset = ood_dataset.map(_remove_fingerprint_id_key(self))
  
        # Combine the two datasets.
        try:
          combined_dataset = dataset.concatenate(ood_dataset)
        except TypeError:
          logging.info(
              'Two datasets have different types, concat feature and label only')
  
          def clean_keys(example):
            # only keep features and labels, remove the rest
            return {
                'features': example['features'],
                'labels': example['labels'],
                'is_in_distribution': example['is_in_distribution']
            }
  
          combined_dataset = dataset.map(clean_keys).concatenate(
              ood_dataset.map(clean_keys))
        if self._shuffle_datasets:
          combined_dataset = combined_dataset.shuffle(self._shuffle_buffer_size)
        return combined_dataset
  
      @property
      def num_examples(self):
        return (
            self._in_distribution_dataset.num_examples +
            super(_OodBaseDataset, self).num_examples)
  
    return _OodBaseDataset
  
  
  # We may be able to refactor this to be part of preprocess_fn so we don't
  # need to do a second `ds.map()`, but the ordering of composed functions is
  # tricky.
  def _remove_fingerprint_id_key(dataset):
  
    def f(example: types.Features) -> types.Features:
      del example[dataset._fingerprint_key]  # pylint: disable=protected-access
      return example
  
    return f
  
  
  def _create_ood_label_fn(is_in_distribution: bool) -> PreProcessFn:
    """Returns a function that adds an `is_in_distribution` key to examles."""
  
    def _add_ood_label(example: types.Features) -> types.Features:
      if is_in_distribution:
        in_dist_label = tf.ones_like(example['labels'], tf.int32)
      else:
        in_dist_label = tf.zeros_like(example['labels'], tf.int32)
      example['is_in_distribution'] = in_dist_label
      return example
  
    return _add_ood_label
  