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-based video classification."""
      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)
  
      dataset_configs = config.get("dataset_configs", {})
  
      tokenizer_config = dataset_configs.get("tokenizer", {})
      tokenizer = util.create_tokenizer(tokenizer_config)
      tokenizer.initialize()
  
      class_names = (
          dataset.meta_data.get("classes")
          or dataset_configs["class_names"]
      )
      class_templates = config.get("class_templates", ["{}"])
      classes_with_templates = []
      for class_ in class_names:
          for template in class_templates:
              classes_with_templates.append(template.format(class_))
  
      tokenized_classes = tokenizer.string_tensor_to_indices(
          classes_with_templates,
          prepend_bos=tokenizer_config.get("prepend_bos", False),
          append_eos=tokenizer_config.get("append_eos", False),
          max_num_tokens=dataset_configs.get("max_num_words", 32),
      )
      tokenized_classes = (
          tokenized_classes._numpy()
      )  # pylint: disable=protected-access
      tokenized_classes = tokenized_classes.reshape(
          -1, len(class_templates), tokenized_classes.shape[-1]
      )
  
      mask = compute_mask(tokenized_classes, config)
  
      # We pmap here, to avoid OOM by optimizations. Also, it's best to pmap
      # everything here, in case this is run in an eval job. If not, still pmapping
      # it shouldn't take long.
      logging.info("Encoding the classes as text…")
  
      tokenized_classes, batch_padding_size = pad_array_to_be_divisible(
          tokenized_classes, jax.local_device_count()
      )
      mask, batch_padding_size = pad_array_to_be_divisible(
          mask, jax.local_device_count()
      )
  
      tokenized_classes, mask = dataset_utils.shard(
          (tokenized_classes, mask)
      )
  
      compute_class_text_embeddings_pmapped = jax.pmap(
          partial_with_cache(
              compute_class_text_embeddings,
              model=model,
              batch_size=config.get(
                  "class_batch_size",
                  config.get("eval_batch_size", config.batch_size),
              ),
              debug=hashable_config.get("debug_eval"),
          ),
          axis_name=NUM_DEVICES_AXIS_NAME,
          donate_argnums=(1, 2),
      )
  
      encoded_classes = jax_utils.unreplicate(
          compute_class_text_embeddings_pmapped(
              train_state, tokenized_classes, mask
          )
      )
  
      encoded_classes = encoded_classes[
          : len(encoded_classes) - batch_padding_size
      ]
  
      logging.info("Classes encoded.")
  
      eval_step_pmapped = jax.pmap(
          # This function would fail to cache because `encoded_classes` is a
          # `DeviceArray`. Besides, this value would change for different params.
          # So we can't cache it.
          functools.partial(
              eval_step,
              model=model,
              encoded_classes=encoded_classes,
              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")
  
      eval_metrics_all = []
  
      for step, batch in zip(
          tqdm(range(total_steps), desc="Evaluating"),
          dataset.valid_iter,
      ):
          with jax.profiler.StepTraceAnnotation("eval", step_num=step):
              eval_metrics_batch = eval_step_pmapped(
                  train_state,
                  batch["inputs"],
                  batch["label"],
                  batch["batch_mask"],
              )
              eval_metrics_all.append(eval_metrics_batch)
  
      return log_eval_summary(
          writer=writer,
          eval_metrics=train_utils.unreplicate_and_get(
              eval_metrics_all
          ),
          step=jax_utils.unreplicate(train_state.global_step),
          prefix=config.get("writer_prefix", "valid"),
      )
  
imports: from collections.abc import Mapping
  import dataclasses
  import functools
  from typing import Optional
  from absl import logging
  from clu import metric_writers
  from flax import jax_utils
  import flax.core
  import jax
  import jax.numpy
  import ml_collections
  import optax
  from scenic.dataset_lib import dataset_utils
  from scenic.model_lib.base_models import model_utils
  from scenic.projects.lang4video import util
  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_array_to_be_divisible
  from scenic.projects.lang4video.trainer.train_utils import partial_with_cache
  from scenic.projects.lang4video.trainer.train_utils import split_in_batches
  from scenic.train_lib import train_utils
  from tqdm.auto import tqdm
label_desc: Use image node_id 4 for calling a global function with example usage: evaluate() and returns: log_eval_summary