Merge pull request #654 from VijayKalmath/Fix-Pwws-QueryBudgetIndexError

Fix Querybudget related IndexError in PWWS Attack
QData · Jun 4, 2022 · dce49a0 · dce49a0
2 parents 2f3e561 + c15aa1a
commit dce49a0
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/textattack/search_methods/greedy_word_swap_wir.py b/textattack/search_methods/greedy_word_swap_wir.py
@@ -73,14 +73,22 @@ def _get_index_order(self, initial_text):
                     # no valid synonym substitutions for this word
                     delta_ps.append(0.0)
                     continue
-                swap_results, _ = self.get_goal_results(transformed_text_candidates)
+                swap_results, search_over = self.get_goal_results(
+                    transformed_text_candidates
+                )
                 score_change = [result.score for result in swap_results]
                 if not score_change:
                     delta_ps.append(0.0)
                     continue
                 max_score_change = np.max(score_change)
                 delta_ps.append(max_score_change)
 
+                # Exit Loop when search_over is True - but we need to make sure delta_ps
+                # is the same size as softmax_saliency_scores
+                if search_over:
+                    delta_ps = delta_ps + [0.0] * (len_text - len(delta_ps))
+                    break
+
             index_scores = softmax_saliency_scores * np.array(delta_ps)
 
         elif self.wir_method == "delete":