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 video-text dataset."""
      del rng
      dataset_configs = dataset_configs or {}
      # RGB related configs.
      num_frames = dataset_configs.get("num_frames", 32)
      stride = dataset_configs.get("stride", 2)
      test_stride = dataset_configs.get("test_stride", 2)
      min_resize = dataset_configs.get("min_resize", 256)
      crop_size = dataset_configs.get("crop_size", 224)
      # General configs.
      num_test_clips = dataset_configs.get("num_test_clips", 1)
      num_test_captions = dataset_configs.get("num_test_captions", 1)
      num_train_val_clips = dataset_configs.get(
          "num_train_val_clips", 1
      )
      num_train_val_captions = dataset_configs.get(
          "num_train_val_captions", 0
      )
      if num_train_val_captions:
          num_train_captions = num_val_captions = num_train_val_captions
      else:
          num_train_captions = dataset_configs.get(
              "num_train_captions", 1
          )
          num_val_captions = dataset_configs.get("num_val_captions", 1)
      zero_centering = dataset_configs.get("zero_centering", True)
      augmentation_params = dataset_configs.get(
          "augmentation_params", None
      )
      # Whether to load the key for each test sample from the data sstables.
      keep_test_key = dataset_configs.get("keep_test_key", False)
      test_split = dataset_configs.get("test_split")
      # For the test set, the actual batch size is test_batch_size * num_test_clips
      test_batch_size = dataset_configs.get(
          "test_batch_size", eval_batch_size
      )
      # Text related configs.
      max_num_words = dataset_configs.get("max_num_words", 16)
      max_num_captions = dataset_configs.get("max_num_captions", 0)
      if max_num_captions > 0:
          num_train_captions = (
              num_val_captions
          ) = num_test_captions = max_num_captions
      caption_string = dataset_configs.get(
          "caption_string", "caption/string"
      )
      num_labels = dataset_configs.get("num_labels", 0)
      caption_string_train = dataset_configs.get(
          "caption_string_train", "clip/label/string"
      )
      rmv_full_stop = dataset_configs.get("rmv_full_stop", False)
      vqa_options_for_ce = dataset_configs.get(
          "vqa_options_for_ce", False
      )
      include_verb = dataset_configs.get("include_verb", False)
  
      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(
          AVTFRecordDatasetFactory,
          base_dir=dataset_configs.get("base_dir", ""),
          tables=dataset_configs.get("tables", {}),
          examples_per_subset=dataset_configs.get(
              "examples_per_subset"
          ),
          num_groups=jax.process_count(),
          group_index=jax.process_index(),
      )
  
      def create_dataset_iterator(
          subset: Text,
          batch_size_local: int,
          num_clips: int,
          num_captions: int,
          caption_string: str,
          stride: int,
          rmv_full_stop: bool,
          num_labels: int,
          include_verb: bool,
          keep_key_local: bool = False,
      ) -> Tuple[Iterator[Batch], int]:
          is_training = subset == "train"
          is_test = subset == "test"
          logging.info("Loading split %s", subset)
  
          dataset, num_examples = load_split_from_dmvr(
              ds_factory,
              batch_size=batch_size_local,
              subset=subset,
              num_frames=num_frames,
              stride=stride,
              num_test_clips=num_clips,
              min_resize=min_resize,
              crop_size=crop_size,
              zero_centering=zero_centering,
              augmentation_params=augmentation_params,
              keep_key=keep_key_local,
              max_num_words=max_num_words,
              max_num_captions=num_captions,
              caption_string=caption_string,
              num_labels=num_labels,
              include_verb=include_verb,
          )
  
          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
              )
  
          current_iter = iter(dataset)
          current_iter = map(dataset_utils.tf_to_numpy, current_iter)
          current_iter = map(map_keys, current_iter)
          current_iter = map(
              functools.partial(
                  clip_tokenize,
                  rmv_full_stop=rmv_full_stop,
                  vqa_options_for_ce=vqa_options_for_ce,
                  include_verb=include_verb,
              ),
              current_iter,
          )
  
          pad_batch_size = batch_size_local
          if is_test:
              pad_batch_size = batch_size_local * num_clips
          maybe_pad_batches = functools.partial(
              maybe_pad_batch,
              train=is_training,
              batch_size=pad_batch_size,
              num_clips=num_clips,
              num_captions=num_captions,
          )
          current_iter = map(maybe_pad_batches, current_iter)
  
          shard_batches = functools.partial(
              dataset_utils.shard, n_devices=num_shards
          )
          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_val_clips,
          num_train_captions,
          caption_string_train,
          stride,
          rmv_full_stop,
          num_labels,
          include_verb,
      )
      eval_iter, n_eval_examples = create_dataset_iterator(
          "validation",
          eval_batch_size,
          num_train_val_clips,
          num_val_captions,
          caption_string,
          test_stride,
          rmv_full_stop,
          num_labels,
          False,
          keep_test_key,
      )
      n_test_examples, test_iter = 0, None
      if test_split:
          test_iter, n_test_examples = create_dataset_iterator(
              test_split,
              test_batch_size,
              num_test_clips,
              num_test_captions,
              caption_string,
              test_stride,
              rmv_full_stop,
              num_labels,
              False,
              keep_test_key,
          )
  
      meta_data = {
          "num_train_examples": (
              n_train_examples * num_train_val_clips
          ),
          "num_eval_examples": (n_eval_examples * num_train_val_clips),
          "num_test_examples": (n_test_examples * num_test_clips),
          "input_dtype": getattr(jnp, dtype_str),
      }
      meta_data["text_shape"] = [-1, max_num_captions, max_num_words]
      meta_data["text_dtype"] = jnp.int32
      meta_data["input_shape"] = {
          "rgb": (-1, num_frames, crop_size, crop_size, 3),
      }
      logging.info("Dataset metadata:\n%s", meta_data)
  
      return dataset_utils.Dataset(
          train_iter, eval_iter, test_iter, meta_data
      )
  
imports: import functools
  from typing import Dict, Iterator, List, Optional, Text, Tuple, Union
  from absl import logging
  from dmvr import modalities
  from flax import jax_utils
  import jax
  import jax.numpy
  import numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  from scenic.dataset_lib import video_ops
  from scenic.projects.baselines.clip import tokenizer
  from scenic.projects.vivit.data import video_tfrecord_dataset
  import tensorflow
  from dmvr import tokenizers
label_desc: Use image node_id 5 for calling a global function with example usage: get_dataset() and returns: dataset_utils, current_iter,  num_examples