text: import argparse
  import json
  import os
  
  from config_utils import DEFAULT_TASK_MAP, get_dataset_config, restore_path
  
  from itrain import ModelArguments, RunArguments, Setup
  
  
  FUSION_OUTPUT_DIR = "fusion_output"
  FUSION_FINETUNE_OUTPUT_DIR = "fusion_finetune_output"
  
  
  def run_fusion(args, config_name=None):
      # init setup
      dataset_manager, config = get_dataset_config(args["target_task"], train_size=args["train_size"])
      if args["train_size"] > 0:
          target_task_name = args["target_task"] + "_n" + str(args["train_size"])
      else:
          target_task_name = args["target_task"]
      if args["finetune_adapters"]:
          output_base = os.path.join(FUSION_FINETUNE_OUTPUT_DIR, "to_" + target_task_name)
      else:
          output_base = os.path.join(FUSION_OUTPUT_DIR, "to_" + target_task_name)
      # patch training args
      config["training"]["learning_rate"] = args["learning_rate"]
      config["training"]["num_train_epochs"] = args["num_train_epochs"]
  
      # load results if existing
      final_results_file = os.path.join(output_base, "eval_results.json")
      if os.path.exists(final_results_file):
          with open(final_results_file, "r") as f:
              results = json.load(f)
      else:
          results = {}
  
      with open(os.path.expanduser(args["task_map"]), "r") as f:
          task_map = json.load(f)
  
      fusion_name = ",".join(args["source_tasks"])
      print(f"*** Running fusion from {fusion_name} to {target_task_name} ***")
      output_dir = os.path.join(output_base, fusion_name)
      # skip this iteration if no overwrites requested & existing
      if args["overwrite_mode"] == 0 and os.path.exists(output_dir):
          print(f"Skipping task {fusion_name} as it already exists.")
          return
  
      # setup the dataset and training params
      setup = Setup(id=args["id"])
      setup.dataset(dataset_manager)
      config["training"]["output_dir"] = output_dir
      setup.training(RunArguments(**config["training"]))
      if isinstance(config["evaluation"], str):
          setup.evaluation(split=config["evaluation"])
      else:
          setup.evaluation()
      setup.notify(config["notify"])
      setup._config_name = config_name or "fusion_" + fusion_name + "_to_" + target_task_name
  
      # setup model
      load_adapters_map = {}
      # iterate over all adapters for fusion
      for source_task in args["source_tasks"]:
          source_dataset_manager, _ = get_dataset_config(source_task)
          load_adapters_map[source_task] = restore_path(task_map, source_task, source_dataset_manager)
      # patch model args
      config["model"]["load_adapters"] = load_adapters_map
      config["model"]["train_adapter_fusion"] = fusion_name
      config["model"]["drop_last_fusion_layer"] = True
      config["model"]["train_adapter"] = args["finetune_adapters"]
      setup.model(ModelArguments(**config["model"]))
      # start!
      if fusion_name in results and args["overwrite_mode"] == 1:
          # append to existing
          run_results = setup.run(restarts=args["restarts"], first_run_index=len(results[fusion_name]["seeds"]))
          for k, v in run_results.items():
              results[fusion_name][k] += v
      else:
          run_results = setup.run(restarts=args["restarts"])
          results[fusion_name] = run_results
  
      # save results
      with open(final_results_file, "w") as f:
          json.dump(results, f)
  
  
  if __name__ == "__main__":
      parser = argparse.ArgumentParser()
  
      parser.add_argument("target_task", type=str, help="Name of the target task training setup.")
      parser.add_argument("--id", type=int, default=0, help="ID of this run.")
      parser.add_argument("--task_map", type=str, default=DEFAULT_TASK_MAP)
      parser.add_argument(
          "--overwrite_mode", type=int, choices=[0, 1, 2], default=0, help="0: no overwrite; 1: append; 2: overwrite"
      )
      parser.add_argument("--source_tasks", type=lambda s: s.split(","), required=True)
      parser.add_argument("--learning_rate", type=float, default=5e-5)
      parser.add_argument("--num_train_epochs", type=int, default=15)
      parser.add_argument("--train_size", type=int, default=-1)
      parser.add_argument("--restarts", type=int, default=None)
      parser.add_argument("--finetune_adapters", action="store_true")
      args = vars(parser.parse_args())
  
      run_fusion(args)
  