code: def main():
      # See all possible arguments in src/transformers/training_args.py
      # or by passing the --help flag to this script.
      # We now keep distinct sets of args, for a cleaner separation of concerns.
  
      parser = HfArgumentParser(
          (ModelArguments, DataTrainingArguments, TrainingArguments)
      )
      if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
          # If we pass only one argument to the script and it's the path to a json file,
          # let's parse it to get our arguments.
          model_args, data_args, training_args = parser.parse_json_file(
              json_file=os.path.abspath(sys.argv[1])
          )
      else:
          (
              model_args,
              data_args,
              training_args,
          ) = parser.parse_args_into_dataclasses()
  
      # Sending telemetry. Tracking the example usage helps us better allocate resources to maintain them. The
      # information sent is the one passed as arguments along with your Python/PyTorch versions.
      send_example_telemetry(
          "run_summarization", model_args, data_args, framework="flax"
      )
  
      if (
          os.path.exists(training_args.output_dir)
          and os.listdir(training_args.output_dir)
          and training_args.do_train
          and not training_args.overwrite_output_dir
      ):
          raise ValueError(
              f"Output directory ({training_args.output_dir}) already exists and is not empty."
              "Use --overwrite_output_dir to overcome."
          )
  
      # Make one log on every process with the configuration for debugging.
      logging.basicConfig(
          format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
          datefmt="%m/%d/%Y %H:%M:%S",
          level=logging.INFO,
      )
      # Setup logging, we only want one process per machine to log things on the screen.
      logger.setLevel(
          logging.INFO if jax.process_index() == 0 else logging.ERROR
      )
      if jax.process_index() == 0:
          datasets.utils.logging.set_verbosity_warning()
          transformers.utils.logging.set_verbosity_info()
      else:
          datasets.utils.logging.set_verbosity_error()
          transformers.utils.logging.set_verbosity_error()
  
      # Set the verbosity to info of the Transformers logger (on main process only):
      logger.info(f"Training/evaluation parameters {training_args}")
  
      # Handle the repository creation
      if training_args.push_to_hub:
          if training_args.hub_model_id is None:
              repo_name = get_full_repo_name(
                  Path(training_args.output_dir).absolute().name,
                  token=training_args.hub_token,
              )
          else:
              repo_name = training_args.hub_model_id
          repo = Repository(
              training_args.output_dir, clone_from=repo_name
          )
  
      # Get the datasets: you can either provide your own CSV/JSON training and evaluation files (see below)
      # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
      # (the dataset will be downloaded automatically from the datasets Hub).
      #
      # For CSV/JSON files this script will use the first column for the full texts and the second column for the
      # summaries (unless you specify column names for this with the `text_column` and `summary_column` arguments).
      #
      if data_args.dataset_name is not None:
          # Downloading and loading a dataset from the hub.
          dataset = load_dataset(
              data_args.dataset_name,
              data_args.dataset_config_name,
              cache_dir=model_args.cache_dir,
              keep_in_memory=False,
              use_auth_token=True
              if model_args.use_auth_token
              else None,
          )
      else:
          data_files = {}
          if data_args.train_file is not None:
              data_files["train"] = data_args.train_file
              extension = data_args.train_file.split(".")[-1]
          if data_args.validation_file is not None:
              data_files["validation"] = data_args.validation_file
              extension = data_args.validation_file.split(".")[-1]
          if data_args.test_file is not None:
              data_files["test"] = data_args.test_file
              extension = data_args.test_file.split(".")[-1]
          dataset = load_dataset(
              extension,
              data_files=data_files,
              cache_dir=model_args.cache_dir,
              use_auth_token=True
              if model_args.use_auth_token
              else None,
          )
      # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
      # https://huggingface.co/docs/datasets/loading_datasets.html.
  
      # Load pretrained model and tokenizer
  
      if model_args.config_name:
          config = AutoConfig.from_pretrained(
              model_args.config_name,
              cache_dir=model_args.cache_dir,
              use_auth_token=True
              if model_args.use_auth_token
              else None,
          )
      elif model_args.model_name_or_path:
          config = AutoConfig.from_pretrained(
              model_args.model_name_or_path,
              cache_dir=model_args.cache_dir,
              use_auth_token=True
              if model_args.use_auth_token
              else None,
          )
      else:
          config = CONFIG_MAPPING[model_args.model_type]()
          logger.warning(
              "You are instantiating a new config instance from scratch."
          )
  
      if model_args.tokenizer_name:
          tokenizer = AutoTokenizer.from_pretrained(
              model_args.tokenizer_name,
              cache_dir=model_args.cache_dir,
              use_fast=model_args.use_fast_tokenizer,
              use_auth_token=True
              if model_args.use_auth_token
              else None,
          )
      elif model_args.model_name_or_path:
          tokenizer = AutoTokenizer.from_pretrained(
              model_args.model_name_or_path,
              cache_dir=model_args.cache_dir,
              use_fast=model_args.use_fast_tokenizer,
              use_auth_token=True
              if model_args.use_auth_token
              else None,
          )
      else:
          raise ValueError(
              "You are instantiating a new tokenizer from scratch. This is not supported by this script."
              "You can do it from another script, save it, and load it from here, using --tokenizer_name."
          )
  
      if model_args.model_name_or_path:
          model = FlaxAutoModelForSeq2SeqLM.from_pretrained(
              model_args.model_name_or_path,
              config=config,
              seed=training_args.seed,
              dtype=getattr(jnp, model_args.dtype),
              use_auth_token=True
              if model_args.use_auth_token
              else None,
          )
      else:
          model = FlaxAutoModelForSeq2SeqLM.from_config(
              config,
              seed=training_args.seed,
              dtype=getattr(jnp, model_args.dtype),
          )
  
      if training_args.gradient_checkpointing:
          model.enable_gradient_checkpointing()
  
      if model.config.decoder_start_token_id is None:
          raise ValueError(
              "Make sure that `config.decoder_start_token_id` is correctly defined"
          )
  
      prefix = (
          data_args.source_prefix
          if data_args.source_prefix is not None
          else ""
      )
  
      # Preprocessing the datasets.
      # We need to tokenize inputs and targets.
      if training_args.do_train:
          column_names = dataset["train"].column_names
      elif training_args.do_eval:
          column_names = dataset["validation"].column_names
      elif training_args.do_predict:
          column_names = dataset["test"].column_names
      else:
          logger.info(
              "There is nothing to do. Please pass `do_train`, `do_eval` and/or `do_predict`."
          )
          return
  
      # Get the column names for input/target.
      dataset_columns = summarization_name_mapping.get(
          data_args.dataset_name, None
      )
      if data_args.text_column is None:
          text_column = (
              dataset_columns[0]
              if dataset_columns is not None
              else column_names[0]
          )
      else:
          text_column = data_args.text_column
          if text_column not in column_names:
              raise ValueError(
                  f"--text_column' value '{data_args.text_column}' needs to be one of: {', '.join(column_names)}"
              )
      if data_args.summary_column is None:
          summary_column = (
              dataset_columns[1]
              if dataset_columns is not None
              else column_names[1]
          )
      else:
          summary_column = data_args.summary_column
          if summary_column not in column_names:
              raise ValueError(
                  f"--summary_column' value '{data_args.summary_column}' needs to be one of: {', '.join(column_names)}"
              )
  
      # Temporarily set max_target_length for training.
      max_target_length = data_args.max_target_length
  
      # In Flax, for seq2seq models we need to pass `decoder_input_ids`
      # as the Flax models don't accept `labels`, we need to prepare the decoder_input_ids here
      # for that dynamically import the `shift_tokens_right` function from the model file
      model_module = __import__(
          model.__module__, fromlist=["shift_tokens_tight"]
      )
      shift_tokens_right_fn = getattr(
          model_module, "shift_tokens_right"
      )
  
      # Setting padding="max_length" as we need fixed length inputs for jitted functions
      def preprocess_function(examples):
          inputs = examples[text_column]
          targets = examples[summary_column]
          inputs = [prefix + inp for inp in inputs]
          model_inputs = tokenizer(
              inputs,
              max_length=data_args.max_source_length,
              padding="max_length",
              truncation=True,
              return_tensors="np",
          )
  
          # Setup the tokenizer for targets
          labels = tokenizer(
              text_target=targets,
              max_length=max_target_length,
              padding="max_length",
              truncation=True,
              return_tensors="np",
          )
  
          model_inputs["labels"] = labels["input_ids"]
          decoder_input_ids = shift_tokens_right_fn(
              labels["input_ids"],
              config.pad_token_id,
              config.decoder_start_token_id,
          )
          model_inputs["decoder_input_ids"] = np.asarray(
              decoder_input_ids
          )
  
          # We need decoder_attention_mask so we can ignore pad tokens from loss
          model_inputs["decoder_attention_mask"] = labels[
              "attention_mask"
          ]
  
          return model_inputs
  
      if training_args.do_train:
          if "train" not in dataset:
              raise ValueError("--do_train requires a train dataset")
          train_dataset = dataset["train"]
          if data_args.max_train_samples is not None:
              max_train_samples = min(
                  len(train_dataset), data_args.max_train_samples
              )
              train_dataset = train_dataset.select(
                  range(max_train_samples)
              )
          train_dataset = train_dataset.map(
              preprocess_function,
              batched=True,
              num_proc=data_args.preprocessing_num_workers,
              remove_columns=column_names,
              load_from_cache_file=not data_args.overwrite_cache,
              desc="Running tokenizer on train dataset",
          )
  
      if training_args.do_eval:
          max_target_length = data_args.val_max_target_length
          if "validation" not in dataset:
              raise ValueError(
                  "--do_eval requires a validation dataset"
              )
          eval_dataset = dataset["validation"]
          if data_args.max_eval_samples is not None:
              max_eval_samples = min(
                  len(eval_dataset), data_args.max_eval_samples
              )
              eval_dataset = eval_dataset.select(
                  range(max_eval_samples)
              )
          eval_dataset = eval_dataset.map(
              preprocess_function,
              batched=True,
              num_proc=data_args.preprocessing_num_workers,
              remove_columns=column_names,
              load_from_cache_file=not data_args.overwrite_cache,
              desc="Running tokenizer on validation dataset",
          )
  
      if training_args.do_predict:
          max_target_length = data_args.val_max_target_length
          if "test" not in dataset:
              raise ValueError("--do_predict requires a test dataset")
          predict_dataset = dataset["test"]
          if data_args.max_predict_samples is not None:
              max_predict_samples = min(
                  len(predict_dataset), data_args.max_predict_samples
              )
              predict_dataset = predict_dataset.select(
                  range(max_predict_samples)
              )
          predict_dataset = predict_dataset.map(
              preprocess_function,
              batched=True,
              num_proc=data_args.preprocessing_num_workers,
              remove_columns=column_names,
              load_from_cache_file=not data_args.overwrite_cache,
              desc="Running tokenizer on prediction dataset",
          )
  
      # Metric
      metric = evaluate.load("rouge")
  
      def postprocess_text(preds, labels):
          preds = [pred.strip() for pred in preds]
          labels = [label.strip() for label in labels]
  
          # rougeLSum expects newline after each sentence
          preds = [
              "\n".join(nltk.sent_tokenize(pred)) for pred in preds
          ]
          labels = [
              "\n".join(nltk.sent_tokenize(label)) for label in labels
          ]
  
          return preds, labels
  
      def compute_metrics(preds, labels):
          decoded_preds = tokenizer.batch_decode(
              preds, skip_special_tokens=True
          )
          decoded_labels = tokenizer.batch_decode(
              labels, skip_special_tokens=True
          )
  
          # Some simple post-processing
          decoded_preds, decoded_labels = postprocess_text(
              decoded_preds, decoded_labels
          )
  
          result = metric.compute(
              predictions=decoded_preds,
              references=decoded_labels,
              use_stemmer=True,
          )
          result = {k: round(v * 100, 4) for k, v in result.items()}
          prediction_lens = [
              np.count_nonzero(pred != tokenizer.pad_token_id)
              for pred in preds
          ]
          result["gen_len"] = np.mean(prediction_lens)
          return result
  
      # Enable tensorboard only on the master node
      has_tensorboard = is_tensorboard_available()
      if has_tensorboard and jax.process_index() == 0:
          try:
              from flax.metrics.tensorboard import SummaryWriter
  
              summary_writer = SummaryWriter(
                  log_dir=Path(training_args.output_dir)
              )
          except ImportError as ie:
              has_tensorboard = False
              logger.warning(
                  f"Unable to display metrics through TensorBoard because some package are not installed: {ie}"
              )
      else:
          logger.warning(
              "Unable to display metrics through TensorBoard because the package is not installed: "
              "Please run pip install tensorboard to enable."
          )
  
      # Initialize our training
      rng = jax.random.PRNGKey(training_args.seed)
      rng, dropout_rng = jax.random.split(rng)
  
      # Store some constant
      num_epochs = int(training_args.num_train_epochs)
      train_batch_size = (
          int(training_args.per_device_train_batch_size)
          * jax.device_count()
      )
      per_device_eval_batch_size = int(
          training_args.per_device_eval_batch_size
      )
      eval_batch_size = per_device_eval_batch_size * jax.device_count()
      steps_per_epoch = len(train_dataset) // train_batch_size
      total_train_steps = steps_per_epoch * num_epochs
  
      # Create learning rate schedule
      linear_decay_lr_schedule_fn = create_learning_rate_fn(
          len(train_dataset),
          train_batch_size,
          training_args.num_train_epochs,
          training_args.warmup_steps,
          training_args.learning_rate,
      )
  
      # We use Optax's "masking" functionality to not apply weight decay
      # to bias and LayerNorm scale parameters. decay_mask_fn returns a
      # mask boolean with the same structure as the parameters.
      # The mask is True for parameters that should be decayed.
      def decay_mask_fn(params):
          flat_params = traverse_util.flatten_dict(params)
          # find out all LayerNorm parameters
          layer_norm_candidates = ["layernorm", "layer_norm", "ln"]
          layer_norm_named_params = set(
              [
                  layer[-2:]
                  for layer_norm_name in layer_norm_candidates
                  for layer in flat_params.keys()
                  if layer_norm_name in "".join(layer).lower()
              ]
          )
          flat_mask = {
              path: (
                  path[-1] != "bias"
                  and path[-2:] not in layer_norm_named_params
              )
              for path in flat_params
          }
          return traverse_util.unflatten_dict(flat_mask)
  
      # create adam optimizer
      adamw = optax.adamw(
          learning_rate=linear_decay_lr_schedule_fn,
          b1=training_args.adam_beta1,
          b2=training_args.adam_beta2,
          eps=training_args.adam_epsilon,
          weight_decay=training_args.weight_decay,
          mask=decay_mask_fn,
      )
  
      # Setup train state
      state = TrainState.create(
          apply_fn=model.__call__,
          params=model.params,
          tx=adamw,
          dropout_rng=dropout_rng,
      )
  
      # label smoothed cross entropy
      def loss_fn(
          logits, labels, padding_mask, label_smoothing_factor=0.0
      ):
          """
          The label smoothing implementation is adapted from Flax's official example:
          https://github.com/google/flax/blob/87a211135c6a377c8f29048a1cac3840e38b9da4/examples/wmt/train.py#L104
          """
          vocab_size = logits.shape[-1]
          confidence = 1.0 - label_smoothing_factor
          low_confidence = (1.0 - confidence) / (vocab_size - 1)
          normalizing_constant = -(
              confidence * jnp.log(confidence)
              + (vocab_size - 1)
              * low_confidence
              * jnp.log(low_confidence + 1e-20)
          )
          soft_labels = onehot(
              labels,
              vocab_size,
              on_value=confidence,
              off_value=low_confidence,
          )
  
          loss = optax.softmax_cross_entropy(logits, soft_labels)
          loss = loss - normalizing_constant
  
          # ignore padded tokens from loss
          loss = loss * padding_mask
          loss = loss.sum()
          num_labels = padding_mask.sum()
          return loss, num_labels
  
      # Define gradient update step fn
      def train_step(state, batch, label_smoothing_factor=0.0):
          dropout_rng, new_dropout_rng = jax.random.split(
              state.dropout_rng
          )
  
          def compute_loss(params):
              labels = batch.pop("labels")
              logits = state.apply_fn(
                  **batch,
                  params=params,
                  dropout_rng=dropout_rng,
                  train=True,
              )[0]
              loss, num_labels = loss_fn(
                  logits,
                  labels,
                  batch["decoder_attention_mask"],
                  label_smoothing_factor,
              )
              return loss, num_labels
  
          grad_fn = jax.value_and_grad(compute_loss, has_aux=True)
          (loss, num_labels), grad = grad_fn(state.params)
          num_labels = jax.lax.psum(num_labels, "batch")
  
          # true loss = total loss / total samples
          loss = jax.lax.psum(loss, "batch")
          loss = jax.tree_util.tree_map(lambda x: x / num_labels, loss)
  
          # true grad = total grad / total samples
          grad = jax.lax.psum(grad, "batch")
          grad = jax.tree_util.tree_map(lambda x: x / num_labels, grad)
          new_state = state.apply_gradients(
              grads=grad, dropout_rng=new_dropout_rng
          )
  
          metrics = {
              "loss": loss,
              "learning_rate": linear_decay_lr_schedule_fn(state.step),
          }
          return new_state, metrics
  
      # Define eval fn
      def eval_step(params, batch, label_smoothing_factor=0.0):
          labels = batch.pop("labels")
          logits = model(**batch, params=params, train=False)[0]
  
          loss, num_labels = loss_fn(
              logits,
              labels,
              batch["decoder_attention_mask"],
              label_smoothing_factor,
          )
          num_labels = jax.lax.psum(num_labels, "batch")
  
          # true loss = total loss / total samples
          loss = jax.lax.psum(loss, "batch")
          loss = jax.tree_util.tree_map(lambda x: x / num_labels, loss)
  
          metrics = {"loss": loss}
          return metrics
  
      # Define generation function
      max_length = (
          data_args.val_max_target_length
          if data_args.val_max_target_length is not None
          else model.config.max_length
      )
      num_beams = (
          data_args.num_beams
          if data_args.num_beams is not None
          else model.config.num_beams
      )
      gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
  
      def generate_step(params, batch):
          model.params = params
          output_ids = model.generate(
              batch["input_ids"],
              attention_mask=batch["attention_mask"],
              **gen_kwargs,
          )
          return output_ids.sequences
  
      # Create parallel version of the train and eval step
      p_train_step = jax.pmap(
          partial(
              train_step,
              label_smoothing_factor=training_args.label_smoothing_factor,
          ),
          "batch",
          donate_argnums=(0,),
      )
      p_eval_step = jax.pmap(
          partial(
              eval_step,
              label_smoothing_factor=training_args.label_smoothing_factor,
          ),
          "batch",
      )
      p_generate_step = jax.pmap(generate_step, "batch")
  
      # Replicate the train state on each device
      state = state.replicate()
  
      logger.info("***** Running training *****")
      logger.info(f"  Num examples = {len(train_dataset)}")
      logger.info(f"  Num Epochs = {num_epochs}")
      logger.info(
          f"  Instantaneous batch size per device = {training_args.per_device_train_batch_size}"
      )
      logger.info(
          f"  Total train batch size (w. parallel & distributed) = {train_batch_size}"
      )
      logger.info(f"  Total optimization steps = {total_train_steps}")
  
      train_time = 0
      epochs = tqdm(
          range(num_epochs),
          desc=f"Epoch ... (1/{num_epochs})",
          position=0,
      )
      for epoch in epochs:
          # ======================== Training ================================
          train_start = time.time()
  
          # Create sampling rng
          rng, input_rng = jax.random.split(rng)
          train_metrics = []
  
          # Generate an epoch by shuffling sampling indices from the train dataset
          train_loader = data_loader(
              input_rng, train_dataset, train_batch_size, shuffle=True
          )
          steps_per_epoch = len(train_dataset) // train_batch_size
          # train
          for _ in tqdm(
              range(steps_per_epoch),
              desc="Training...",
              position=1,
              leave=False,
          ):
              batch = next(train_loader)
              batch = shard(batch)
              state, train_metric = p_train_step(state, batch)
              train_metrics.append(train_metric)
  
          train_time += time.time() - train_start
  
          train_metric = unreplicate(train_metric)
  
          epochs.write(
              f"Epoch... ({epoch + 1}/{num_epochs} | Loss: {train_metric['loss']}, Learning Rate:"
              f" {train_metric['learning_rate']})"
          )
  
          # ======================== Evaluating ==============================
          eval_metrics = []
          eval_preds = []
          eval_labels = []
  
          eval_loader = data_loader(
              input_rng, eval_dataset, eval_batch_size, drop_last=False
          )
          eval_steps = math.ceil(len(eval_dataset) / eval_batch_size)
          for _ in tqdm(
              range(eval_steps),
              desc="Evaluating...",
              position=2,
              leave=False,
          ):
              # Model forward
              batch = next(eval_loader)
              labels = batch["labels"]
  
              metrics = pad_shard_unpad(
                  p_eval_step, static_return=True
              )(
                  state.params,
                  batch,
                  min_device_batch=per_device_eval_batch_size,
              )
              eval_metrics.append(metrics)
  
              # generation
              if data_args.predict_with_generate:
                  generated_ids = pad_shard_unpad(p_generate_step)(
                      state.params, batch
                  )
                  eval_preds.extend(
                      jax.device_get(
                          generated_ids.reshape(
                              -1, gen_kwargs["max_length"]
                          )
                      )
                  )
                  eval_labels.extend(labels)
  
          # normalize eval metrics
          eval_metrics = get_metrics(eval_metrics)
          eval_metrics = jax.tree_util.tree_map(jnp.mean, eval_metrics)
  
          # compute ROUGE metrics
          rouge_desc = ""
          if data_args.predict_with_generate:
              rouge_metrics = compute_metrics(eval_preds, eval_labels)
              eval_metrics.update(rouge_metrics)
              rouge_desc = " ".join(
                  [
                      f"Eval {key}: {value} |"
                      for key, value in rouge_metrics.items()
                  ]
              )
  
          # Print metrics and update progress bar
          desc = f"Epoch... ({epoch + 1}/{num_epochs} | Eval Loss: {eval_metrics['loss']} | {rouge_desc})"
          epochs.write(desc)
          epochs.desc = desc
  
          # Save metrics
          if has_tensorboard and jax.process_index() == 0:
              cur_step = epoch * (
                  len(train_dataset) // train_batch_size
              )
              write_metric(
                  summary_writer,
                  train_metrics,
                  eval_metrics,
                  train_time,
                  cur_step,
              )
  
          # save checkpoint after each epoch and push checkpoint to the hub
          if jax.process_index() == 0:
              params = jax.device_get(
                  jax.tree_util.tree_map(lambda x: x[0], state.params)
              )
              model.save_pretrained(
                  training_args.output_dir, params=params
              )
              tokenizer.save_pretrained(training_args.output_dir)
              if training_args.push_to_hub:
                  repo.push_to_hub(
                      commit_message=f"Saving weights and logs of epoch {epoch}",
                      blocking=False,
                  )
  
      # ======================== Prediction loop ==============================
      if training_args.do_predict:
          logger.info("*** Predict ***")
  
          pred_metrics = []
          pred_generations = []
          pred_labels = []
  
          pred_loader = data_loader(
              input_rng,
              predict_dataset,
              eval_batch_size,
              drop_last=False,
          )
          pred_steps = math.ceil(len(predict_dataset) / eval_batch_size)
          for _ in tqdm(
              range(pred_steps),
              desc="Predicting...",
              position=2,
              leave=False,
          ):
              # Model forward
              batch = next(pred_loader)
              labels = batch["labels"]
  
              metrics = pad_shard_unpad(
                  p_eval_step, static_return=True
              )(
                  state.params,
                  batch,
                  min_device_batch=per_device_eval_batch_size,
              )
              pred_metrics.append(metrics)
  
              # generation
              if data_args.predict_with_generate:
                  generated_ids = pad_shard_unpad(p_generate_step)(
                      state.params, batch
                  )
                  pred_generations.extend(
                      jax.device_get(
                          generated_ids.reshape(
                              -1, gen_kwargs["max_length"]
                          )
                      )
                  )
                  pred_labels.extend(labels)
  
          # normalize prediction metrics
          pred_metrics = get_metrics(pred_metrics)
          pred_metrics = jax.tree_util.tree_map(jnp.mean, pred_metrics)
  
          # compute ROUGE metrics
          rouge_desc = ""
          if data_args.predict_with_generate:
              rouge_metrics = compute_metrics(
                  pred_generations, pred_labels
              )
              pred_metrics.update(rouge_metrics)
              rouge_desc = " ".join(
                  [
                      f"Predict {key}: {value} |"
                      for key, value in rouge_metrics.items()
                  ]
              )
  
          # Print metrics
          desc = f"Predict Loss: {pred_metrics['loss']} | {rouge_desc})"
          logger.info(desc)
  
          # save final metrics in json
          if jax.process_index() == 0:
              rouge_metrics = {
                  f"test_{metric_name}": value
                  for metric_name, value in rouge_metrics.items()
              }
              path = os.path.join(
                  training_args.output_dir, "test_results.json"
              )
              with open(path, "w") as f:
                  json.dump(rouge_metrics, f, indent=4, sort_keys=True)
  
imports: import json
  import logging
  import math
  import os
  import sys
  import time
  from dataclasses import asdict, dataclass, field
  from enum import Enum
  from functools import partial
  from pathlib import Path
  from typing import Callable, Optional
  import datasets
  import nltk
  import numpy
  from datasets import Dataset, load_dataset
  from tqdm import tqdm
  import evaluate
  import jax
  import jax.numpy
  import optax
  import transformers
  from filelock import FileLock
  from flax import jax_utils, traverse_util
  from flax.jax_utils import pad_shard_unpad, unreplicate
  from flax.training import train_state
  from flax.training.common_utils import get_metrics, onehot, shard, shard_prng_key
  from huggingface_hub import Repository
  from transformers import CONFIG_MAPPING, FLAX_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING, AutoConfig, AutoTokenizer, FlaxAutoModelForSeq2SeqLM, HfArgumentParser, is_tensorboard_available
  from transformers.utils import get_full_repo_name, is_offline_mode, send_example_telemetry
label_desc: Use image node_id 4 for calling a global function with example usage: main() and returns: model_inputs, preds,  labels, result, traverse_util, loss,  num_labels, new_state,  metrics, metrics, output_ids, loss,  num_labels