code: def get_model_cls(model_name):
      """ "Selects Vivit model type."""
      if model_name == "vivit_multilabel_classification":
          return ViViTMultilabelClassificationModel
      elif model_name == "vivit_classification":
          return ViViTClassificationModel
      elif model_name == "vivit_multihead_classification":
          return ViViTMultiHeadClassificationModel
      else:
          raise ValueError("Unrecognized model: {}".format(model_name))
  
imports: import functools
  from typing import Any, Optional, Callable, Sequence
  from absl import logging
  import flax.linen
  from flax.linen.linear import default_kernel_init
  from immutabledict import immutabledict
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.common_lib import video_utils
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.base_models.classification_model import ClassificationModel
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.baselines import vit
  from scenic.projects.vivit import model_utils
label_desc: Use image node_id 1 for calling a global function with example usage: get_model_cls(model_name) and returns: ViViTMultilabelClassificationModel, ViViTClassificationModel, ViViTMultiHeadClassificationModel