diff --git a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py index eca7208aab..b0a70d57be 100644 --- a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py +++ b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py @@ -445,6 +445,9 @@ def _token_slice( return slice(indices[0], indices[-1] + 1) def _update_ids(self) -> None: + # Gemma and byte-level BPE tokenizers decode the control with the separator space in front; + # rendering that again would add a second space and an extra token before the control. + self.control = self.control.removeprefix(" ") if not self.control.strip() or not self.target.strip(): raise ValueError("GCG control and target must contain non-whitespace text.") if not self.tokenizer.is_fast: diff --git a/tests/unit/executor/promptgen/gcg/test_gcg_core.py b/tests/unit/executor/promptgen/gcg/test_gcg_core.py index e787e84b99..d0419df180 100644 --- a/tests/unit/executor/promptgen/gcg/test_gcg_core.py +++ b/tests/unit/executor/promptgen/gcg/test_gcg_core.py @@ -787,6 +787,18 @@ def test_control_and_content_setters_recompute_boundaries(self, byte_level: bool assert prompt.target_toks.tolist() == control_ids assert prompt._control_slice.stop <= prompt._target_slice.start + def test_decoded_candidate_with_the_separator_space_renders_like_the_original(self) -> None: + tokenizer = _fast_tokenizer(byte_level=True) + original = AttackPrompt(goal="Say it", target="done", tokenizer=tokenizer, control_init="! !") + prompt = AttackPrompt(goal="Say it", target="done", tokenizer=tokenizer, control_init="! !") + + # Gemma and byte-level BPE tokenizers decode the control tokens with the separator space in front. + prompt.control_str = " ! !" + + assert prompt.control == "! !" + assert prompt.input_ids.tolist() == original.input_ids.tolist() + assert prompt._control_slice == original._control_slice + def test_template_without_a_separator_uses_content_positions(self) -> None: tokenizer = _fast_tokenizer(chat_template="{{ messages[0]['content'] }}{{ messages[1]['content'] }}") prompt = AttackPrompt(goal="Say it", target="done", tokenizer=tokenizer, control_init="!")