code: def get_eval_step(flax_model, loss_and_metrics_fn, debug=False):
      """Runs a single step of training.
  
      Note that in this code, the buffer of the second argument (batch) is donated
      to the computation.
  
      Args:
        flax_model: Flax model (an instance of nn.Module).
        loss_and_metrics_fn: A function that given model predictions, a batch, and
          parameters of the model calculates the loss as well as metrics.
        debug: bool; Whether the debug mode is enabled during evaluation.
          `debug=True` enables model specific logging/storing some values using
          jax.host_callback.
  
      Returns:
        Eval step function which returns validation metrics.
      """
  
      def metrics_fn(train_state, batch, predictions):
          _, metrics = loss_and_metrics_fn(
              predictions, batch, model_params=train_state.params
          )
          return metrics
  
      def eval_step(train_state, batch):
          predictions = flax_model.apply(
              {"params": train_state.params, **train_state.model_state},
              batch["inputs"],
              train=False,
              debug=debug,
              text_queries=batch["queries"],
          )
          return metrics_fn(train_state, batch, predictions)
  
      return eval_step
  
imports: from concurrent import futures
  import functools
  import time
  from typing import Any
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from clu import platform
  import flax
  from flax import jax_utils
  from flax.training import checkpoints
  import jax
  import jax.numpy
  import jax.profiler
  import ml_collections
  import numpy
  import optax
  from scenic.dataset_lib import dataset_utils
  from scenic.projects.owl_vit import utils
  from scenic.train_lib import optax
  from scenic.train_lib import train_utils
label_desc: Use image node_id 2 for calling a global function with example usage: get_eval_step(flax_model, loss_and_metrics_fn, debug) and returns: eval_step, metrics, metrics_fn