code: def _create_model_and_train_state(
      config: ml_collections.ConfigDict,
      dataset: dataset_utils.Dataset,
      model_cls: type[ImageTextModel],
      input_spec: InputSpec,
      is_video: bool,
      rng: Optional[jnp.ndarray] = None,
  ) -> tuple[ImageTextModel, train_utils.TrainState]:
      """Creates the model and train state."""
      model = model_cls(config, dataset.meta_data)
      encoder = model.flax_model
  
      params, model_state = init_encoder(
          encoder=encoder,
          input_spec=input_spec,
          method=encoder.encode_video_and_text if is_video else None,
          config=config,
          rng=rng,
      )
  
      train_state = train_utils.TrainState(
          params=params, model_state=model_state
      )
  
      return model, train_state
  
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 3 for calling a global function with example usage: _create_model_and_train_state(config, dataset, model_cls, input_spec, is_video, rng) and returns: model,  train_state