Merge pull request #591 from QData/fix-StopIteration-bug

Fix CLARE StopIteration Bug
QData · Dec 3, 2021 · 5f825f6 · 5f825f6
2 parents a0ec175 + dbd31e8
commit 5f825f6
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 4 deletions.
diff --git a/textattack/transformations/word_insertions/word_insertion_masked_lm.py b/textattack/transformations/word_insertions/word_insertion_masked_lm.py
@@ -2,7 +2,7 @@
 WordInsertionMaskedLM Class
 -------------------------------
 """
-
+import re
 
 import torch
 from transformers import AutoModelForMaskedLM, AutoTokenizer
@@ -162,7 +162,7 @@ def _get_transformations(self, current_text, indices_to_modify):
             word_at_index = current_text.words[index_to_modify]
             for word in new_words[i]:
                 word = word.strip("Ġ")
-                if word != word_at_index:
+                if word != word_at_index and re.search("[a-zA-Z]", word):
                     transformed_texts.append(
                         current_text.insert_text_before_word_index(
                             index_to_modify, word

diff --git a/textattack/transformations/word_merges/word_merge_masked_lm.py b/textattack/transformations/word_merges/word_merge_masked_lm.py
@@ -3,6 +3,7 @@
 ------------------------------------------------
 
 """
+import re
 
 import torch
 from transformers import AutoModelForMaskedLM, AutoTokenizer
@@ -169,7 +170,7 @@ def _get_transformations(self, current_text, indices_to_modify):
             word_at_index = current_text.words[index_to_modify]
             for word in merged_words[i]:
                 word = word.strip("Ġ")
-                if word != word_at_index:
+                if word != word_at_index and re.search("[a-zA-Z]", word):
                     temp_text = current_text.delete_word_at_index(index_to_modify + 1)
                     transformed_texts.append(
                         temp_text.replace_word_at_index(index_to_modify, word)

diff --git a/textattack/transformations/word_swaps/word_swap_masked_lm.py b/textattack/transformations/word_swaps/word_swap_masked_lm.py
@@ -5,6 +5,7 @@
 
 
 import itertools
+import re
 
 import torch
 from transformers import AutoModelForMaskedLM, AutoTokenizer
@@ -290,7 +291,11 @@ def _get_transformations(self, current_text, indices_to_modify):
                 word_at_index = current_text.words[index_to_modify]
                 for word in replacement_words[i]:
                     word = word.strip("Ġ")
-                    if word != word_at_index and len(utils.words_from_text(word)) == 1:
+                    if (
+                        word != word_at_index
+                        and re.search("[a-zA-Z]", word)
+                        and len(utils.words_from_text(word)) == 1
+                    ):
                         transformed_texts.append(
                             current_text.replace_word_at_index(index_to_modify, word)
                         )