code: def train_and_evaluate(
      config: ml_collections.ConfigDict, workdir: str
  ):
      """Runs a training and evaluation loop.
  
      Args:
        config: Configuration to use.
        workdir: Working directory for checkpoints and TF summaries. If this
          contains checkpoint training will be resumed from the latest checkpoint.
      """
      tf.io.gfile.makedirs(workdir)
  
      vocab_path = config.vocab_path
      if vocab_path is None:
          vocab_path = os.path.join(workdir, "sentencepiece_model")
          config.vocab_path = vocab_path
      tf.io.gfile.makedirs(os.path.split(vocab_path)[0])
  
      # Load Dataset
      # ---------------------------------------------------------------------------
      logging.info("Initializing dataset.")
      train_ds, eval_ds, _, encoder = input_pipeline.get_datasets(
          n_devices=jax.local_device_count(),
          config=config,
          vocab_path=vocab_path,
      )
  
      train_iter = iter(train_ds)
      vocab_size = int(encoder.vocab_size())
      eos_id = (
          temperature_sampler.EOS_ID
      )  # Default Sentencepiece EOS token.
  
      def decode_tokens(toks):
          valid_toks = toks[: np.argmax(toks == eos_id) + 1].astype(
              np.int32
          )
          return encoder.detokenize(valid_toks).numpy().decode("utf-8")
  
      def encode_strings(strs, max_len):
          tokenized_batch = np.zeros((len(strs), max_len), np.int32)
          for i, s in enumerate(strs):
              toks = encoder.tokenize(s).numpy()
              # Remove EOS token in prompt.
              tokenized_batch[i, : toks.shape[0] - 1] = toks[:-1]
          return tokenized_batch
  
      tokenized_prompts = encode_strings(
          [config.prompts], config.max_predict_length
      )
  
      logging.info("Initializing model, optimizer, and step functions.")
      # Build Model and Optimizer
      # ---------------------------------------------------------------------------
      train_config = models.TransformerConfig(
          vocab_size=vocab_size,
          output_vocab_size=vocab_size,
          logits_via_embedding=config.logits_via_embedding,
          dtype=jnp.bfloat16 if config.use_bfloat16 else jnp.float32,
          emb_dim=config.emb_dim,
          num_heads=config.num_heads,
          num_layers=config.num_layers,
          qkv_dim=config.qkv_dim,
          mlp_dim=config.mlp_dim,
          max_len=max(
              config.max_target_length, config.max_eval_target_length
          ),
          dropout_rate=config.dropout_rate,
          attention_dropout_rate=config.attention_dropout_rate,
          deterministic=False,
          decode=False,
          kernel_init=nn.initializers.xavier_uniform(),
          bias_init=nn.initializers.normal(stddev=1e-6),
      )
      eval_config = train_config.replace(deterministic=True)
      predict_config = train_config.replace(
          deterministic=True, decode=True
      )
  
      # Mesh definition
      devices_array = utils.create_device_mesh(config)
      mesh = Mesh(devices_array, config.mesh_axes)
  
      start_step = 0
      rng = jax.random.PRNGKey(config.seed)
      rng, init_rng = jax.random.split(rng)
      rng, inference_rng = random.split(rng)
  
      m = models.TransformerLM(eval_config)
  
      learning_rate_fn = create_learning_rate_schedule(
          learning_rate=config.learning_rate,
          warmup_steps=config.warmup_steps,
      )
  
      optimizer = optax.adamw(
          learning_rate_fn,
          b1=0.9,
          b2=0.98,
          eps=1e-9,
          weight_decay=config.weight_decay,
      )
  
      state, state_mesh_annotations = utils.setup_initial_state(
          m, optimizer, config, init_rng, mesh
      )
      data_sharding = NamedSharding(mesh, P(config.data_sharding))
  
      if config.restore_checkpoints:
          # Restore unreplicated optimizer + model state from last checkpoint.
          state = checkpoints.restore_checkpoint(workdir, state)
          # Grab last step.
          start_step = int(state.step)
  
      writer = metric_writers.create_default_writer(
          workdir, just_logging=jax.process_index() > 0
      )
      if start_step == 0:
          writer.write_hparams(dict(config))
  
      # compile multidevice versions of train/eval/predict step fn.
      jit_train_step = jax.jit(
          train_step,
          in_shardings=(
              state_mesh_annotations,
              data_sharding,
              None,
          ),  # type: ignore
          out_shardings=(state_mesh_annotations, None),  # type: ignore
          static_argnums=(2, 3, 4),
          donate_argnums=0,
      )
  
      jit_eval_step = jax.jit(
          eval_step,
          in_shardings=(
              state_mesh_annotations.params,
              data_sharding,
          ),  # type: ignore
          out_shardings=None,  # type: ignore
          static_argnums=(2, 3),
      )
  
      # Since the inputs and rngkey args for predict_step will be batched,
      # we must vmap them, otherwise the global arrays will be seen in each device
      jit_pred_step = jax.jit(
          jax.vmap(
              predict_step,
              in_axes=(
                  0,
                  jax.tree_map(lambda x: None, state.params),
                  0,
                  None,
                  None,
                  jax.tree_map(lambda x: None, predict_config),
                  None,
                  None,
              ),
          ),
          in_shardings=(
              data_sharding,
              state_mesh_annotations.params,
              data_sharding,
          ),  # type: ignore
          out_shardings=data_sharding,  # type: ignore
          static_argnums=(3, 4, 5, 6, 7),
      )
  
      # Main Train Loop
      # ---------------------------------------------------------------------------
  
      # We init the first set of dropout PRNG keys, but update it afterwards inside
      # the main pmap"d training update for performance.
      dropout_rngs = rng
  
      logging.info("Starting training loop.")
      hooks = []
      report_progress = periodic_actions.ReportProgress(
          num_train_steps=config.num_train_steps, writer=writer
      )
      if jax.process_index() == 0:
          hooks += [
              report_progress,
              periodic_actions.Profile(
                  logdir=workdir, num_profile_steps=5
              ),
          ]
      train_metrics = []
      with metric_writers.ensure_flushes(writer):
          for step in range(start_step, config.num_train_steps):
              is_last_step = step == config.num_train_steps - 1
  
              # Shard data to devices and do a training step.
              with jax.profiler.StepTraceAnnotation(
                  "train", step_num=step
              ):
                  batch = next(train_iter)
                  batch = jax.tree_map(lambda x: jnp.array(x), batch)
                  state, metrics = jit_train_step(
                      state,
                      batch,
                      train_config,
                      learning_rate_fn,
                      0.0,
                      dropout_rngs,
                  )
                  train_metrics.append(metrics)
  
              # Quick indication that training is happening.
              logging.log_first_n(
                  logging.INFO, "Finished training step %d.", 5, step
              )
              for h in hooks:
                  h(step)
  
              # Periodic metric handling.
              if step % config.eval_every_steps == 0 or is_last_step:
                  with report_progress.timed("training_metrics"):
                      logging.info("Gathering training metrics.")
                      train_metrics = common_utils.stack_forest(
                          train_metrics
                      )
                      lr = train_metrics.pop("learning_rate").mean()
                      metrics_sums = jax.tree_util.tree_map(
                          jnp.sum, train_metrics
                      )
                      denominator = metrics_sums.pop("denominator")
                      summary = jax.tree_util.tree_map(
                          lambda x: x / denominator, metrics_sums
                      )  # pylint: disable=cell-var-from-loop
                      summary["learning_rate"] = lr
                      summary["perplexity"] = jnp.clip(
                          jnp.exp(summary["loss"]), a_max=1.0e4
                      )
                      summary = {
                          "train_" + k: v for k, v in summary.items()
                      }
                      writer.write_scalars(step, summary)
                      train_metrics = []
  
                  with report_progress.timed("eval"):
                      eval_results = evaluate(
                          jit_eval_step=jit_eval_step,
                          params=state.params,
                          eval_ds=eval_ds,
                          num_eval_steps=config.num_eval_steps,
                          config=eval_config,
                      )
                      # (clipped) perplexity after averaging log-perplexitie
                      eval_results["perplexity"] = jnp.clip(
                          jnp.exp(eval_results["loss"]), a_max=1.0e4
                      )
                      writer.write_scalars(
                          step,
                          {
                              "eval_" + k: v
                              for k, v in eval_results.items()
                          },
                      )
  
                  with report_progress.timed("generate_text"):
                      exemplars = generate_prediction(
                          jit_pred_step=jit_pred_step,
                          params=state.params,
                          tokenized_prompts=tokenized_prompts,
                          eos_id=eos_id,
                          inference_rng=inference_rng,
                          decode_tokens=decode_tokens,
                          config=config,
                          predict_config=predict_config,
                      )
                      writer.write_texts(step, {"samples": exemplars})
  
              # Save a checkpoint on one host after every checkpoint_freq steps.
              save_checkpoint = (
                  step % config.checkpoint_every_steps == 0
                  or is_last_step
              )
              if config.save_checkpoints and save_checkpoint:
                  logging.info("Saving checkpoint step %d.", step)
                  with report_progress.timed("checkpoint"):
                      checkpoints.save_checkpoint_multiprocess(
                          workdir, state, step
                      )
  
imports: import collections
  import os
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from flax import linen
  from flax.training import checkpoints
  from flax.training import common_utils
  import jax
  from jax import random
  import jax.numpy
  from jax.sharding import PartitionSpec, Mesh, NamedSharding
  import ml_collections
  import numpy
  import optax
  import tensorflow
  import input_pipeline
  import models
  import temperature_sampler
  import utils
label_desc: Use image node_id 14 for calling a global function with example usage: train_and_evaluate(config, workdir) and returns: encoder, tokenized_batch