mirror of
https://github.com/QData/TextAttack.git
synced 2021-10-13 00:05:06 +03:00
114 lines
5.0 KiB
Python
114 lines
5.0 KiB
Python
import lru
|
|
import numpy as np
|
|
import torch
|
|
import math
|
|
|
|
from textattack.goal_functions import GoalFunctionResult
|
|
from textattack.shared import utils
|
|
|
|
class GoalFunction:
|
|
|
|
def __init__(self, model):
|
|
self.model = model
|
|
self.num_queries = 0
|
|
self._call_model_cache = lru.LRU(2**18)
|
|
|
|
def should_skip(self, tokenized_text, correct_output):
|
|
self.original_text = tokenized_text
|
|
self.correct_output = correct_output
|
|
model_outputs = self._call_model([tokenized_text])
|
|
return self._is_goal_complete(model_outputs[0])
|
|
|
|
def set_original_attrs(self, tokenized_text, correct_output):
|
|
self.original_text = tokenized_text
|
|
self.correct_output = correct_output
|
|
self.num_queries = 1
|
|
|
|
def get_results(self, tokenized_text_list):
|
|
model_outputs = self._call_model(tokenized_text_list)
|
|
results = []
|
|
for tokenized_text, raw_output in zip(tokenized_text_list, model_outputs):
|
|
succeeded = self._is_goal_complete(raw_output)
|
|
score = self._get_score(raw_output)
|
|
output = self._get_output(raw_output)
|
|
results.append(GoalFunctionResult(tokenized_text, output, succeeded, score))
|
|
return results
|
|
|
|
def _is_goal_complete(self, correct_output, model_output):
|
|
raise NotImplementedError()
|
|
|
|
def _get_score(self, model_output):
|
|
raise NotImplementedError()
|
|
|
|
def _get_output(self, raw_output):
|
|
return raw_output
|
|
|
|
def _call_model_uncached(self, tokenized_text_list, batch_size=8):
|
|
""" Queries model and returns predictions for a list of TokenizedText
|
|
objects.
|
|
"""
|
|
if not len(tokenized_text_list):
|
|
return torch.tensor([])
|
|
ids = [t.ids for t in tokenized_text_list]
|
|
ids = torch.tensor(ids).to(utils.get_device())
|
|
#
|
|
# shape of `ids` is (n, m, d)
|
|
# - n: number of elements in `tokenized_text_list`
|
|
# - m: number of vectors per element
|
|
# ex: most classification models take a single vector, so m=1
|
|
# ex: some entailment models take three vectors, so m=3
|
|
# - d: dimensionality of each vector
|
|
# (a typical model might set d=128 or d=256)
|
|
num_fields = ids.shape[1]
|
|
num_batches = int(math.ceil(len(tokenized_text_list) / float(batch_size)))
|
|
scores = []
|
|
for batch_i in range(num_batches):
|
|
batch_start = batch_i * batch_size
|
|
batch_stop = (batch_i + 1) * batch_size
|
|
batch_ids = ids[batch_start:batch_stop]
|
|
batch = [batch_ids[:, x, :] for x in range(num_fields)]
|
|
with torch.no_grad():
|
|
preds = self.model(*batch)
|
|
scores.append(preds)
|
|
scores = torch.cat(scores, dim=0)
|
|
# Validation check on model score dimensions
|
|
if scores.ndim == 1:
|
|
# Unsqueeze prediction, if it's been squeezed by the model.
|
|
if len(tokenized_text_list == 1):
|
|
scores = scores.unsqueeze(dim=0)
|
|
else:
|
|
raise ValueError(f'Model return score of shape {scores.shape} for {len(tokenized_text_list)} inputs.')
|
|
elif scores.ndim != 2:
|
|
# If model somehow returns too may dimensions, throw an error.
|
|
raise ValueError(f'Model return score of shape {scores.shape} for {len(tokenized_text_list)} inputs.')
|
|
elif scores.shape[0] != len(tokenized_text_list):
|
|
# If model returns an incorrect number of scores, throw an error.
|
|
raise ValueError(f'Model return score of shape {scores.shape} for {len(tokenized_text_list)} inputs.')
|
|
elif not ((scores.sum(dim=1) - 1).abs() < 1e-6).all():
|
|
# Values in each row should sum up to 1. The model should return a
|
|
# set of numbers corresponding to probabilities, which should add
|
|
# up to 1. Since they are `torch.float` values, allow a small
|
|
# error in the summation.
|
|
raise ValueError('Model scores do not add up to 1.')
|
|
return scores
|
|
|
|
def _call_model(self, tokenized_text_list):
|
|
""" Gets predictions for a list of `TokenizedText` objects.
|
|
|
|
Gets prediction from cache if possible. If prediction is not in the
|
|
cache, queries model and stores prediction in cache.
|
|
"""
|
|
try:
|
|
self.num_queries += len(tokenized_text_list)
|
|
except AttributeError:
|
|
# If some outside class is just using the attack for its `call_model`
|
|
# function, then `self.num_queries` will not have been initialized.
|
|
# In this case, just continue.
|
|
pass
|
|
uncached_list = [text for text in tokenized_text_list if text not in self._call_model_cache]
|
|
scores = self._call_model_uncached(uncached_list)
|
|
for text, score in zip(uncached_list, scores):
|
|
self._call_model_cache[text] = score.cpu()
|
|
final_scores = [self._call_model_cache[text].to(utils.get_device()) for text in tokenized_text_list]
|
|
return torch.stack(final_scores)
|