code: def evaluate(
      *,
      config: ml_collections.ConfigDict,
      model_cls: type[ImageTextModel],
      dataset: dataset_utils.Dataset,
      rng: Optional[jnp.ndarray] = None,
      workdir: Optional[str] = None,  # pylint: disable=unused-argument
      writer: metric_writers.MetricWriter,
  ) -> Mapping[str, float]:
      """Evaluates a model on zero-shot text-to-visual retrieval."""
      input_spec = get_input_spec(
          dataset_meta_data=dataset.meta_data,
          dataset_configs=config.get("dataset_configs", {}),
          train=False,
      )
  
      is_video = is_video_input(input_spec)
  
      hashable_config = config.copy_and_resolve_references()
      hashable_config = ml_collections.FrozenConfigDict(hashable_config)
  
      # Note that different calls of `_replace` with the same contents will yield
      # the same hash.
      dataset = dataclasses.replace(
          dataset, meta_data=flax.core.freeze(dataset.meta_data)
      )
      model, train_state = _create_model_and_train_state(
          config=hashable_config,
          dataset=dataset,
          model_cls=model_cls,
          input_spec=input_spec,
          is_video=is_video,
          rng=rng,
      )
  
      if config.checkpoint:
          train_state = train_utils.restore_checkpoint(
              workdir, train_state
          )[0]
  
      train_state = jax_utils.replicate(train_state)
  
      eval_step_pmapped = jax.pmap(
          partial_with_cache(
              eval_step,
              model=model,
              is_video=is_video,
              debug=config.get("debug_eval"),
          ),
          axis_name=NUM_DEVICES_AXIS_NAME,
          donate_argnums=(1, 2, 3),
      )
  
      total_steps = get_epoch_steps(config, dataset, split="eval")
  
      encoded_visual_list = []
      encoded_text_list = []
  
      for step, batch in zip(
          tqdm(range(total_steps), desc="Evaluating"),
          dataset.valid_iter,
      ):
          with jax.profiler.StepTraceAnnotation("eval", step_num=step):
              batch.pop("key", None)
  
              if (text := batch.get("text_indices")) is None:
                  text = batch["label"]
              else:
                  text = text.squeeze(axis=-2)
  
              mask = compute_mask(text, config)
  
              encoded_visual, encoded_text = eval_step_pmapped(
                  train_state,
                  batch["inputs"],
                  text,
                  mask,
              )
              encoded_visual = jax_utils.unreplicate(encoded_visual)
              encoded_text = jax_utils.unreplicate(encoded_text)
  
              batch_mask = batch["batch_mask"].reshape(-1).astype(bool)
  
              for (
                  mask_instance,
                  encoded_visual_instance,
                  encoded_text_instance,
              ) in zip(batch_mask, encoded_visual, encoded_text):
                  if mask_instance:
                      encoded_visual_list.append(
                          encoded_visual_instance
                      )
                      encoded_text_list.append(encoded_text_instance)
  
      encoded_visuals = jnp.stack(encoded_visual_list)
      encoded_texts = jnp.stack(encoded_text_list)
  
      # This evaluation may be done multiple times in an eval job, so we jit here.
      # Still, in case it's run only once, it shouldn't take long to jit, as it's
      # only evaluation.
      metrics_summary = jax.jit(
          partial_with_cache(
              _compute_metrics,
              model=model,
              # This is a workaround for checking if we are evaluating the
              # validation split from the dataset used for training. In that case,
              # we want to compute the loss with the same function as in training.
              # Otherwise, we want a retrieval-like loss.
              loss_fn=model.loss_function
              if "retrieval_batch_size" in config
              else loss.nce_loss,
              retrieval_batch_size=config.get("retrieval_batch_size"),
          ),
          backend="cpu",  # To avoid OOM while compiling.
          donate_argnums=(0, 1),
      )(encoded_visuals, encoded_texts)
  
      step = jax_utils.unreplicate(train_state.global_step)
      prefix = config.get("writer_prefix", "valid")
  
      ranks = metrics_summary.pop("ranks")
      if step % (config.get("log_eval_steps", 100) * 10) == 0:
          writer.write_histograms(step, {f"{prefix}/ranks": ranks})
  
      return log_eval_summary(
          writer=writer,
          extra_eval_summary=metrics_summary,
          step=step,
          prefix=prefix,
      )
  
imports: from collections.abc import Callable, Mapping, MutableMapping, Sequence
  import dataclasses
  from typing import Optional
  from clu import metric_writers
  from flax import jax_utils
  import flax.core
  import jax
  import jax.numpy
  import ml_collections
  from scenic.dataset_lib import dataset_utils
  from scenic.projects.lang4video.model import loss
  from scenic.projects.lang4video.model.image_text_model import ImageTextModel
  from scenic.projects.lang4video.trainer.train_utils import axis_name_exists
  from scenic.projects.lang4video.trainer.train_utils import compute_mask
  from scenic.projects.lang4video.trainer.train_utils import get_cached_fn
  from scenic.projects.lang4video.trainer.train_utils import get_epoch_steps
  from scenic.projects.lang4video.trainer.train_utils import get_input_spec
  from scenic.projects.lang4video.trainer.train_utils import init_encoder
  from scenic.projects.lang4video.trainer.train_utils import InputSpec
  from scenic.projects.lang4video.trainer.train_utils import is_video_input
  from scenic.projects.lang4video.trainer.train_utils import log_eval_summary
  from scenic.projects.lang4video.trainer.train_utils import NUM_DEVICES_AXIS_NAME
  from scenic.projects.lang4video.trainer.train_utils import pad_and_batch
  from scenic.projects.lang4video.trainer.train_utils import partial_with_cache
  from scenic.train_lib import train_utils
  from tqdm.auto import tqdm
label_desc: Use image node_id 5 for calling a global function with example usage: evaluate() and returns: log_eval_summary