class_docstr: Dense video captioning model with a video encoder and a text decoder.
docstr: Returns a callable metric function for the model.
  
  Args:
    split: The split for which we calculate the metrics. It should be one of
      the ['train',  'validation', 'test'].
  Returns: A metric function with the following API: ```metrics_fn(preds,
    label, weights)```
docstr_tok: ["Returns","a","callable","metric","function","for","the","model",".","Args",":","split",":","The","split","for","which","we","calculate","the","metrics",".","It","should","be","one","of","the","[","'train","'",",","'validation","'",",","'test","'","]",".","Returns",":","A","metric","function","with","the","following","API",":","``","`","metrics_fn","(","preds",",","label",",","weights",")","``","`"]
code: def get_metrics_fn(
      self, split: Optional[str] = None
  ) -> base_model.MetricFn:
      """Returns a callable metric function for the model.
  
      Args:
        split: The split for which we calculate the metrics. It should be one of
          the ['train',  'validation', 'test'].
      Returns: A metric function with the following API: ```metrics_fn(preds,
        label, weights)```
      """
  
      def token_accuracy(
          logits, batch: Batch
      ) -> Dict[str, Tuple[float, int]]:
          # This metric is computed only during teacher forcing training mode.
          if split != "train":
              return {}
  
          targets = batch["decoder_inputs"]["decoder_target_tokens"]
  
          batch_mask = batch["batch_mask"]
  
          # logits: [batch_size, seq_len, vocab_size]
          # targets: [batch_size, seq_len]
          # batch_mask: [batch_size]
          one_hot_targets = common_utils.onehot(
              targets, logits.shape[-1]
          )
          masks = (
              jnp.greater(targets, 0).astype(jnp.int32)
              * batch_mask[:, None]
          )
  
          n_corrects = base_model_utils.weighted_correctly_classified(
              logits, one_hot_targets, masks
          )
          n_valids = base_model_utils.num_examples(
              logits, one_hot_targets, masks
          )
  
          key = "token_accuracy"
          return {  # pytype: disable=bad-return-type  # jax-ndarray
              key: base_model_utils.psum_metric_normalizer(
                  (n_corrects, n_valids)
              )
          }
  
      return token_accuracy
  
imports: import functools
  from typing import Any, Dict, Mapping, Optional, Tuple
  from absl import logging
  import flax.linen
  from flax.training import common_utils
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  from scenic.projects.baselines import vit
  from scenic.projects.t5 import layers
  from scenic.projects.t5 import model
  from t5x import decoding
  from t5x.models import DecodeFnCallable
label_desc: Use image node_id 3 for calling the DenseVideoCaptioningModel obj's underlying member method code with example usage: obj.get_metrics_fn(split) and returns: token_accuracy, dict, dict