Spaces:
Running
Running
| """ | |
| PyTorch Model Wrapper | |
| -------------------------- | |
| """ | |
| import torch | |
| from torch.nn import CrossEntropyLoss | |
| import textattack | |
| from .model_wrapper import ModelWrapper | |
| torch.cuda.empty_cache() | |
| class PyTorchModelWrapper(ModelWrapper): | |
| """Loads a PyTorch model (`nn.Module`) and tokenizer. | |
| Args: | |
| 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. | |
| """ | |
| def __init__(self, model, tokenizer): | |
| if not isinstance(model, torch.nn.Module): | |
| raise TypeError( | |
| f"PyTorch model must be torch.nn.Module, got type {type(model)}" | |
| ) | |
| self.model = model | |
| self.tokenizer = tokenizer | |
| def to(self, device): | |
| self.model.to(device) | |
| def __call__(self, text_input_list, batch_size=32): | |
| model_device = next(self.model.parameters()).device | |
| ids = self.tokenizer(text_input_list) | |
| ids = torch.tensor(ids).to(model_device) | |
| with torch.no_grad(): | |
| outputs = textattack.shared.utils.batch_model_predict( | |
| self.model, ids, batch_size=batch_size | |
| ) | |
| return outputs | |
| def get_grad(self, text_input, loss_fn=CrossEntropyLoss()): | |
| """Get gradient of loss with respect to input tokens. | |
| Args: | |
| 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. | |
| """ | |
| if not hasattr(self.model, "get_input_embeddings"): | |
| raise AttributeError( | |
| f"{type(self.model)} must have method `get_input_embeddings` that returns `torch.nn.Embedding` object that represents input embedding layer" | |
| ) | |
| if not isinstance(loss_fn, torch.nn.Module): | |
| raise ValueError("Loss function must be of type `torch.nn.Module`.") | |
| self.model.train() | |
| embedding_layer = self.model.get_input_embeddings() | |
| original_state = embedding_layer.weight.requires_grad | |
| embedding_layer.weight.requires_grad = True | |
| emb_grads = [] | |
| def grad_hook(module, grad_in, grad_out): | |
| emb_grads.append(grad_out[0]) | |
| emb_hook = embedding_layer.register_backward_hook(grad_hook) | |
| self.model.zero_grad() | |
| model_device = next(self.model.parameters()).device | |
| ids = self.tokenizer([text_input]) | |
| ids = torch.tensor(ids).to(model_device) | |
| predictions = self.model(ids) | |
| output = predictions.argmax(dim=1) | |
| loss = loss_fn(predictions, output) | |
| loss.backward() | |
| # grad w.r.t to word embeddings | |
| # Fix for Issue #601 | |
| # Check if gradient has shape [max_sequence,1,_] ( when model input in transpose of input sequence) | |
| if emb_grads[0].shape[1] == 1: | |
| grad = torch.transpose(emb_grads[0], 0, 1)[0].cpu().numpy() | |
| else: | |
| # gradient has shape [1,max_sequence,_] | |
| grad = emb_grads[0][0].cpu().numpy() | |
| embedding_layer.weight.requires_grad = original_state | |
| emb_hook.remove() | |
| self.model.eval() | |
| output = {"ids": ids[0].tolist(), "gradient": grad} | |
| return output | |
| def _tokenize(self, inputs): | |
| """Helper method that for `tokenize` | |
| Args: | |
| inputs (list[str]): list of input strings | |
| Returns: | |
| tokens (list[list[str]]): List of list of tokens as strings | |
| """ | |
| return [self.tokenizer.convert_ids_to_tokens(self.tokenizer(x)) for x in inputs] | |