code: def main():
      args = parse_args()
  
      # Sanity checks
      if not Path(args.checkpoint).exists():
          raise ValueError(
              f"Checkpoint {args.checkpoint} does not exist."
          )
  
      if not args.output_path:
          args.output_path = Path(args.checkpoint) / "merged_model"
          print(f"Output path not specified. Using {args.output_path}")
  
      Path(args.output_path).mkdir(parents=True, exist_ok=True)
  
      # Load orignal model
      s = time.time()
      model_id = f"meta-llama/Llama-2-{args.model_name}-hf"
      s3_bucket = get_mirror_link(model_id)
      ckpt_path, _ = get_checkpoint_and_refs_dir(
          model_id=model_id, bucket_uri=s3_bucket
      )
  
      print(
          f"Downloading original model {model_id} from {s3_bucket} to {ckpt_path} ..."
      )
      print("Loading tokenizer...")
  
      tokenizer = AutoTokenizer.from_pretrained(
          args.checkpoint, legacy=True
      )
      tokenizer.save_pretrained(Path(args.output_path))
  
      print(f"Saved tokenizer to {args.output_path}")
  
      download_model(
          model_id=model_id,
          bucket_uri=s3_bucket,
          s3_sync_args=["--no-sign-request"],
      )
  
      print(
          f"Downloading to {ckpt_path} finished after {time.time() - s} seconds."
      )
      print(f"Loading original model from {ckpt_path} ...")
  
      s2 = time.time()
  
      model = AutoModelForCausalLM.from_pretrained(
          ckpt_path,
          trust_remote_code=True,
          torch_dtype=torch.bfloat16,
          use_cache=False,
      )
      model.resize_token_embeddings(len(tokenizer))
  
      print(
          f"Done downloading and loading model after {time.time() - s2} seconds."
      )
      print("Loading and merging peft weights...")
      s3 = time.time()
  
      # Load LoRA weights
      model: peft.PeftModel = peft.PeftModel.from_pretrained(
          model=model,
          model_id=args.checkpoint,
      )
  
      # Merge weights and save
      model = model.merge_and_unload()
      output_path = Path(args.output_path)
      model.save_pretrained(output_path, safe_serialization=True)
      model.config.save_pretrained(output_path)
  
      print(
          f"Saved merged model to {args.output_path} after {time.time() - s3} seconds."
      )
      print(f"This script took {time.time() - s} seconds to execute.")
  
      if TEST_EVAL:
          test_eval(model, tokenizer)
  
file_path: ray/release/workspace_templates/04_finetuning_llms_with_deepspeed/merge_lora_weights.py
imports: import torch
  import argparse
  import time
  import peft
  from pathlib import Path
  from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteriaList
  from utils import download_model, get_mirror_link, get_checkpoint_and_refs_dir
label_desc: Use image node_id 3 for calling a global function with example usage: main() without return types