code: def finalize_word_mask_info(
      batch,
      spectrogram_feature_name,
      patch_size,
      max_num_word_masks,
      max_num_masked_input_indices,
  ):
      """Format the word mask related information in the batch dict."""
      spectrogram = batch[spectrogram_feature_name]
      len_spec = tf.shape(spectrogram)[0]
      num_feats = tf.shape(spectrogram)[1]
  
      len_spec = (
          tf.cast(len_spec / patch_size[0], tf.int32) * patch_size[0]
      )
  
      start_idxs = batch["mwp_start_indices"]
      end_idxs = batch["mwp_end_indices"]
      masked_targets = batch["masked_targets"]
  
      oor_mask = tf.less(start_idxs, len_spec)
  
      start_idxs = tf.boolean_mask(start_idxs, oor_mask)
      end_idxs = tf.boolean_mask(end_idxs, oor_mask)
      masked_targets = tf.boolean_mask(masked_targets, oor_mask, axis=0)
  
      # now convert spectrogram index to token index using patch size
      start_idxs = tf.math.divide(
          tf.cast(start_idxs, tf.float32),
          tf.cast(patch_size[0], tf.float32),
      )
      end_idxs = tf.math.divide(
          tf.cast(end_idxs, tf.float32),
          tf.cast(patch_size[0], tf.float32),
      )
  
      num_tokens_per_step = tf.cast(
          tf.math.divide(
              tf.cast(num_feats, tf.int32),
              tf.cast(patch_size[1], tf.int32),
          ),
          tf.int32,
      )
      start_token_index = (
          tf.cast(start_idxs, tf.int32) * num_tokens_per_step
      )
      end_token_index = (
          tf.cast(end_idxs + 1, tf.int32) * num_tokens_per_step
      )
  
      token_length = (
          tf.cast(len_spec / patch_size[0], tf.int32)
          * num_tokens_per_step
      )
      idx_pool = tf.range(token_length, dtype=tf.int32)
      word_token_idx_mask = tf.logical_and(
          tf.greater_equal(
              idx_pool[None, :], start_token_index[:, None]
          ),
          tf.less(idx_pool[None, :], end_token_index[:, None]),
      )
  
      masked_input_idxs = tf.ragged.boolean_mask(
          tf.tile(idx_pool[None, :], [tf.size(start_token_index), 1]),
          word_token_idx_mask,
      )
      valid_input_idx_mask = tf.ones_like(masked_input_idxs)
  
      batch["masked_targets"] = pad_first_dim(
          masked_targets, max_num_word_masks, 0
      )
      batch["masked_input_token_indices"] = masked_input_idxs.to_tensor(
          shape=[max_num_word_masks, max_num_masked_input_indices]
      )
      batch[
          "valid_input_token_index_mask"
      ] = valid_input_idx_mask.to_tensor(
          shape=[max_num_word_masks, max_num_masked_input_indices]
      )
  
      del batch["mwp_start_indices"], batch["mwp_end_indices"]
  
      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 8 for calling a global function with example usage: finalize_word_mask_info(batch, spectrogram_feature_name, patch_size, max_num_word_masks, max_num_masked_input_indices) and returns: batch