textattack.models.wrappers package
Model Wrappers Package
TextAttack can attack any model that takes a list of strings as input and outputs a list of predictions. This is the idea behind model wrappers: to help your model conform to this API, we’ve provided the textattack.models.wrappers.ModelWrapper
abstract class.
We’ve also provided implementations of model wrappers for common patterns in some popular machine learning frameworks:
HuggingFace Model Wrapper
- class textattack.models.wrappers.huggingface_model_wrapper.HuggingFaceModelWrapper(model, tokenizer)[source]
Bases:
PyTorchModelWrapper
Loads a HuggingFace
transformers
model and tokenizer.
ModelWrapper class
- class textattack.models.wrappers.model_wrapper.ModelWrapper[source]
Bases:
ABC
A model wrapper queries a model with a list of text inputs.
Classification-based models return a list of lists, where each sublist represents the model’s scores for a given input.
Text-to-text models return a list of strings, where each string is the output – like a translation or summarization – for a given input.
- tokenize(inputs, strip_prefix=False)[source]
Helper method that tokenizes input strings :param inputs: list of input strings :type inputs: list[str] :param strip_prefix: If True, we strip auxiliary characters added to tokens as prefixes (e.g. “##” for BERT, “Ġ” for RoBERTa) :type strip_prefix: bool
- Returns:
List of list of tokens as strings
- Return type:
tokens (list[list[str]])
PyTorch Model Wrapper
- class textattack.models.wrappers.pytorch_model_wrapper.PyTorchModelWrapper(model, tokenizer)[source]
Bases:
ModelWrapper
Loads a PyTorch model (nn.Module) and tokenizer.
- Parameters:
model (torch.nn.Module) – PyTorch model
tokenizer – tokenizer whose output can be packed as a tensor and passed to the model. No type requirement, but most have tokenizer method that accepts list of strings.
scikit-learn Model Wrapper
- class textattack.models.wrappers.sklearn_model_wrapper.SklearnModelWrapper(model, tokenizer)[source]
Bases:
ModelWrapper
Loads a scikit-learn model and tokenizer (tokenizer implements transform and model implements predict_proba).
May need to be extended and modified for different types of tokenizers.
TensorFlow Model Wrapper
- class textattack.models.wrappers.tensorflow_model_wrapper.TensorFlowModelWrapper(model)[source]
Bases:
ModelWrapper
Loads a TensorFlow model and tokenizer.
TensorFlow models can use many different architectures and tokenization strategies. This assumes that the model takes an np.array of strings as input and returns a tf.Tensor of outputs, as is typical with Keras modules. You may need to subclass this for models that have dedicated tokenizers or otherwise take input differently.