code: def main():
      parser = argparse.ArgumentParser()
  
      # Required parameters
      parser.add_argument(
          "--data_dir",
          default=None,
          type=str,
          required=True,
          help="The input data dir. Should contain the .tsv files (or other data files) for the task.",
      )
      parser.add_argument(
          "--model_type",
          default=None,
          type=str,
          required=True,
          help="Model type selected in the list: "
          + ", ".join(MODEL_CLASSES.keys()),
      )
      parser.add_argument(
          "--model_name_or_path",
          default=None,
          type=str,
          required=True,
          help="Path to pre-trained model or shortcut name.",
      )
      parser.add_argument(
          "--task_name",
          default=None,
          type=str,
          required=True,
          help="The name of the task to train selected in the list: "
          + ", ".join(processors.keys()),
      )
      parser.add_argument(
          "--output_dir",
          default=None,
          type=str,
          required=True,
          help="The output directory where the model predictions and checkpoints will be written.",
      )
      parser.add_argument(
          "--patience",
          default="0",
          type=str,
          required=False,
      )
      parser.add_argument(
          "--regression_threshold",
          default=0,
          type=float,
          required=False,
      )
  
      # Other parameters
      parser.add_argument(
          "--config_name",
          default="",
          type=str,
          help="Pretrained config name or path if not the same as model_name",
      )
      parser.add_argument(
          "--tokenizer_name",
          default="",
          type=str,
          help="Pretrained tokenizer name or path if not the same as model_name",
      )
      parser.add_argument(
          "--cache_dir",
          default="",
          type=str,
          help="Where do you want to store the pre-trained models downloaded from huggingface.co",
      )
      parser.add_argument(
          "--max_seq_length",
          default=128,
          type=int,
          help=(
              "The maximum total input sequence length after tokenization. Sequences longer "
              "than this will be truncated, sequences shorter will be padded."
          ),
      )
      parser.add_argument(
          "--do_train",
          action="store_true",
          help="Whether to run training.",
      )
      parser.add_argument(
          "--do_eval",
          action="store_true",
          help="Whether to run eval on the dev set.",
      )
      parser.add_argument(
          "--evaluate_during_training",
          action="store_true",
          help="Run evaluation during training at each logging step.",
      )
      parser.add_argument(
          "--do_lower_case",
          action="store_true",
          help="Set this flag if you are using an uncased model.",
      )
  
      parser.add_argument(
          "--per_gpu_train_batch_size",
          default=8,
          type=int,
          help="Batch size per GPU/CPU for training.",
      )
      parser.add_argument(
          "--per_gpu_eval_batch_size",
          default=1,
          type=int,
          help="Batch size per GPU/CPU for evaluation.",
      )
      parser.add_argument(
          "--gradient_accumulation_steps",
          type=int,
          default=1,
          help="Number of updates steps to accumulate before performing a backward/update pass.",
      )
      parser.add_argument(
          "--learning_rate",
          default=5e-5,
          type=float,
          help="The initial learning rate for Adam.",
      )
      parser.add_argument(
          "--weight_decay",
          default=0.0,
          type=float,
          help="Weight decay if we apply some.",
      )
      parser.add_argument(
          "--adam_epsilon",
          default=1e-8,
          type=float,
          help="Epsilon for Adam optimizer.",
      )
      parser.add_argument(
          "--max_grad_norm",
          default=1.0,
          type=float,
          help="Max gradient norm.",
      )
      parser.add_argument(
          "--num_train_epochs",
          default=3.0,
          type=float,
          help="Total number of training epochs to perform.",
      )
      parser.add_argument(
          "--max_steps",
          default=-1,
          type=int,
          help="If > 0: set total number of training steps to perform. Override num_train_epochs.",
      )
      parser.add_argument(
          "--warmup_steps",
          default=0,
          type=int,
          help="Linear warmup over warmup_steps.",
      )
  
      parser.add_argument(
          "--logging_steps",
          type=int,
          default=500,
          help="Log every X updates steps.",
      )
      parser.add_argument(
          "--save_steps",
          type=int,
          default=500,
          help="Save checkpoint every X updates steps.",
      )
      parser.add_argument(
          "--eval_all_checkpoints",
          action="store_true",
          help="Evaluate all checkpoints starting with the same prefix as model_name ending and ending with step number",
      )
      parser.add_argument(
          "--no_cuda",
          action="store_true",
          help="Avoid using CUDA when available",
      )
      parser.add_argument(
          "--overwrite_output_dir",
          action="store_true",
          help="Overwrite the content of the output directory",
      )
      parser.add_argument(
          "--overwrite_cache",
          action="store_true",
          help="Overwrite the cached training and evaluation sets",
      )
      parser.add_argument(
          "--seed",
          type=int,
          default=42,
          help="random seed for initialization",
      )
  
      parser.add_argument(
          "--fp16",
          action="store_true",
          help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit",
      )
      parser.add_argument(
          "--fp16_opt_level",
          type=str,
          default="O1",
          help=(
              "For fp16: Apex AMP optimization level selected in ['O0', 'O1', 'O2', and 'O3']. "
              "See details at https://nvidia.github.io/apex/amp.html"
          ),
      )
      parser.add_argument(
          "--local_rank",
          type=int,
          default=-1,
          help="For distributed training: local_rank",
      )
      parser.add_argument(
          "--server_ip",
          type=str,
          default="",
          help="For distant debugging.",
      )
      parser.add_argument(
          "--server_port",
          type=str,
          default="",
          help="For distant debugging.",
      )
      args = parser.parse_args()
  
      if (
          os.path.exists(args.output_dir)
          and os.listdir(args.output_dir)
          and args.do_train
          and not args.overwrite_output_dir
      ):
          raise ValueError(
              "Output directory ({}) already exists and is not empty. Use --overwrite_output_dir to overcome.".format(
                  args.output_dir
              )
          )
  
      # Setup distant debugging if needed
      if args.server_ip and args.server_port:
          # Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
          import ptvsd
  
          print("Waiting for debugger attach")
          ptvsd.enable_attach(
              address=(args.server_ip, args.server_port),
              redirect_output=True,
          )
          ptvsd.wait_for_attach()
  
      # Setup CUDA, GPU & distributed training
      if args.local_rank == -1 or args.no_cuda:
          device = torch.device(
              "cuda"
              if torch.cuda.is_available() and not args.no_cuda
              else "cpu"
          )
          args.n_gpu = torch.cuda.device_count()
      else:  # Initializes the distributed backend which will take care of sychronizing nodes/GPUs
          torch.cuda.set_device(args.local_rank)
          device = torch.device("cuda", args.local_rank)
          torch.distributed.init_process_group(backend="nccl")
          args.n_gpu = 1
      args.device = device
  
      # Setup logging
      logging.basicConfig(
          format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
          datefmt="%m/%d/%Y %H:%M:%S",
          level=logging.INFO
          if args.local_rank in [-1, 0]
          else logging.WARN,
      )
      logger.warning(
          "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
          args.local_rank,
          device,
          args.n_gpu,
          bool(args.local_rank != -1),
          args.fp16,
      )
      # Set the verbosity to info of the Transformers logger (on main process only):
      if is_main_process(args.local_rank):
          transformers.utils.logging.set_verbosity_info()
          transformers.utils.logging.enable_default_handler()
          transformers.utils.logging.enable_explicit_format()
      # Set seed
      set_seed(args)
  
      # Prepare GLUE task
      args.task_name = args.task_name.lower()
      if args.task_name not in processors:
          raise ValueError("Task not found: %s" % (args.task_name))
      processor = processors[args.task_name]()
      args.output_mode = output_modes[args.task_name]
      label_list = processor.get_labels()
      num_labels = len(label_list)
  
      if args.patience != "0" and args.per_gpu_eval_batch_size != 1:
          raise ValueError(
              "The eval batch size must be 1 with PABEE inference on."
          )
  
      # Load pretrained model and tokenizer
      if args.local_rank not in [-1, 0]:
          torch.distributed.barrier()  # Make sure only the first process in distributed training will download model & vocab
  
      args.model_type = args.model_type.lower()
      config_class, model_class, tokenizer_class = MODEL_CLASSES[
          args.model_type
      ]
      config = config_class.from_pretrained(
          args.config_name
          if args.config_name
          else args.model_name_or_path,
          num_labels=num_labels,
          finetuning_task=args.task_name,
          cache_dir=args.cache_dir if args.cache_dir else None,
      )
      tokenizer = tokenizer_class.from_pretrained(
          args.tokenizer_name
          if args.tokenizer_name
          else args.model_name_or_path,
          do_lower_case=args.do_lower_case,
          cache_dir=args.cache_dir if args.cache_dir else None,
      )
      model = model_class.from_pretrained(
          args.model_name_or_path,
          from_tf=bool(".ckpt" in args.model_name_or_path),
          config=config,
          cache_dir=args.cache_dir if args.cache_dir else None,
      )
  
      if args.local_rank == 0:
          torch.distributed.barrier()  # Make sure only the first process in distributed training will download model & vocab
  
      model.to(args.device)
  
      print(
          "Total Model Parameters:",
          sum(param.numel() for param in model.parameters()),
      )
      output_layers_param_num = sum(
          param.numel() for param in model.classifiers.parameters()
      )
      print("Output Layers Parameters:", output_layers_param_num)
      single_output_layer_param_num = sum(
          param.numel() for param in model.classifiers[0].parameters()
      )
      print(
          "Added Output Layers Parameters:",
          output_layers_param_num - single_output_layer_param_num,
      )
  
      logger.info("Training/evaluation parameters %s", args)
  
      # Training
      if args.do_train:
          train_dataset = load_and_cache_examples(
              args, args.task_name, tokenizer, evaluate=False
          )
          global_step, tr_loss = train(
              args, train_dataset, model, tokenizer
          )
          logger.info(
              " global_step = %s, average loss = %s",
              global_step,
              tr_loss,
          )
  
      # Saving best-practices: if you use defaults names for the model, you can reload it using from_pretrained()
      if args.do_train and (
          args.local_rank == -1 or torch.distributed.get_rank() == 0
      ):
          logger.info("Saving model checkpoint to %s", args.output_dir)
          # Save a trained model, configuration and tokenizer using `save_pretrained()`.
          # They can then be reloaded using `from_pretrained()`
          model_to_save = (
              model.module if hasattr(model, "module") else model
          )  # Take care of distributed/parallel training
          model_to_save.save_pretrained(args.output_dir)
          tokenizer.save_pretrained(args.output_dir)
  
          # Good practice: save your training arguments together with the trained model
          torch.save(
              args, os.path.join(args.output_dir, "training_args.bin")
          )
  
          # Load a trained model and vocabulary that you have fine-tuned
          model = model_class.from_pretrained(args.output_dir)
          tokenizer = tokenizer_class.from_pretrained(args.output_dir)
          model.to(args.device)
  
      # Evaluation
      results = {}
      if args.do_eval and args.local_rank in [-1, 0]:
          patience_list = [int(x) for x in args.patience.split(",")]
          tokenizer = tokenizer_class.from_pretrained(
              args.output_dir, do_lower_case=args.do_lower_case
          )
          checkpoints = [args.output_dir]
          if args.eval_all_checkpoints:
              checkpoints = [
                  os.path.dirname(c)
                  for c in sorted(
                      glob.glob(
                          args.output_dir + "/**/" + WEIGHTS_NAME,
                          recursive=True,
                      )
                  )
              ]
  
          logger.info(
              "Evaluate the following checkpoints: %s", checkpoints
          )
  
          for checkpoint in checkpoints:
              global_step = (
                  checkpoint.split("-")[-1]
                  if len(checkpoints) > 1
                  else ""
              )
              prefix = (
                  checkpoint.split("/")[-1]
                  if checkpoint.find("checkpoint") != -1
                  else ""
              )
  
              model = model_class.from_pretrained(checkpoint)
              model.to(args.device)
  
              print(f"Evaluation for checkpoint {prefix}")
              for patience in patience_list:
                  result = evaluate(
                      args,
                      model,
                      tokenizer,
                      prefix=prefix,
                      patience=patience,
                  )
                  result = {
                      k + "_{}".format(global_step): v
                      for k, v in result.items()
                  }
                  results.update(result)
      return results
  
imports: import argparse
  import glob
  import json
  import logging
  import os
  import random
  import numpy
  import torch
  from pabee.modeling_pabee_albert import AlbertForSequenceClassificationWithPabee
  from pabee.modeling_pabee_bert import BertForSequenceClassificationWithPabee
  from torch import nn
  from torch.utils.data import DataLoader, RandomSampler, SequentialSampler, TensorDataset
  from torch.utils.data.distributed import DistributedSampler
  from tqdm import tqdm, trange
  import transformers
  from transformers import WEIGHTS_NAME, AdamW, AlbertConfig, AlbertTokenizer, BertConfig, BertTokenizer, get_linear_schedule_with_warmup
  from transformers import glue_compute_metrics
  from transformers import glue_convert_examples_to_features
  from transformers import glue_output_modes
  from transformers import glue_processors
  from transformers.trainer_utils import is_main_process
label_desc: Use image node_id 5 for calling a global function with example usage: main() and returns: results