class_docstr: MultiProblem base class.
class_docstr_tok: ["MultiProblem","base","class","."]
code: def dataset(
      self,
      mode,
      data_dir=None,
      num_threads=None,
      output_buffer_size=None,
      shuffle_files=None,
      hparams=None,
      preprocess=True,
      dataset_split=None,
      shard=None,
      partition_id=0,
      num_partitions=1,
      shuffle_buffer_size=1024,
      max_records=-1,
  ):
      # A list of datasets corresponding to the tasks in the task_list object
      # that need to be mixed.
      datasets = []
      is_training = mode == tf_estimator.ModeKeys.TRAIN
      is_infer = mode == tf_estimator.ModeKeys.PREDICT
      enc = self.task_list[0].feature_encoders(data_dir=data_dir)[
          "targets"
      ]
      self.update_task_ids(enc.vocab_size)
  
      for task in self.task_list:
          task_dataset = task.dataset(
              mode=mode,
              data_dir=data_dir,
              num_threads=num_threads,
              output_buffer_size=output_buffer_size,
              shuffle_files=shuffle_files,
              hparams=hparams,
              preprocess=preprocess,
              dataset_split=dataset_split,
              shard=shard,
              partition_id=partition_id,
              num_partitions=num_partitions,
              shuffle_buffer_size=shuffle_buffer_size,
              max_records=max_records,
          )
  
          if is_training:
              task_dataset = task_dataset.repeat()
  
          # pylint: disable=cell-var-from-loop
          task_dataset = task_dataset.map(
              lambda x: self.normalize_example(
                  task, x, enc, hparams, is_infer
              )
          )
          # pylint: enable=cell-var-from-loop
  
          # To run evaluation, we want to zip datasets from different tasks,
          # but zipping will cut off at the shortest dataset in tf.Datasets.
          # For this reason, we add zero padding to the shorter datasets as
          # it will be ignored in metrics but it provides space for larger data.
          if not is_training and not is_infer:
              zeros = tf.zeros(
                  [self._ADDED_EVAL_COUNT, 1], dtype=tf.int64
              )
              pad_data = tf.data.Dataset.from_tensor_slices(
                  {
                      "targets": zeros,
                      "batch_prediction_key": zeros,
                      "task_id": zeros,
                  }
              )
              task_dataset = task_dataset.concatenate(pad_data)
  
          datasets.append(task_dataset)
  
      # Setup the problem hparams by setting them to the LM task hparams.
      self.get_hparams(model_hparams=hparams)
  
      if is_training:
          # Using tf.Variable instead of get_variable to work around issues with
          # queues on multiple hosts. Note that this will separately count steps
          # on each host that's feeding the data, so in a large-scale setting you
          # may need to adjust hparams for that. For example, a 4x4 slice of a TPU
          # pod may use 2 data hosts, so we'll be only adding 1 here once for 2
          # examples -- divide the corresponding hparams by 2 to compensate.
          problem_step = tf.Variable(
              tf.constant(0, dtype=tf.int64),
              trainable=False,
              use_resource=True,
              dtype=tf.int64,
              name="problem_step",
          )
          dataset_iterators = [
              d.make_one_shot_iterator() for d in datasets
          ]
  
          def get_next_from_dataset(dataset_iter):
              return dataset_iter.get_next()
  
          def get_exp_sched_prob():
              """Inverse decay exponential to mix datasets."""
              with tf.control_dependencies(
                  [problem_step.assign_add(1)]
              ):
                  inv_exp_decay = common_layers.inverse_exp_decay(
                      max_step=hparams.multiproblem_schedule_max_examples,
                      min_value=1e-4,
                      step=tf.to_float(problem_step),
                  )
                  # inv_exp_decay is bounded above by 1.0
                  return (
                      inv_exp_decay
                      * hparams.multiproblem_schedule_threshold
                  )
  
          def get_const_sched_prob():
              return hparams.multiproblem_schedule_threshold
  
          def get_pretrain_sched_prob():
              """Pretrain the primary tasks for max examples."""
              with tf.control_dependencies(
                  [problem_step.assign_add(1)]
              ):
                  return tf.cond(
                      tf.greater(
                          problem_step,
                          tf.cast(
                              hparams.multiproblem_schedule_max_examples,
                              dtype=tf.int64,
                          ),
                      ),
                      lambda: 1.0,
                      lambda: 0.0,
                  )
  
          def mix_data(example):
              """Function to mix the different datasets according to a schedule."""
              del example
              # This block computes the probability of mixing the primary task with
              # the secondary tasks. 0 = only the primary task, 1 = only the secondary
              # tasks.
              if (
                  hparams.multiproblem_mixing_schedule
                  == MixingSchedule.EXPONENTIAL
              ):
                  prob = get_exp_sched_prob()
                  prob = tf.cond(
                      tf.equal(
                          tf.floormod(
                              problem_step, tf.cast(5e6, dtype=tf.int64)
                          ),
                          0,
                      ),
                      lambda: tf.Print(
                          prob, [prob], message="Probability"
                      ),
                      lambda: prob,
                  )
              elif (
                  hparams.multiproblem_mixing_schedule
                  == MixingSchedule.CONSTANT
              ):
                  prob = get_const_sched_prob()
              elif (
                  hparams.multiproblem_mixing_schedule
                  == MixingSchedule.PRETRAIN
              ):
                  prob = get_pretrain_sched_prob()
              else:
                  raise ValueError(
                      "Unknown schedule %s"
                      % str(hparams.multiproblem_mixing_schedule)
                  )
              tf.logging.info(
                  "Using the %s schedule to "
                  "train the MultiProblem."
                  % str(hparams.multiproblem_mixing_schedule)
              )
              tf.logging.info(
                  "Schedule mixing threshold "
                  "%.2f" % hparams.multiproblem_schedule_threshold
              )
  
              # If per-task thresholds are specified, use them.
              thresholds = None
              if hparams.multiproblem_per_task_threshold:
                  thresholds = (
                      hparams.multiproblem_per_task_threshold.split(",")
                  )
                  thresholds = [
                      float(t) for t in thresholds
                  ]  # Convert to floats.
                  thresholds_sum = sum(thresholds)
                  tf.logging.info(
                      "Per-task thresholds: %s." % str(thresholds)
                  )
                  thresholds = [
                      t / thresholds_sum for t in thresholds
                  ]  # Normalize.
                  thresholds = [
                      sum(thresholds[: i + 1])
                      for i in range(len(thresholds))
                  ]
                  tf.logging.info(
                      "Per-task threshold sums: %s." % str(thresholds)
                  )
                  if len(thresholds) != len(self.task_list):
                      tf.logging.warn(
                          "Specified %d thresholds but encountered %d tasks."
                          % (len(thresholds), len(self.task_list))
                      )
                      thresholds = None
  
              def sample_task(curr_task, num_tasks_left, randnum):
                  """A recursive function to sample a task.
  
                  This function treats the probability as the threshold for the primary
                  task and divides the remaining probability mass across the other
                  tasks.
  
                  Args:
                    curr_task: The index of the task being considered for sampling.
                    num_tasks_left: Number of tasks remaining to possibly sample from.
                    randnum: The random number used to select the dataset.
  
                  Returns:
                    A Tensor representing an example from the task that was sampled
                    from.
                  """
                  if num_tasks_left == 0:
                      return get_next_from_dataset(
                          dataset_iterators[curr_task]
                      )
  
                  if (
                      thresholds is not None
                  ):  # Use per-task thresholds if specified.
                      prob_sum = thresholds[curr_task]
                      return tf.cond(
                          randnum < prob_sum,
                          lambda: get_next_from_dataset(
                              dataset_iterators[curr_task]
                          ),
                          lambda: sample_task(
                              curr_task + 1, num_tasks_left - 1, randnum
                          ),
                      )
  
                  # When curr_task is 0, the primary task, the new prob is the same as
                  # the original probability. `tf.greater` indicates that the primary
                  # task receives (1-prob) of the probability mass.
                  # Otherwise, `prob` is divided equally amongst all the secondary
                  # tasks.
                  new_prob = prob - (
                      curr_task * prob / (len(self.task_list) - 1)
                  )
                  return tf.cond(
                      tf.greater(randnum, new_prob),
                      lambda: get_next_from_dataset(
                          dataset_iterators[curr_task]
                      ),
                      lambda: sample_task(
                          curr_task + 1, num_tasks_left - 1, randnum
                      ),
                  )
  
              return tf.data.Dataset.from_tensors(
                  sample_task(
                      0, len(self.task_list) - 1, tf.random_uniform([])
                  )
              )
  
          single_mtl_dataset = tf.data.Dataset.from_tensors(
              tf.zeros([1])
          ).repeat()
          single_mtl_dataset = single_mtl_dataset.flat_map(mix_data)
  
      else:
          if hparams.multiproblem_target_eval_only:
              single_mtl_dataset = datasets[1]
          else:
              single_mtl_dataset = tf.data.Dataset.zip(
                  tuple(datasets)
              ).flat_map(flatten_zip_dataset)
  
      return single_mtl_dataset
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from tensor2tensor.data_generators import problem
  from tensor2tensor.data_generators import text_problems
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers import discretization
  from tensor2tensor.layers import modalities
  from tensor2tensor.utils import metrics
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 6 for calling the MultiProblem obj's underlying member method code with example usage: obj.dataset(mode, data_dir, num_threads, output_buffer_size, shuffle_files, hparams, preprocess, dataset_split, shard, partition_id, num_partitions, shuffle_buffer_size, max_records) and returns: single_mtl_dataset, dataset_iter, hparams, tf, unknown, tf, tf, get_next_from_dataset, tf