code: def map_input_keys(
      batch, modalities=("features",), return_as_dict=False
  ):
      """DMVR dataset returns 'image' and 'label'. We want 'inputs' and 'label'."""
      if not return_as_dict:
          if len(modalities) == 1 and modalities[0] == "features":
              batch["encoder_inputs"] = batch["features"]
          elif len(modalities) == 1 and "text" == modalities[0]:
              batch["encoder_inputs"] = batch["asr_indices"]
          else:
              raise NotImplementedError(
                  "modality not supported by map_keys."
              )
      else:
          batch["encoder_inputs"] = {}
          if "text" in modalities:
              if "asr_indices" in batch:
                  batch["encoder_inputs"]["text"] = batch["asr_indices"]
          if "features" in modalities:
              batch["encoder_inputs"]["features"] = batch["features"]
      return batch
  
imports: import functools
  from typing import Dict, Iterator, List, Optional, Text, Tuple, Union
  from absl import logging
  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.projects.t5 import tokenizer
  from scenic.projects.vid2seq import data_utils
  from scenic.projects.vivit.data import video_tfrecord_dataset
  import tensorflow
label_desc: Use image node_id 2 for calling a global function with example usage: map_input_keys(batch, modalities, return_as_dict) and returns: batch