code: def add_visual_word_masking_noise(
      batch,
      wave_feature_name: str,
      snr: float = 1.0,
      sample_rate: int = 16000,
      vis_threshold: float = 0.2,
      max_word_len: int = 128,
      mask_rate: float = 0.4,
      max_num_masks: int = -1,
      extend_boundaries_ms: float = 0.0,
      add_word_mask: bool = True,
      add_word_mask_info: bool = False,
  ):
      """Add word masking noise to the wave signal in the input batch.
  
      Masks out wave signal that correspond to words indicated by word_mask. Each
      masked part is filled with Gaussian white noise whose power is determined by
      snr. If snr is None, Gaussian white noise filling is not performed. When
      mask_rate < 1.0, the target words are randomly selected with the probability
      of mask_rate and only the selected words are masked out.
  
      Args:
        batch: The input batch dictionary.
        wave_feature_name: Dictionary key for the raw wave feature.
        snr: The signal to noise ratio to determine the power of the white noise. If
          None, the target regions are simply zeroed out.
        sample_rate: The sample rate of the input wave signal.
        vis_threshold: The threshold for the visualness scores of the input words.
        max_word_len: Maximum length of the words. This is used to pad and truncate
          noise_word_mask which is added in the batch dictionary.
        mask_rate: The masking ratio used for random mask generation.
        max_num_masks: The maximum number of masks to keep. If the number of mask
          candidates are larger than this, the masks are randomly subsampled to
          match this number. If -1, all masks are kept.
        extend_boundaries_ms: If set, the start and end boundaries of each maksing
          region in the signal is extended.
        add_word_mask: If set, add raw word mask into the batch dict. Used to
          compute recovery rate.
        add_word_mask_info: If set, add word masking related information including
          timestamps and length to be used for computing masked word prediction
          loss.
  
      Returns:
        The updated batch dictionary where the wave signal is updated with the noise
        injected one.
      """
  
      vis_score = batch["visualness_score"]
      word_sts = batch["word_start_timestamp"]
      word_ets = batch["word_end_timestamp"]
      wave_tensor = batch[wave_feature_name]
  
      word_mask = vis_score > vis_threshold
  
      logging.info("mask_rate: %f", mask_rate)
      if mask_rate < 1.0:
          logging.info("doing random masking")
          random_mask = (
              tf.random.uniform(tf.shape(word_mask)) < mask_rate
          )
          word_mask = tf.logical_and(random_mask, word_mask)
  
      logging.info("max_num_masks: %d", max_num_masks)
      logging.info("word_mask: %s", word_mask)
  
      if max_num_masks > -1:
          if (
              tf.reduce_sum(tf.cast(word_mask, tf.int32))
              > max_num_masks
          ):
              indices = tf.where(word_mask)
              indices = tf.random.shuffle(indices)[:max_num_masks]
              word_mask = tf.cast(
                  tf.scatter_nd(
                      tf.cast(indices, tf.int32),
                      tf.ones(max_num_masks),
                      tf.shape(word_mask),
                  ),
                  tf.bool,
              )
  
      if add_word_mask:
          logging.info("adding noise_word_mask")
          # This is for computing word recovery rate
          batch["noise_word_mask"] = pad_first_dim(
              word_mask[:max_word_len], max_word_len, 0
          )
  
      (
          batch[wave_feature_name],
          word_start_timestamp,
          word_end_timestamp,
      ) = mask_word_with_white_noise(
          wave_tensor,
          snr,
          word_mask,
          word_sts,
          word_ets,
          sample_rate,
          extend_boundaries_ms,
      )
  
      if add_word_mask_info:
          # This is for masked word prediction loss
          batch["mwp_word_mask"] = word_mask
          batch["mwp_start_timestamp"] = word_start_timestamp
          batch["mwp_end_timestamp"] = word_end_timestamp
          batch["mwp_wave_length"] = (
              tf.math.divide(
                  tf.cast(tf.size(wave_tensor), tf.float32),
                  tf.cast(sample_rate, tf.float32),
              )
              * 1000.0
          )
  
      del (
          batch["word_start_timestamp"],
          batch["word_end_timestamp"],
          batch["visualness_score"],
      )
  
      return batch
  
imports: import functools
  from typing import Any, Dict, Optional, Sequence
  from absl import logging
  from dmvr import builders
  from dmvr import processors
  from dmvr import tokenizers
  import numpy
  from scenic.projects.mbt.datasets.dataset_utils import _decode_spectrogram
  from scenic.projects.mbt.datasets.dataset_utils import apply_specaugment
  import tensorflow
label_desc: Use image node_id 5 for calling a global function with example usage: add_visual_word_masking_noise(batch, wave_feature_name, snr, sample_rate, vis_threshold, max_word_len, mask_rate, max_num_masks, extend_boundaries_ms, add_word_mask, add_word_mask_info) and returns: batch