code: def get_model_cls(model_name: str) -> Any:
      """Returns model class given its name."""
      if model_name == "xvit_classification":
          return xvit.XViTClassificationModel
      elif model_name == "vit_classification":
          return vit.ViTClassificationModel
      elif model_name == "topological_vit_classification":
          return vit.TopologicalViTClassificationModel
      elif model_name == "duplex_vit_classification":
          return duplex_vit.DuplexViTClassificationModel
      else:
          raise ValueError(f"Unrecognized model: {model_name}.")
  
imports: from typing import Any
  from absl import flags
  from clu import metric_writers
  import jax
  import jax.numpy
  import ml_collections
  from scenic import app
  from scenic.projects.tasseo import classification_trainer
  from scenic.projects.tasseo import duplex_vit
  from scenic.projects.tasseo import duplex_vit_classification_trainer
  from scenic.projects.tasseo import inference
  from scenic.projects.tasseo import transfer_trainer
  from scenic.projects.tasseo import vit
  from scenic.projects.tasseo import xvit
  from scenic.projects.tasseo.datasets import abnormality_baseline_dataset
  from scenic.projects.tasseo.datasets import chrmID_baseline_dataset
  from scenic.projects.tasseo.datasets import chrmID_big_metaphase_context_dataset
  from scenic.projects.tasseo.datasets import chrmID_dataset
  from scenic.projects.tasseo.datasets import chrmID_metaphase_context_dataset
  from scenic.projects.tasseo.datasets import longtail_baseline_dataset
  from scenic.projects.tasseo.datasets import longtail_rhs_baseline_dataset
  from scenic.projects.tasseo.datasets import metaphase_sexid_dataset
  from scenic.train_lib import train_utils
  from scenic.train_lib import trainers
label_desc: Use image node_id 1 for calling a global function with example usage: get_model_cls(model_name) and returns: xvit, vit, vit, duplex_vit