text: """
  ModelArgs Class
  ===============
  """
  
  
  from dataclasses import dataclass
  import json
  import os
  import transformers
  import textattack
  from textattack.shared.utils import ARGS_SPLIT_TOKEN, load_module_from_file
  from common_pyutil.monitor import Timer
  
  HUGGINGFACE_MODELS = {
      #
      # bert-base-uncased
      #
      "bert-base-uncased": "bert-base-uncased",
      "bert-base-uncased-ag-news": "textattack/bert-base-uncased-ag-news",
      "bert-base-uncased-cola": "textattack/bert-base-uncased-CoLA",
      "bert-base-uncased-imdb": "textattack/bert-base-uncased-imdb",
      "bert-base-uncased-mnli": "textattack/bert-base-uncased-MNLI",
      "bert-base-uncased-mrpc": "textattack/bert-base-uncased-MRPC",
      "bert-base-uncased-qnli": "textattack/bert-base-uncased-QNLI",
      "bert-base-uncased-qqp": "textattack/bert-base-uncased-QQP",
      "bert-base-uncased-rte": "textattack/bert-base-uncased-RTE",
      "bert-base-uncased-sst2": "textattack/bert-base-uncased-SST-2",
      "bert-base-uncased-stsb": "textattack/bert-base-uncased-STS-B",
      "bert-base-uncased-wnli": "textattack/bert-base-uncased-WNLI",
      "bert-base-uncased-mr": "textattack/bert-base-uncased-rotten-tomatoes",
      "bert-base-uncased-snli": "textattack/bert-base-uncased-snli",
      "bert-base-uncased-yelp": "textattack/bert-base-uncased-yelp-polarity",
      #
      # distilbert-base-cased
      #
      "distilbert-base-uncased": "distilbert-base-uncased",
      "distilbert-base-cased-cola": "textattack/distilbert-base-cased-CoLA",
      "distilbert-base-cased-mrpc": "textattack/distilbert-base-cased-MRPC",
      "distilbert-base-cased-qqp": "textattack/distilbert-base-cased-QQP",
      "distilbert-base-cased-snli": "textattack/distilbert-base-cased-snli",
      "distilbert-base-cased-sst2": "textattack/distilbert-base-cased-SST-2",
      "distilbert-base-cased-stsb": "textattack/distilbert-base-cased-STS-B",
      "distilbert-base-uncased-ag-news": "textattack/distilbert-base-uncased-ag-news",
      "distilbert-base-uncased-cola": "textattack/distilbert-base-cased-CoLA",
      "distilbert-base-uncased-imdb": "textattack/distilbert-base-uncased-imdb",
      "distilbert-base-uncased-mnli": "textattack/distilbert-base-uncased-MNLI",
      "distilbert-base-uncased-mr": "textattack/distilbert-base-uncased-rotten-tomatoes",
      "distilbert-base-uncased-mrpc": "textattack/distilbert-base-uncased-MRPC",
      "distilbert-base-uncased-qnli": "textattack/distilbert-base-uncased-QNLI",
      "distilbert-base-uncased-rte": "textattack/distilbert-base-uncased-RTE",
      "distilbert-base-uncased-wnli": "textattack/distilbert-base-uncased-WNLI",
      #
      # roberta-base (RoBERTa is cased by default)
      #
      "roberta-base": "roberta-base",
      "roberta-base-ag-news": "textattack/roberta-base-ag-news",
      "roberta-base-cola": "textattack/roberta-base-CoLA",
      "roberta-base-imdb": "textattack/roberta-base-imdb",
      "roberta-base-mr": "textattack/roberta-base-rotten-tomatoes",
      "roberta-base-mrpc": "textattack/roberta-base-MRPC",
      "roberta-base-qnli": "textattack/roberta-base-QNLI",
      "roberta-base-rte": "textattack/roberta-base-RTE",
      "roberta-base-sst2": "textattack/roberta-base-SST-2",
      "roberta-base-stsb": "textattack/roberta-base-STS-B",
      "roberta-base-wnli": "textattack/roberta-base-WNLI",
      #
      # albert-base-v2 (ALBERT is cased by default)
      #
      "albert-base-v2": "albert-base-v2",
      "albert-base-v2-ag-news": "textattack/albert-base-v2-ag-news",
      "albert-base-v2-cola": "textattack/albert-base-v2-CoLA",
      "albert-base-v2-imdb": "textattack/albert-base-v2-imdb",
      "albert-base-v2-mr": "textattack/albert-base-v2-rotten-tomatoes",
      "albert-base-v2-rte": "textattack/albert-base-v2-RTE",
      "albert-base-v2-qqp": "textattack/albert-base-v2-QQP",
      "albert-base-v2-snli": "textattack/albert-base-v2-snli",
      "albert-base-v2-sst2": "textattack/albert-base-v2-SST-2",
      "albert-base-v2-stsb": "textattack/albert-base-v2-STS-B",
      "albert-base-v2-wnli": "textattack/albert-base-v2-WNLI",
      "albert-base-v2-yelp": "textattack/albert-base-v2-yelp-polarity",
      #
      # xlnet-base-cased
      #
      "xlnet-base-cased": "xlnet-base-cased",
      "xlnet-base-cased-cola": "textattack/xlnet-base-cased-CoLA",
      "xlnet-base-cased-imdb": "textattack/xlnet-base-cased-imdb",
      "xlnet-base-cased-mr": "textattack/xlnet-base-cased-rotten-tomatoes",
      "xlnet-base-cased-mrpc": "textattack/xlnet-base-cased-MRPC",
      "xlnet-base-cased-rte": "textattack/xlnet-base-cased-RTE",
      "xlnet-base-cased-stsb": "textattack/xlnet-base-cased-STS-B",
      "xlnet-base-cased-wnli": "textattack/xlnet-base-cased-WNLI",
  }
  
  
  #
  # Models hosted by textattack.
  # `models` vs `models_v2`: `models_v2` is simply a new dir in S3 that contains models' `config.json`.
  # Fixes issue https://github.com/QData/TextAttack/issues/485
  # Model parameters has not changed.
  #
  TEXTATTACK_MODELS = {
      #
      # LSTMs
      #
      "lstm-ag-news": "models_v2/classification/lstm/ag-news",
      "lstm-imdb": "models_v2/classification/lstm/imdb",
      "lstm-mr": "models_v2/classification/lstm/mr",
      "lstm-sst2": "models_v2/classification/lstm/sst2",
      "lstm-yelp": "models_v2/classification/lstm/yelp",
      #
      # CNNs
      #
      "cnn-ag-news": "models_v2/classification/cnn/ag-news",
      "cnn-imdb": "models_v2/classification/cnn/imdb",
      "cnn-mr": "models_v2/classification/cnn/rotten-tomatoes",
      "cnn-sst2": "models_v2/classification/cnn/sst",
      "cnn-yelp": "models_v2/classification/cnn/yelp",
      #
      # T5 for translation
      #
      "t5-en-de": "english_to_german",
      "t5-en-fr": "english_to_french",
      "t5-en-ro": "english_to_romanian",
      #
      # T5 for summarization
      #
      "t5-summarization": "summarization",
  }
  
  
  @dataclass
  class ModelArgs:
      """Arguments for loading base/pretrained or trained models."""
  
      model: str = None
      model_from_file: str = None
      model_from_huggingface: str = None
  
      @classmethod
      def _add_parser_args(cls, parser):
          """Adds model-related arguments to an argparser."""
          model_group = parser.add_mutually_exclusive_group()
  
          model_names = list(HUGGINGFACE_MODELS.keys()) + list(TEXTATTACK_MODELS.keys())
          model_group.add_argument(
              "--model",
              type=str,
              required=False,
              default=None,
              help="Name of or path to a pre-trained TextAttack model to load. Choices: "
              + str(model_names),
          )
          model_group.add_argument(
              "--model-from-file",
              type=str,
              required=False,
              help="File of model and tokenizer to import.",
          )
          model_group.add_argument(
              "--model-from-huggingface",
              type=str,
              required=False,
              help="Name of or path of pre-trained HuggingFace model to load.",
          )
  
          return parser
  
      @classmethod
      def _create_model_from_args(cls, args):
          """Given ``ModelArgs``, return specified
          ``textattack.models.wrappers.ModelWrapper`` object."""
  
          timer = Timer()
          with timer:
              assert isinstance(
                  args, cls
              ), f"Expect args to be of type `{type(cls)}`, but got type `{type(args)}`."
  
              if args.model_from_file:
                  # Support loading the model from a .py file where a model wrapper
                  # is instantiated.
                  colored_model_name = textattack.shared.utils.color_text(
                      args.model_from_file, color="blue", method="ansi"
                  )
                  textattack.shared.logger.info(
                      f"Loading model and tokenizer from file: {colored_model_name}"
                  )
                  if ARGS_SPLIT_TOKEN in args.model_from_file:
                      model_file, model_name = args.model_from_file.split(ARGS_SPLIT_TOKEN)
                  else:
                      _, model_name = args.model_from_file, "model"
                  try:
                      model_module = load_module_from_file(args.model_from_file)
                  except Exception:
                      raise ValueError(f"Failed to import file {args.model_from_file}.")
                  try:
                      model = getattr(model_module, model_name)
                  except AttributeError:
                      raise AttributeError(
                          f"Variable `{model_name}` not found in module {args.model_from_file}."
                      )
  
                  if not isinstance(model, textattack.models.wrappers.ModelWrapper):
                      raise TypeError(
                          f"Variable `{model_name}` must be of type "
                          f"``textattack.models.ModelWrapper``, got type {type(model)}."
                      )
              elif (args.model in HUGGINGFACE_MODELS) or args.model_from_huggingface:
                  # Support loading models automatically from the HuggingFace model hub.
  
                  model_name = (
                      HUGGINGFACE_MODELS[args.model]
                      if (args.model in HUGGINGFACE_MODELS)
                      else args.model_from_huggingface
                  )
                  colored_model_name = textattack.shared.utils.color_text(
                      model_name, color="blue", method="ansi"
                  )
                  textattack.shared.logger.info(
                      f"Loading pre-trained model from HuggingFace model repository: {colored_model_name}"
                  )
                  model = transformers.AutoModelForSequenceClassification.from_pretrained(
                      model_name
                  )
                  tokenizer = transformers.AutoTokenizer.from_pretrained(
                      model_name, use_fast=True
                  )
                  model = textattack.models.wrappers.HuggingFaceModelWrapper(model, tokenizer)
              elif args.model in TEXTATTACK_MODELS:
                  # Support loading TextAttack pre-trained models via just a keyword.
                  colored_model_name = textattack.shared.utils.color_text(
                      args.model, color="blue", method="ansi"
                  )
                  if args.model.startswith("lstm"):
                      textattack.shared.logger.info(
                          f"Loading pre-trained TextAttack LSTM: {colored_model_name}"
                      )
                      model = textattack.models.helpers.LSTMForClassification.from_pretrained(
                          args.model
                      )
                  elif args.model.startswith("cnn"):
                      textattack.shared.logger.info(
                          f"Loading pre-trained TextAttack CNN: {colored_model_name}"
                      )
                      model = (
                          textattack.models.helpers.WordCNNForClassification.from_pretrained(
                              args.model
                          )
                      )
                  elif args.model.startswith("t5"):
                      model = textattack.models.helpers.T5ForTextToText.from_pretrained(
                          args.model
                      )
                  else:
                      raise ValueError(f"Unknown textattack model {args.model}")
  
                  # Choose the approprate model wrapper (based on whether or not this is
                  # a HuggingFace model).
                  if isinstance(model, textattack.models.helpers.T5ForTextToText):
                      model = textattack.models.wrappers.HuggingFaceModelWrapper(
                          model, model.tokenizer
                      )
                  else:
                      model = textattack.models.wrappers.PyTorchModelWrapper(
                          model, model.tokenizer
                      )
              elif args.model and os.path.exists(args.model):
                  # Support loading TextAttack-trained models via just their folder path.
                  # If `args.model` is a path/directory, let's assume it was a model
                  # trained with textattack, and try and load it.
                  if os.path.exists(os.path.join(args.model, "t5-wrapper-config.json")):
                      model = textattack.models.helpers.T5ForTextToText.from_pretrained(
                          args.model
                      )
                      model = textattack.models.wrappers.HuggingFaceModelWrapper(
                          model, model.tokenizer
                      )
                  elif os.path.exists(os.path.join(args.model, "config.json")):
                      with open(os.path.join(args.model, "config.json")) as f:
                          config = json.load(f)
                      model_class = config["architectures"]
                      if (
                          model_class == "LSTMForClassification"
                          or model_class == "WordCNNForClassification"
                      ):
                          model = eval(
                              f"textattack.models.helpers.{model_class}.from_pretrained({args.model})"
                          )
                          model = textattack.models.wrappers.PyTorchModelWrapper(
                              model, model.tokenizer
                          )
                      else:
                          # assume the model is from HuggingFace.
                          model = (
                              transformers.AutoModelForSequenceClassification.from_pretrained(
                                  args.model
                              )
                          )
                          tokenizer = transformers.AutoTokenizer.from_pretrained(
                              args.model, use_fast=True
                          )
                          model = textattack.models.wrappers.HuggingFaceModelWrapper(
                              model, tokenizer
                          )
              else:
                  raise ValueError(f"Error: unsupported TextAttack model {args.model}")
  
              assert isinstance(
                  model, textattack.models.wrappers.ModelWrapper
              ), "`model` must be of type `textattack.models.wrappers.ModelWrapper`."
  
          print(f"Model create time: {timer.time}")
          return model
  