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.

get_grad(text_input)[source]

Get gradient of loss with respect to input tokens.

Parameters:

text_input (str) – input string

Returns:

Dict of ids, tokens, and gradient as numpy array.

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.

get_grad(text_input)[source]

Get gradient of loss with respect to input tokens.

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.

get_grad(text_input, loss_fn=CrossEntropyLoss())[source]

Get gradient of loss with respect to input tokens.

Parameters:
  • text_input (str) – input string

  • loss_fn (torch.nn.Module) – loss function. Default is torch.nn.CrossEntropyLoss

Returns:

Dict of ids, tokens, and gradient as numpy array.

to(device)[source]

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.

get_grad(text_input)[source]

Get gradient of loss with respect to input tokens.

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.

get_grad(text_input)[source]

Get gradient of loss with respect to input tokens.