code: def load_split_from_dmvr(
      ds_factory,
      batch_size,
      subset="train",
      modalities=("rgb"),
      num_frames=32,
      stride=2,
      num_spec_frames=5,
      num_test_clips=1,
      min_resize=256,
      crop_size=224,
      spec_shape=(96, 64),
      dataset_spec_mean=0.0,
      dataset_spec_stddev=1.0,
      spec_augment=False,
      spec_augment_params=None,
      num_waveform_samples=32000,
      waveform_stride=1,
      zero_centering=True,
      spec_compute_online: bool = False,
      spec_from_wave_sample_rate: int = 16000,
      spec_from_wave_snr: Optional[float] = None,
      spec_from_wave_add_gaussian_noise: bool = False,
      spec_from_wave_add_masking_noise: bool = False,
      spec_from_wave_visualness_threshold: float = 0.2,
      spec_from_wave_random_mask_noise_rate: float = 0.4,
      spec_from_wave_max_word_len: int = 128,
      spec_from_wave_extend_mask_boundaries_ms: float = 0.0,
      spec_from_wave_max_num_masks: int = -1,
      spec_from_wave_add_word_mask: bool = True,
      spec_from_wave_add_word_mask_info: bool = False,
      spec_from_wave_eval_noise_types: Optional[Sequence[str]] = None,
      spec_from_wave_environment_noise_path: Optional[str] = None,
      spec_from_wave_eval_noise_configs: Optional[
          Dict[str, Dict[str, Any]]
      ] = None,
      spectrogram_type: str = "logmf",
      spec_frame_length: int = 400,
      spec_frame_step: int = 160,
      spec_num_features: int = 80,
      spec_lower_edge_hertz: float = 0.0,
      spec_upper_edge_hertz: float = 7600.0,
      augmentation_params=None,
      keep_key=False,
      max_num_words: int = 16,
      max_num_captions: int = 1,
      tokenizer_type="bert",
      tokenizer_vocab=None,
      prepend_bos: bool = False,
      append_eos: bool = True,
      caption_string="caption/string",
      masked_word_pred_aligned_caption_string: str = "caption/label/string",
      masked_word_pred_max_num_tokens: int = 8,
      masked_word_pred_patch_size: Sequence[int] = (),
      masked_word_pred_max_num_masked_input_indices: int = 64,
  ):
      """Loads dataset using DMVR for pre-processing.
  
      DMVR dataset loader already does basic augmentation (random crop and flip in
        train mode. It also already shuffles and batches the data.
  
      Args:
        ds_factory: A DMVR factory to instantiate with the subset.
        batch_size: The batch_size to use.
        subset: train, validation or test.
        modalities: list of input modalities.
        num_frames: Number of RGB frames per subclip.
        stride: Temporal stride to sample RGB frames.
        num_spec_frames: Number of spectrogram frames per subclip.
        num_test_clips: Number of test clips (1 by default). If more than 1, this
          will sample multiple linearly spaced clips within each video at test time.
          If 1, then a single clip in the middle of the video is sampled. The clips
          are aggreagated in the batch dimension.
        min_resize: Frames are resized so that min(height, width) is min_resize.
        crop_size: Final size of the frame after cropping the resized frames. Both
          height and width are the same.
        spec_shape: Input size of spectrogram per frame.
        dataset_spec_mean: Mean of spectrograms in the dataset.
        dataset_spec_stddev: Std dev of spectrograms in the dataset.
        spec_augment: whether to apply augmentation using SpecAugment.
        spec_augment_params: dict; augmentation configurations for SpecAugment
        num_waveform_samples: Number of waveform samples to use, defaults to 32000
          which corresponds to two seconds at 16kHz sampling rate.
        waveform_stride: Temporal stride to sample waveform.
        zero_centering: If True, frames are normalized to values in [-1, 1]. If
          False, values in [0, 1].
        spec_compute_online: whether to compute spectrograms on the fly.
        spec_from_wave_sample_rate: spectrogram computation parameter. The sample
          rate of the input audio.
        spec_from_wave_snr: spectrogram computation parameter. Signal-to-noise ratio
          used to inject white noise. The default value None or non-positive values
          disable the noise injection. Noise is add to non-training examples only by
          default. Training noise is added if `add_train_noise` is `True`.
        spec_from_wave_add_gaussian_noise: Whether to add Gaussian noise.
        spec_from_wave_add_masking_noise: If the noise type is word masking noise.
        spec_from_wave_visualness_threshold: The threshold for determining the
          visual words.
        spec_from_wave_random_mask_noise_rate: The random mask noise sampling
          ratio. If 1, the word masking is deterministic. If below 1,
          the masking noise is random and each word is masked out with this
          chance.
        spec_from_wave_max_word_len: Maximum length of the words. This is used to
          pad and truncate noise_word_mask added to the batch dictionary in the word
          masking noise addition.
        spec_from_wave_extend_mask_boundaries_ms: If set, the start and end
          boundaries of each maksing region in the signal is extended.
        spec_from_wave_max_num_masks: Maximum number of word masks to apply. -1
          means unlimited.
        spec_from_wave_add_word_mask: If set, add masks indicating masked words
          into the batch dict. Used for computing recovery rate of the masked
          words.
        spec_from_wave_add_word_mask_info: If set, add word masking related
          information such as timestamps and masks to the batch dict. Used for
          applying masked word prediction loss.
        spec_from_wave_eval_noise_types: A tuple of noise type strings used in
          evaluation. Each string should be either `environment_noise` or
          `packet_loss_noise`.
        spec_from_wave_environment_noise_path: Path to the npy file containing
          noise waveforms in numpy array.
        spec_from_wave_eval_noise_configs: A dict of noise config dicts for each
          noise type listed in `spec_from_wave_eval_noise_types`.
        spectrogram_type: spectrogram computation parameter. The type of the
          spectrogram to be extracted from the waveform. Can be either
          `spectrogram`, `logmf`, and `mfcc`.
        spec_frame_length: spectrogram computation parameter. The length of each
          spectroram frame.
        spec_frame_step: spectrogram computation parameter. The stride of
          spectrogram frames.
        spec_num_features: spectrogram computation parameter. The number of
          spectrogram features.
        spec_lower_edge_hertz: spectrogram computation parameter. Lowest frequency
          to consider.
        spec_upper_edge_hertz: spectrogram computation parameter. Highest frequency
          to consider.
        augmentation_params: dict; augmentation configurations in train mode.
        keep_key: bool; If true, also return the key for each example.
        max_num_words: Maximum number of tokens to keep from the text for each
          caption. If there are more tokens, sequence is cropped, if less, the
          caption is padded using the tokenizer pad id.
        max_num_captions: Maximum number of captions to keep. If there are more
          captions in the proto, only the first `max_num_captions` will be returned
          is `is_training` is set to `False`. If `is_training` is `True`, then
          `max_num_captions` will be randomly sampled. Finally if the proto contains
          less than `max_num_captions`, we pad with empty srings to make sure there
          are `max_num_captions` in total.
        tokenizer_type: The type of tokenizer. Supported types: ('bert',).
        tokenizer_vocab: The path to the tokenizer vocabulary.
        prepend_bos: Whether to prepend BOS token.
        append_eos: Whether to append EOS token.
        caption_string: Input feature name in sstable for caption.
        masked_word_pred_aligned_caption_string: The feature name to parse to
          extract the caption aligned with the timestamps. Used for the masked
          word prediction.
        masked_word_pred_max_num_tokens: Maximum number of tokens for each masked
          word.
        masked_word_pred_patch_size: The patch size used in the network architecture
          for spectrogram token embedding. Used to compute the token index.
        masked_word_pred_max_num_masked_input_indices: Maximum number of masked
          input tokens (spectrogram tokens).
  
      Returns:
        A pair `(ds, num_examples)` with
          ds: A `tf.data.Dataset` object
          num_examples: Number of examples in the dataset.
      """
      is_training = subset == "train"
  
      if tokenizer_type == "bert":
          assert tokenizer_vocab
          tokenizer = tokenizers.BertTokenizer(tokenizer_vocab)
      else:
          raise ValueError("Tokenizer not supported")
      vocab_size = int(tokenizer.vocab_size)
      logging.info("vocab_size %d", vocab_size)
      logging.info("EOS token: %d", tokenizer.eos_token)
      # Init the TF models of the tokenizer.
      tokenizer.initialize()
  
      ds_factory = ds_factory(
          subset=subset, modalities=modalities
      ).configure(
          is_training=is_training,
          num_frames=num_frames,
          stride=stride,
          num_spec_frames=num_spec_frames,
          num_test_clips=num_test_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=spec_from_wave_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=spec_from_wave_eval_noise_types,
          spec_from_wave_environment_noise_path=spec_from_wave_environment_noise_path,
          spec_from_wave_eval_noise_configs=spec_from_wave_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_image=zero_centering,
          max_num_words=max_num_words,
          max_num_captions=max_num_captions,
          tokenizer=tokenizer,
          prepend_bos=prepend_bos,
          append_eos=append_eos,
          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 "rgb" in modalities and is_training and augmentation_params:
          # additional augmentation for the RGB features.
          ds_factory = video_ops.additional_augmentations(
              ds_factory,
              augmentation_params,
              crop_size,
              num_frames,
              zero_centering,
          )
  
      logging.info(
          "Preprocessing graph: %s",
          ds_factory.preprocessor_builder.get_summary(),
      )
      logging.info(
          "Postprocessing graph: %s",
          ds_factory.postprocessor_builder.get_summary(),
      )
  
      # Val and test splits are a single epoch otherwise the last
      # batch is not padded with zeros but with valid examples
      num_examples = ds_factory.num_examples
      ds = ds_factory.make_dataset(
          batch_size=batch_size,
          shuffle=is_training,
          num_epochs=None if is_training else 1,
          drop_remainder=is_training,
          keep_key=(not is_training and keep_key),
      )
  
      if not is_training:
          # Repeat indefinitely
          ds = ds.repeat(None)
  
      options = tf.data.Options()
      options.experimental_threading.private_threadpool_size = 48
      ds = ds.with_options(options)
  
      return ds, num_examples
  
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 6 for calling a global function with example usage: load_split_from_dmvr(ds_factory, batch_size, subset, modalities, num_frames, stride, num_spec_frames, num_test_clips, min_resize, crop_size, spec_shape, dataset_spec_mean, dataset_spec_stddev, spec_augment, spec_augment_params, num_waveform_samples, waveform_stride, zero_centering, spec_compute_online, spec_from_wave_sample_rate, spec_from_wave_snr, spec_from_wave_add_gaussian_noise, spec_from_wave_add_masking_noise, spec_from_wave_visualness_threshold, spec_from_wave_random_mask_noise_rate, spec_from_wave_max_word_len, spec_from_wave_extend_mask_boundaries_ms, spec_from_wave_max_num_masks, spec_from_wave_add_word_mask, spec_from_wave_add_word_mask_info, spec_from_wave_eval_noise_types, spec_from_wave_environment_noise_path, spec_from_wave_eval_noise_configs, spectrogram_type, spec_frame_length, spec_frame_step, spec_num_features, spec_lower_edge_hertz, spec_upper_edge_hertz, augmentation_params, keep_key, max_num_words, max_num_captions, tokenizer_type, tokenizer_vocab, prepend_bos, append_eos, caption_string, masked_word_pred_aligned_caption_string, masked_word_pred_max_num_tokens, masked_word_pred_patch_size, masked_word_pred_max_num_masked_input_indices) and returns: ds,  num_examples