Spaces:
Sleeping
Sleeping
| """ | |
| CoLA for Grammaticality | |
| -------------------------- | |
| """ | |
| import lru | |
| import nltk | |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer | |
| from textattack.constraints import Constraint | |
| from textattack.models.wrappers import HuggingFaceModelWrapper | |
| class COLA(Constraint): | |
| """Constrains an attack to text that has a similar number of linguistically | |
| accecptable sentences as the original text. Linguistic acceptability is | |
| determined by a model pre-trained on the `CoLA dataset <https://nyu- | |
| mll.github.io/CoLA/>`_. By default a BERT model is used, see the `pre- | |
| trained models README <https://github.com/QData/TextAttack/tree/master/ | |
| textattack/models>`_ for a full list of available models or provide your | |
| own model from the huggingface model hub. | |
| Args: | |
| max_diff (float or int): The absolute (if int or greater than or equal to 1) or percent (if float and less than 1) | |
| maximum difference allowed between the number of valid sentences in the reference | |
| text and the number of valid sentences in the attacked text. | |
| model_name (str): The name of the pre-trained model to use for classification. The model must be in huggingface model hub. | |
| compare_against_original (bool): If `True`, compare against the original text. | |
| Otherwise, compare against the most recent text. | |
| """ | |
| def __init__( | |
| self, | |
| max_diff, | |
| model_name="textattack/bert-base-uncased-CoLA", | |
| compare_against_original=True, | |
| ): | |
| super().__init__(compare_against_original) | |
| if not isinstance(max_diff, float) and not isinstance(max_diff, int): | |
| raise TypeError("max_diff must be a float or int") | |
| if max_diff < 0.0: | |
| raise ValueError("max_diff must be a value greater or equal to than 0.0") | |
| self.max_diff = max_diff | |
| self.model_name = model_name | |
| self._reference_score_cache = lru.LRU(2**10) | |
| model = AutoModelForSequenceClassification.from_pretrained(model_name) | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| self.model = HuggingFaceModelWrapper(model, tokenizer) | |
| def clear_cache(self): | |
| self._reference_score_cache.clear() | |
| def _check_constraint(self, transformed_text, reference_text): | |
| if reference_text not in self._reference_score_cache: | |
| # Split the text into sentences before predicting validity | |
| reference_sentences = nltk.sent_tokenize(reference_text.text) | |
| # A label of 1 indicates the sentence is valid | |
| num_valid = self.model(reference_sentences).argmax(axis=1).sum() | |
| self._reference_score_cache[reference_text] = num_valid | |
| sentences = nltk.sent_tokenize(transformed_text.text) | |
| predictions = self.model(sentences) | |
| num_valid = predictions.argmax(axis=1).sum() | |
| reference_score = self._reference_score_cache[reference_text] | |
| if isinstance(self.max_diff, int) or self.max_diff >= 1: | |
| threshold = reference_score - self.max_diff | |
| else: | |
| threshold = reference_score - (reference_score * self.max_diff) | |
| if num_valid < threshold: | |
| return False | |
| return True | |
| def extra_repr_keys(self): | |
| return [ | |
| "max_diff", | |
| "model_name", | |
| ] + super().extra_repr_keys() | |
| def __getstate__(self): | |
| state = self.__dict__.copy() | |
| state["_reference_score_cache"] = self._reference_score_cache.get_size() | |
| return state | |
| def __setstate__(self, state): | |
| self.__dict__ = state | |
| self._reference_score_cache = lru.LRU(state["_reference_score_cache"]) | |