textflint.input_layer.model.flint_model.torch_model¶

PyTorch Model Wrapper¶

class textflint.input_layer.model.flint_model.torch_model.TorchModel(model, tokenizer, task='SA', batch_size=32)[source]¶

Loads a PyTorch model (nn.Module) and tokenizer.

__init__(model, tokenizer, task='SA', batch_size=32)[source]¶

Parameters

model (torch.nn.Module) – target 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.
task (str) – task name
batch_size (int) – batch size to apply evaluation

get_model_grad(text_inputs, loss_fn=CrossEntropyLoss())[source]¶

Get gradient of loss with respect to input tokens.

Parameters

text_inputs (str|[str]) – input string or input string list
loss_fn (torch.nn.Module) – loss function. Default is torch.nn.CrossEntropyLoss

Returns

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

encode(inputs)[source]¶

Tokenize inputs and convert it to ids.