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,
      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", ["rgb"])
      return_as_dict = dataset_configs.get("return_as_dict", True)
      # RGB related configs.
      num_frames = dataset_configs.get("num_frames", 32)
      stride = dataset_configs.get("stride", 2)
      eval_stride = dataset_configs.get("eval_stride", stride)
      min_resize = dataset_configs.get("min_resize", 256)
      crop_size = dataset_configs.get("crop_size", 224)
      # Spectrogram related configs.
      num_spec_frames = dataset_configs.get("num_spec_frames", 25)
      eval_num_spec_frames = dataset_configs.get(
          "eval_num_spec_frames", num_spec_frames
      )
      spec_shape = dataset_configs.get("spec_shape", (100, 80))
      spec_augment = dataset_configs.get("spec_augment", False)
      spec_augment_params = dataset_configs.get(
          "spec_augment_params", None
      )
      dataset_spec_mean = dataset_configs.get("spec_mean", 0.0)
      dataset_spec_stddev = dataset_configs.get("spec_stddev", 1.0)
      # Waveform related configs
      num_waveform_samples = dataset_configs.get(
          "num_waveform_samples", 16000
      )
      eval_num_waveform_samples = dataset_configs.get(
          "eval_num_waveform_samples", num_waveform_samples
      )
      waveform_stride = dataset_configs.get("waveform_stride", 1)
      # Spectrogram computation related configs
      spec_compute_online = dataset_configs.get(
          "spec_compute_online", False
      )
      spec_from_wave_sample_rate = dataset_configs.get(
          "spec_from_wave_sample_rate", 16000
      )
      spec_from_wave_snr = dataset_configs.get(
          "spec_from_wave_snr", None
      )
      spec_from_wave_add_gaussian_noise = dataset_configs.get(
          "spec_from_wave_add_gaussian_noise", False
      )
      spec_from_wave_add_train_masking_noise = dataset_configs.get(
          "spec_from_wave_add_train_masking_noise", False
      )
      spec_from_wave_visualness_threshold = dataset_configs.get(
          "spec_from_wave_visualness_threshold", None
      )
      spec_from_wave_random_mask_noise_rate = dataset_configs.get(
          "spec_from_wave_random_mask_noise_rate", 0.4
      )
      spec_from_wave_max_word_len = dataset_configs.get(
          "spec_from_wave_max_word_len", 0
      )
      spec_from_wave_extend_mask_boundaries_ms = dataset_configs.get(
          "spec_from_wave_extend_mask_boundaries_ms", 0.0
      )
      spec_from_wave_max_num_masks = dataset_configs.get(
          "spec_from_wave_max_num_masks", 16
      )
      spec_from_wave_add_word_mask = dataset_configs.get(
          "spec_from_wave_add_word_mask", False
      )
      spec_from_wave_add_word_mask_info = dataset_configs.get(
          "spec_from_wave_add_word_mask_info", False
      )
      spec_from_wave_eval_noise_configs = dataset_configs.get(
          "spec_from_wave_eval_noise_configs", None
      )
      spec_from_wave_environment_noise_path = dataset_configs.get(
          "spec_from_wave_environment_noise_path", None
      )
      spectrogram_type = dataset_configs.get(
          "spectrogram_type", "logmf"
      )
      spec_frame_length = dataset_configs.get("spec_frame_length", 400)
      spec_frame_step = dataset_configs.get("spec_frame_step", 160)
      spec_num_features = dataset_configs.get("spec_num_features", 80)
      spec_lower_edge_hertz = dataset_configs.get(
          "spec_lower_edge_hertz", 0.0
      )
      spec_upper_edge_hertz = dataset_configs.get(
          "spec_upper_edge_hertz", 7600.0
      )
      masked_word_pred_aligned_caption_string = dataset_configs.get(
          "masked_word_pred_aligned_caption_string",
          "caption/label/string",
      )
      masked_word_pred_max_num_tokens = dataset_configs.get(
          "masked_word_pred_max_num_tokens", 8
      )
      masked_word_pred_patch_size = dataset_configs.get(
          "masked_word_pred_patch_size", []
      )
      masked_word_pred_max_num_masked_input_indices = (
          dataset_configs.get(
              "masked_word_pred_max_num_masked_input_indices", 32
          )
      )
      # General configs.
      num_eval_clips = dataset_configs.get("num_eval_clips", 1)
      zero_centre_data = dataset_configs.get("zero_centering", True)
      augmentation_params = dataset_configs.get(
          "augmentation_params", None
      )
      num_train_val_clips = dataset_configs.get(
          "num_train_val_clips", 1
      )
      do_three_spatial_crops = dataset_configs.get(
          "do_three_spatial_crops", False
      )
      num_spatial_crops = 3 if do_three_spatial_crops else 1
  
      # Should hold two fields: tokenizer_type and tokenizer_vocab.
      tokenizer_config = dataset_configs.get("tokenizer", {})
      max_num_words = dataset_configs.get("max_num_words", 16)
      eval_max_num_words = dataset_configs.get(
          "eval_max_num_words", max_num_words
      )
      max_num_captions = dataset_configs.get("max_num_captions", 1)
      caption_string = dataset_configs.get(
          "caption_string", "clip/label/string"
      )
  
      if (
          spec_from_wave_add_train_masking_noise
          and spec_from_wave_visualness_threshold is None
      ):
          raise ValueError(
              "`spec_from_wave_visualness_threshold` must be specified when "
              "`spec_from_wave_add_train_masking_noise` is True"
          )
  
      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")
  
      ds_factory = functools.partial(
          ASRTFRecordDatasetFactory,
          base_dir=dataset_configs.get("base_dir"),
          tables=dataset_configs.get("tables"),
          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,
          num_spec_frames: int,
          num_waveform_samples: int,
          max_num_words: int,
          keep_key_local: bool = True,
          add_masking_noise: bool = False,
          eval_noise_types: Optional[Sequence[str]] = None,
          eval_noise_configs: Optional[
              Dict[str, Dict[str, Any]]
          ] = None,
      ) -> Tuple[Iterator[Batch], int]:
          is_training = subset == "train"
          is_test = subset == "test"
          logging.info("Loading split %s", subset)
  
          # TODO(phseo): Remove duplicates and pass the dict itself.
          dataset, num_examples = load_split_from_dmvr(
              ds_factory,
              batch_size=batch_size_local,
              subset=subset,
              modalities=modalities,
              num_frames=num_frames,
              stride=stride,
              num_spec_frames=num_spec_frames,
              num_test_clips=num_clips,
              min_resize=min_resize,
              crop_size=crop_size,
              spec_shape=spec_shape,
              dataset_spec_mean=dataset_spec_mean,
              dataset_spec_stddev=dataset_spec_stddev,
              spec_augment=spec_augment,
              spec_augment_params=spec_augment_params,
              spec_compute_online=spec_compute_online,
              spec_from_wave_sample_rate=spec_from_wave_sample_rate,
              spec_from_wave_snr=spec_from_wave_snr,
              spec_from_wave_add_gaussian_noise=spec_from_wave_add_gaussian_noise,
              spec_from_wave_add_masking_noise=add_masking_noise,
              spec_from_wave_visualness_threshold=(
                  spec_from_wave_visualness_threshold
              ),
              spec_from_wave_random_mask_noise_rate=(
                  spec_from_wave_random_mask_noise_rate
              ),
              spec_from_wave_max_word_len=spec_from_wave_max_word_len,
              spec_from_wave_extend_mask_boundaries_ms=(
                  spec_from_wave_extend_mask_boundaries_ms
              ),
              spec_from_wave_max_num_masks=spec_from_wave_max_num_masks,
              spec_from_wave_add_word_mask=spec_from_wave_add_word_mask,
              spec_from_wave_add_word_mask_info=spec_from_wave_add_word_mask_info,
              spec_from_wave_eval_noise_types=eval_noise_types,
              spec_from_wave_environment_noise_path=spec_from_wave_environment_noise_path,
              spec_from_wave_eval_noise_configs=eval_noise_configs,
              spectrogram_type=spectrogram_type,
              spec_frame_length=spec_frame_length,
              spec_frame_step=spec_frame_step,
              spec_num_features=spec_num_features,
              spec_lower_edge_hertz=spec_lower_edge_hertz,
              spec_upper_edge_hertz=spec_upper_edge_hertz,
              num_waveform_samples=num_waveform_samples,
              waveform_stride=waveform_stride,
              zero_centering=zero_centre_data,
              augmentation_params=augmentation_params,
              keep_key=keep_key_local,
              max_num_words=max_num_words,
              max_num_captions=max_num_captions,
              **tokenizer_config,
              caption_string=caption_string,
              masked_word_pred_aligned_caption_string=masked_word_pred_aligned_caption_string,
              masked_word_pred_max_num_tokens=masked_word_pred_max_num_tokens,
              masked_word_pred_patch_size=masked_word_pred_patch_size,
              masked_word_pred_max_num_masked_input_indices=masked_word_pred_max_num_masked_input_indices,
          )
  
          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
              )
  
          pad_batch_size = batch_size_local
          if is_test:
              pad_batch_size = (
                  batch_size_local * num_clips * num_spatial_crops
              )
          maybe_pad_batches = functools.partial(
              maybe_pad_batch,
              train=is_training,
              batch_size=pad_batch_size,
              return_as_dict=return_as_dict,
          )
  
          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_keys,
                  modalities=modalities,
                  return_as_dict=return_as_dict,
              ),
              current_iter,
          )
          current_iter = map(_convert_key_string_to_int, current_iter)
          if not is_training:
              current_iter = map(
                  _convert_caption_string_to_int, 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_val_clips,
          caption_string,
          stride,
          num_spec_frames,
          num_waveform_samples,
          max_num_words,
          add_masking_noise=spec_from_wave_add_train_masking_noise,
      )
      eval_iter, n_eval_examples = create_dataset_iterator(
          "val",
          eval_batch_size,
          num_eval_clips,
          caption_string,
          eval_stride,
          eval_num_spec_frames,
          eval_num_waveform_samples,
          eval_max_num_words,
      )
      test_iter, n_test_examples = create_dataset_iterator(
          "test",
          eval_batch_size,
          num_eval_clips,
          caption_string,
          eval_stride,
          eval_num_spec_frames,
          eval_num_waveform_samples,
          eval_max_num_words,
      )
      if spec_from_wave_eval_noise_configs:
          test_iter = [test_iter]
          for (
              noise_types,
              configs,
          ) in spec_from_wave_eval_noise_configs.items():
              noise_types = noise_types.split(",")
              for config in configs:
                  test_iter.append(
                      create_dataset_iterator(
                          "test",
                          eval_batch_size,
                          num_eval_clips,
                          caption_string,
                          eval_stride,
                          eval_num_spec_frames,
                          eval_num_waveform_samples,
                          eval_max_num_words,
                          eval_noise_types=noise_types,
                          eval_noise_configs=config,
                      )[0]
                  )
  
      meta_data = {
          # pylint:disable=protected-access
          "num_train_examples": (
              n_train_examples * num_train_val_clips
          ),
          "num_eval_examples": (n_eval_examples * num_eval_clips),
          "num_test_examples": (
              n_test_examples * num_eval_clips * num_spatial_crops
          ),
          "input_dtype": getattr(jnp, dtype_str),
      }
  
      # Set the input shapes
      input_shapes = {
          "rgb": (-1, num_frames, crop_size, crop_size, 3),
          "spectrogram": (
              -1,
              num_spec_frames * spec_shape[0],
              spec_shape[1],
              3,
          ),
          "waveform": (-1, num_waveform_samples),
      }
      meta_data["input_shape"] = {}
      for modality, shape in input_shapes.items():
          if modality in modalities:
              meta_data["input_shape"][modality] = shape
  
      meta_data["target_shape"] = (-1, max_num_words)
      meta_data["target_dtype"] = jnp.int32
  
      if spec_from_wave_add_word_mask_info:
          meta_data["masked_token_idxs_shape"] = (
              -1,
              spec_from_wave_max_num_masks,
              masked_word_pred_max_num_masked_input_indices,
          )
          meta_data["masked_token_idx_masks_shape"] = (
              -1,
              spec_from_wave_max_num_masks,
              masked_word_pred_max_num_masked_input_indices,
          )
          meta_data["masked_word_targets_shape"] = (
              -1,
              spec_from_wave_max_num_masks,
              masked_word_pred_max_num_tokens,
          )
          meta_data["masked_token_idxs_dtype"] = jnp.int32
          meta_data["masked_token_idx_masks_dtype"] = jnp.int32
          meta_data["masked_word_targets_dtype"] = jnp.int32
  
      logging.info(
          "Number of training examples: %d",
          meta_data["num_train_examples"],
      )
      logging.info(
          "Number of validation examples: %d",
          meta_data["num_eval_examples"],
      )
      logging.info(
          "Number of test examples: %d", meta_data["num_test_examples"]
      )
  
      return dataset_utils.Dataset(
          train_iter, eval_iter, test_iter, meta_data
      )
  
imports: import functools
  from typing import Any, Dict, Iterator, Optional, Sequence, Text, Tuple
  from absl import logging
  from dmvr import builders
  from dmvr import modalities
  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.dataset_lib import datasets
  from scenic.dataset_lib import video_ops
  from scenic.projects.avatar.datasets.dataset_utils import add_int64
  from scenic.projects.avatar.datasets.dataset_utils import add_spectrogram
  from scenic.projects.avatar.datasets.dataset_utils import add_spectrogram_from_audio
  from scenic.projects.avatar.datasets.dataset_utils import add_text
  from scenic.projects.vivit.data import video_tfrecord_dataset
  import tensorflow
label_desc: Use image node_id 8 for calling a global function with example usage: get_dataset() and returns: dataset_utils, current_iter,  num_examples