code: def parse_args():
      # Parse command line arguments
      parser = argparse.ArgumentParser(
          description="""
      To use, download the full checkpoint archive, extract and upload to a GCS bucket, and set that as --tune-model-path
      Modify the config file:
          - set `model_dir` to where the checkpoints should be written during training
          - set `train_set`, `val_set` to index files for your data
          - set `tpu_size` to 8 (if on a v3-8)
          - set `warmup_steps`, `anneal_steps`, `lr`, `end_lr` to the lr schedule for your finetuning run
          - the global step will reset to 0, keep that in mind when writing your lr schedule
          - set `name` to specify the name of the Weights & Biases run
          - set `wandb_project` to specify the Weights & Biases project to log to
      To prepare data in the expected data format:
          - use the script `create_finetune_tfrecords.py` in this repo to create data in the expected format
          - upload the .tfrecords files to GCS
          - save their GCS paths to a index file under `data/`, see existing files for examples
      """,
          formatter_class=argparse.RawTextHelpFormatter,
      )
      parser.add_argument(
          "--config",
          type=str,
          default=None,
          help="Config file location",
      )
      parser.add_argument(
          "--tune-model-path",
          type=str,
          default=None,
          help="Base model to finetune",
      )
      parser.add_argument(
          "--fresh-opt",
          default=False,
          action="store_true",
          help="Use a newly initialized optimizer, ignoring any optimizer state saved in the base checkpoint",
      )
  
      args = parser.parse_args()
      return args
  
file_path: mesh-transformer-jax/device_train.py
imports: import argparse
  import json
  import time
  import jax
  import numpy
  import optax
  import wandb
  from tqdm import tqdm
  from mesh_transformer import util
  from mesh_transformer.checkpoint import read_ckpt, write_ckpt
  from mesh_transformer.transformer_shard import CausalTransformer
  from tfrecord_loader import TFRecordNewInputs
  from smart_open import open
  from google.cloud import storage
  from google.cloud.exceptions import NotFound
  from mesh_transformer.util import clip_by_global_norm, additive_weight_decay
label_desc: Use image node_id 1 for calling a global function with example usage: parse_args() and returns: args