Source code for textattack.model_args

"""
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

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",
}


[docs]@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.""" 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`." return model