code: def get_dataset(
      *,
      batch_size,
      eval_batch_size,
      num_shards,
      dtype_str="float32",
      shuffle_seed=0,  # pylint:disable=unused-argument
      rng=None,
      dataset_configs=None,
      dataset_service_address: Optional[str] = None,
  ):
      """Returns a generator for the audiovisual dataset."""
      del rng
      dataset_configs = dataset_configs or {}
      modalities = dataset_configs.get("modalities", ["features"])
      return_as_dict = dataset_configs.get("return_as_dict", False)
      num_frames = dataset_configs.get("num_frames", 100)
      stride = dataset_configs.get("stride", 2)
      eval_stride = dataset_configs.get("eval_stride", 2)
      augmentation_params = dataset_configs.get(
          "augmentation_params", None
      )
      num_train_clips = dataset_configs.get("num_train_clips", 1)
      num_eval_clips = dataset_configs.get("num_eval_clips", 1)
      max_num_output_words = dataset_configs.get(
          "max_num_output_words", 128
      )
      max_num_captions = dataset_configs.get("max_num_captions", 1)
      caption_string = dataset_configs.get(
          "caption_string", "caption/string"
      )
      train_caption_string = dataset_configs.get(
          "train_caption_string", "caption/string"
      )
      num_train_captions_per_clip = dataset_configs.get(
          "num_train_captions_per_clip", 1
      )
      features_dim = dataset_configs.get("features_dim", 768)
  
      max_num_input_words = dataset_configs.get(
          "max_num_input_words", 512
      )
  
      def validate_config(field):
          if dataset_configs.get(field) is None:
              raise ValueError(
                  f"{field} must be specified for TFRecord dataset."
              )
  
      validate_config("base_dir")
      validate_config("tables")
      validate_config("examples_per_subset")
  
      ds_factory = functools.partial(
          VTTFRecordDatasetFactory,
          base_dir=dataset_configs.get(
              "base_dir"
          ),  # pytype: disable=attribute-error
          tables=dict(
              dataset_configs.get("tables")
          ),  # pytype: disable=attribute-error
          examples_per_subset=dataset_configs.get(
              "examples_per_subset"
          ),  # pytype: disable=attribute-error
          num_groups=jax.process_count(),
          group_index=jax.process_index(),
      )
  
      def create_dataset_iterator(
          subset: Text,
          batch_size_local: int,
          num_clips: int,
          caption_string: str,
          stride: int,
          keep_key_local: bool = False,
          max_num_captions_local: int = 1,
      ) -> Tuple[Iterator[Batch], int]:
          is_training = subset == "train"
          logging.info("Loading split %s", subset)
  
          dataset, num_examples = load_split_from_dmvr(
              ds_factory,
              dataset_configs=dataset_configs,
              batch_size=batch_size_local,
              subset=subset,
              stride=stride,
              num_test_clips=num_clips,
              keep_key=keep_key_local,
              max_num_captions=max_num_captions_local,
              caption_string=caption_string,
          )
  
          if dataset_service_address and is_training:
              if shuffle_seed is not None:
                  raise ValueError(
                      "Using dataset service with a random seed causes each "
                      "worker to produce exactly the same data. Add "
                      "config.shuffle_seed = None to your config if you "
                      "want to run with dataset service."
                  )
              logging.info(
                  "Using the tf.data service at %s",
                  dataset_service_address,
              )
              dataset = dataset_utils.distribute(
                  dataset, dataset_service_address
              )
  
          maybe_pad_batches = functools.partial(
              dataset_utils.maybe_pad_batch,
              train=is_training,
              batch_size=batch_size_local,
              inputs_key=None,
          )
          shard_batches = functools.partial(
              dataset_utils.shard, n_devices=num_shards
          )
          current_iter = iter(dataset)
          current_iter = map(dataset_utils.tf_to_numpy, current_iter)
          current_iter = map(
              functools.partial(
                  map_input_keys,
                  modalities=modalities,
                  return_as_dict=return_as_dict,
              ),
              current_iter,
          )
          current_iter = map(maybe_pad_batches, current_iter)
  
          if augmentation_params and augmentation_params.get(
              "do_mixup", False
          ):
              raise ValueError("mixup should be done in the trainer.")
  
          current_iter = map(shard_batches, current_iter)
  
          if is_training and dataset_configs.get("prefetch_to_device"):
              # Async bind batch to device which speeds up training.
              current_iter = jax_utils.prefetch_to_device(
                  current_iter,
                  dataset_configs.get("prefetch_to_device"),
              )
  
          return current_iter, num_examples
  
      train_iter, n_train_examples = create_dataset_iterator(
          "train",
          batch_size,
          num_train_clips,
          train_caption_string,
          stride,
      )
  
      eval_iters = {}
      num_eval_examples = {}
      for (
          k
      ) in (
          dataset_configs.tables.keys()
      ):  # pytype: disable=attribute-error
          if k == "train":
              continue
          eval_iter, n_eval_examples = create_dataset_iterator(
              k,
              eval_batch_size,
              num_eval_clips,
              caption_string,
              eval_stride,
              True,
              max_num_captions,
          )
          eval_iters[k] = eval_iter
          num_eval_examples[k] = n_eval_examples * num_eval_clips
  
      meta_data = {
          # The training loop iterates each caption sample individually. Therefore,
          # the number of captions per clip is muliplied for computing the number of
          # training examples. In contrast, for the eval sets, all GT captions are
          # used together for evaluating a single clip.
          "num_train_examples": (
              n_train_examples
              * num_train_clips
              * num_train_captions_per_clip
          ),
          "num_eval_examples": num_eval_examples,
          "encoder_input_dtype": getattr(jnp, dtype_str),
          "encoder_input_text_dtype": jnp.int32,
          # The sahpes below is for the train model where the number of captions is
          # set to 1 and the dimension 1 corresponding to the number of captions is
          # squeezed.
          "decoder_input_shape": {
              "decoder_input_tokens": (-1, max_num_output_words),
              "decoder_target_tokens": (-1, max_num_output_words),
          },
          "decoder_input_dtype": jnp.int32,
      }
  
      if return_as_dict:
          input_shape_dict = {
              "text": (-1, max_num_input_words),
              "features": (-1, num_frames, features_dim),
          }
  
          meta_data["encoder_input_shape"] = {
              m: input_shape_dict[m] for m in modalities
          }
      elif len(modalities) == 1 and modalities[0] == "features":
          meta_data["encoder_input_shape"] = (
              -1,
              num_frames,
              features_dim,
          )
      elif len(modalities) == 1 and modalities[0] == "text":
          meta_data["encoder_input_shape"] = (-1, max_num_input_words)
      else:
          raise NotImplementedError("modality not supported")
      logging.info("Dataset metadata:\n%s", meta_data)
  
      return dataset_utils.Dataset(
          train_iter, eval_iters, None, meta_data
      )
  
imports: import functools
  from typing import Dict, Iterator, List, Optional, Text, Tuple, Union
  from absl import logging
  from dmvr import tokenizers
  from flax import jax_utils
  import jax
  import jax.numpy
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.projects.t5 import tokenizer
  from scenic.projects.vid2seq import data_utils
  from scenic.projects.vivit.data import video_tfrecord_dataset
  import tensorflow
label_desc: Use image node_id 4 for calling a global function with example usage: get_dataset() and returns: dataset_utils, current_iter,  num_examples