code: def get_call_model_func(
      model: lit_model.Model,
      model_input: JsonDict,
      image_field_key: str,
      grad_field_key: str,
      grad_target_field_key: str,
      grad_target_label: str,
  ) -> CallModelFunction:
      """Returns a function that is used by the Saliency library to get gradients.
  
      Args:
        model: LIT model that is used to calculate actual gradients.
        model_input: the model input.
        image_field_key: the name (key) of the field in the model input that
          contains the image data with respect to which the gradients should be
          calculated.
        grad_field_key: the name (key) of the field in the model output that
          contains the computed gradients.
        grad_target_field_key: the name (key) of the field in the model input that
          is used to specify the label for which the gradients should be calculated.
          If the value is None then it is a regression or a single class
          classification model that has only one output.
        grad_target_label: the value of the label that should be passed as the
          `grad_target_field_name` value.
  
      Returns:
        The function that should be passed to the Saliency library.
      """
  
      def call_model_func(
          x_value_batch: np.ndarray,
          call_model_args,
          expected_keys: list[str],
      ) -> dict[str, np.ndarray]:
          """This function is called by the Saliency library to calculate gradients.
  
          Args:
            x_value_batch: a batch of inputs with respect to which the gradients
              should be calculated.
            call_model_args: unused.
            expected_keys: the list of expected keys that the return value should
              contain.
  
          Returns:
            A dictionary with gradients values for the batch.
          """
          del call_model_args  # Unused.
  
          # Iterate through the batch of saliency lib inputs and convert them to
          # a batch acceptable by the LIT model.
          model_inputs = []
          for x_value in x_value_batch:
              updates = {image_field_key: x_value}
              if grad_target_field_key is not None:
                  updates[grad_target_field_key] = grad_target_label
              input_copy = lit_utils.make_modified_input(
                  model_input, updates, "ImageSalience"
              )
              model_inputs.append(input_copy)
  
          # Call the model to obtain gradients.
          predictions = model.predict(inputs=model_inputs)
  
          # Gradient results formatted for the saliency library.
          gradients_batch = [p[grad_field_key] for p in predictions]
  
          assert (
              saliency.core.base.OUTPUT_LAYER_VALUES
              not in expected_keys
          )
  
          # Convert the result to the format acceptable by the saliency library.
          return {
              saliency.core.base.INPUT_OUTPUT_GRADIENTS: np.asarray(
                  gradients_batch
              ),
          }
  
      return call_model_func
  
imports: import abc
  from collections.abc import Callable
  from typing import Any, NamedTuple, Optional, cast
  from absl import logging
  from lit_nlp.api import components
  from lit_nlp.api import dataset
  from lit_nlp.api import model
  from lit_nlp.api import types
  from lit_nlp.lib import image_utils
  from lit_nlp.lib import utils
  import numpy
  import saliency
label_desc: Use image node_id 2 for calling a global function with example usage: get_call_model_func(model, model_input, image_field_key, grad_field_key, grad_target_field_key, grad_target_label) and returns: call_model_func, dict