diff --git a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py index 9914ffd9fc..b08b55c4cd 100644 --- a/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py +++ b/pyrit/executor/promptgen/gcg/attack/base/attack_manager.py @@ -12,6 +12,7 @@ from dataclasses import dataclass from enum import Enum from typing import TYPE_CHECKING, Any, cast +from uuid import uuid4 import numpy as np import pandas as pd @@ -325,13 +326,17 @@ def __init__( target (str): The target of the attack tokenizer (Transformer Tokenizer): - The tokenizer used to convert text into tokens. Must have a configured chat template - (i.e., ``tokenizer.chat_template`` is not ``None``); ``apply_chat_template`` is used - to render the user/assistant exchange instead of model-specific fastchat templates. + A fast tokenizer with a configured chat template. The template must render each + message once and preserve its content, apart from surrounding whitespace. + Unsupported templates or token boundaries raise an error rather than guessing slices. control_init (str, optional): A string used to control the attack (default is "! ! ! ! ! ! ! ! ! ! ! ! ! ! ! ! ! ! ! !") test_prefixes (list, optional): A list of prefixes to test the attack (default is _DEFAULT_TEST_PREFIXES). + + Raises: + ValueError: If the tokenizer or template cannot provide safe token slices, or the control + or target has no non-whitespace text or tokens. Empty goals are supported for target-only datasets. """ if test_prefixes is None: test_prefixes = list(_DEFAULT_TEST_PREFIXES) @@ -347,69 +352,149 @@ def __init__( self._update_ids() + def _content_bounds(self, *, prompt: str, messages: list[dict[str, str]], message_index: int) -> tuple[int, int]: + """ + Locate one message using a probe that leaves the other message unchanged. + + Anchor both ends of the probe to the complete prompt, without searching for role markers + or separators that may also occur inside the content. + + Args: + prompt (str): The complete rendered conversation. + messages (list[dict[str, str]]): The original user and assistant messages. + message_index (int): The message whose content bounds are needed. + + Returns: + tuple[int, int]: Inclusive start and exclusive end of the rendered content. + + Raises: + ValueError: If the template drops, duplicates, or transforms content, or changes its + surrounding scaffolding when the content is replaced. + """ + marker = f"pyrit{uuid4().hex}" + probe_messages = [dict(message) for message in messages] + probe_messages[message_index]["content"] = marker + scaffold = self.tokenizer.apply_chat_template(probe_messages, tokenize=False) + role = messages[message_index]["role"] + error = ( + f"Cannot safely locate {role} content in the chat template. " + "The template must render each message once and preserve its content, " + "apart from surrounding whitespace." + ) + if not isinstance(scaffold, str) or scaffold.count(marker) != 1: + raise ValueError(error) + + prefix, suffix = scaffold.split(marker) + start, stop = len(prefix), len(prompt) - len(suffix) + if ( + start > stop + or not prompt.startswith(prefix) + or not prompt.endswith(suffix) + or prompt[start:stop].strip() != messages[message_index]["content"].strip() + ): + raise ValueError(error) + return start, stop + + def _token_slice( + self, + *, + prompt: str, + offsets: list[tuple[int, int]], + start: int, + stop: int, + name: str, + allow_empty: bool = False, + ) -> slice: + """ + Map a character span to all its tokens, including repeated byte-level offsets. + + Args: + prompt (str): The complete rendered conversation. + offsets (list[tuple[int, int]]): Character offsets for each token. + start (int): Inclusive character start. + stop (int): Exclusive character end. + name (str): Component name for validation errors. + allow_empty (bool): Whether an empty span is valid. + + Returns: + slice: The corresponding token range. + + Raises: + ValueError: If a required span has no tokens or a token crosses a content boundary. + """ + indices: list[int] = [] + for i, (token_start, token_stop) in enumerate(offsets): + if token_start >= token_stop or token_start >= stop or token_stop <= start: + continue + outside = prompt[token_start:start].strip() or prompt[stop:token_stop].strip() + if outside: + # Added role tokens can consume neighboring whitespace through lstrip/rstrip. + if prompt[max(start, token_start) : min(stop, token_stop)].strip(): + raise ValueError(f"GCG {name} token crosses a content boundary in the chat template.") + continue + indices.append(i) + if start == stop or not indices: + if not allow_empty: + raise ValueError(f"GCG {name} contains no tokens in the rendered prompt.") + boundary = next((i for i, (_, token_stop) in enumerate(offsets) if token_stop > start), len(offsets)) + return slice(boundary, boundary) + + return slice(indices[0], indices[-1] + 1) + def _update_ids(self) -> None: - # Render the goal+control as the user turn and the target as the assistant turn using the - # tokenizer's built-in chat template. This replaces fastchat's per-model Conversation logic - # and works for any HuggingFace chat-tuned model (issue #965). + 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: + raise ValueError("GCG requires a fast tokenizer (use_fast=True) for character-to-token alignment.") messages = [ {"role": "user", "content": f"{self.goal} {self.control}"}, - {"role": "assistant", "content": f"{self.target}"}, + {"role": "assistant", "content": self.target}, ] prompt = self.tokenizer.apply_chat_template(messages, tokenize=False) - - encoding = self.tokenizer(prompt) + user_start, user_end = self._content_bounds(prompt=prompt, messages=messages, message_index=0) + assistant_start, assistant_end = self._content_bounds(prompt=prompt, messages=messages, message_index=1) + if not user_start <= user_end <= assistant_start <= assistant_end: + raise ValueError("Cannot safely locate user and assistant content in conversation order.") + + raw_user = messages[0]["content"] + rendered_user = prompt[user_start:user_end] + raw_leading = len(raw_user) - len(raw_user.lstrip()) + rendered_leading = len(rendered_user) - len(rendered_user.lstrip()) + user_origin = user_start + rendered_leading - raw_leading + goal_start = max(user_start, user_origin) + goal_end = max(goal_start, min(user_end, user_origin + len(self.goal))) + control_start = max(user_start, min(user_end, user_origin + len(self.goal) + 1)) + + # Templates already supply their special tokens. Offset spans handle both unmapped + # whitespace and multiple byte-level tokens sharing the same character position. + encoding = self.tokenizer(prompt, add_special_tokens=False, return_offsets_mapping=True) toks = encoding.input_ids - - # Locate goal/control/target substrings in the rendered prompt. - goal_start = prompt.find(self.goal) - control_start = prompt.find(self.control) - target_start = prompt.find(self.target) - if goal_start == -1 or control_start == -1 or target_start == -1: - raise ValueError( - "Could not locate goal/control/target in chat-templated prompt. " - f"prompt={prompt!r}, goal={self.goal!r}, " - f"control={self.control!r}, target={self.target!r}" - ) - - # ``char_to_token`` returns None when the character index has no - # corresponding token (e.g. when the substring ends exactly at the end - # of the prompt or lands on whitespace squashed into a neighbouring - # token). For end positions we clamp to ``len(toks)``; for start - # positions we walk forward to the next character that does map to a - # token. Both are necessary for the slice arithmetic to remain valid - # across tokenizers/templates. - def end_tok(char_pos: int) -> int: - tok: int | None = encoding.char_to_token(char_pos) - return len(toks) if tok is None else tok - - def start_tok(char_pos: int) -> int: - limit = len(prompt) - cur = char_pos - while cur < limit: - tok: int | None = encoding.char_to_token(cur) - if tok is not None: - return tok - cur += 1 - return len(toks) - - self._goal_slice = slice( - start_tok(goal_start), - end_tok(goal_start + len(self.goal)), + offsets = encoding["offset_mapping"] + goal_slice = self._token_slice( + prompt=prompt, + offsets=offsets, + start=goal_start, + stop=goal_end, + name="goal", + allow_empty=not self.goal.strip(), + ) + control_slice = self._token_slice( + prompt=prompt, offsets=offsets, start=control_start, stop=user_end, name="control" ) - self._control_slice = slice( - start_tok(control_start), - end_tok(control_start + len(self.control)), + target_slice = self._token_slice( + prompt=prompt, offsets=offsets, start=assistant_start, stop=assistant_end, name="target" ) - target_start_tok = start_tok(target_start) - target_end_tok = end_tok(target_start + len(self.target)) - self._target_slice = slice(target_start_tok, target_end_tok) - self._loss_slice = slice(target_start_tok - 1, target_end_tok - 1) - # Assistant role tokens are everything between the control end and the target start. - # This works for any chat template (e.g. llama-2 "[/INST]", phi-3 "<|assistant|>", etc.) - # without us needing to know the literal marker text. - self._assistant_role_slice = slice(self._control_slice.stop, self._target_slice.start) - - self.input_ids = torch.tensor(toks[: self._target_slice.stop], device="cpu") + if not self.goal.strip(): + goal_slice = slice(control_slice.start, control_slice.start) + if goal_slice.stop > control_slice.start or control_slice.stop > target_slice.start: + raise ValueError("GCG token slices overlap across a content boundary.") + + self._goal_slice = goal_slice + self._control_slice = control_slice + self._target_slice = target_slice + self._loss_slice = slice(target_slice.start - 1, target_slice.stop - 1) + self._assistant_role_slice = slice(control_slice.stop, target_slice.start) + self.input_ids = torch.tensor(toks[: target_slice.stop], device="cpu") @torch.no_grad() # type: ignore[misc, untyped-decorator, unused-ignore] def generate(self, model: Any, gen_config: Any = None) -> torch.Tensor: diff --git a/tests/integration/executor/promptgen/gcg/test_gcg_attack_wiring_integration.py b/tests/integration/executor/promptgen/gcg/test_gcg_attack_wiring_integration.py index 9490ac9d56..d6e537ec9a 100644 --- a/tests/integration/executor/promptgen/gcg/test_gcg_attack_wiring_integration.py +++ b/tests/integration/executor/promptgen/gcg/test_gcg_attack_wiring_integration.py @@ -6,7 +6,7 @@ These tests construct real ``IndividualPromptAttack`` / ``ProgressiveMultiPromptAttack`` instances and exercise their wiring all the way down through ``GCGAttackPrompt._update_ids``, which calls ``tokenizer.apply_chat_template`` and -walks character positions via ``char_to_token``. They live here (not under +maps character spans through token offsets. They live here (not under ``tests/unit/``) because they require a real HuggingFace tokenizer (gpt2) to back the chat-template pipeline — mocking that out would defeat the test's purpose, which is to catch kwarg-mismatch and template-compatibility bugs that @@ -75,9 +75,9 @@ def _make_mock_worker_with_real_tokenizer() -> MagicMock: """Worker mock backed by a real gpt2 tokenizer. The wiring tests construct real ``GCGAttackPrompt`` instances which call - ``tokenizer.apply_chat_template`` and then walk character positions in the - rendered prompt. We need a real string + a tokenizer that can answer - ``char_to_token`` queries on it, so we back the mock with the smallest + ``tokenizer.apply_chat_template`` and then map character spans in the + rendered prompt. We need a real string and a tokenizer that provides + offset mappings for it, so we back the mock with the smallest workable real HF tokenizer (gpt2) plus an explicit llama-2-style chat template (gpt2 ships without one). """ diff --git a/tests/unit/executor/promptgen/gcg/test_gcg_core.py b/tests/unit/executor/promptgen/gcg/test_gcg_core.py index efa145b1d2..c9983d65de 100644 --- a/tests/unit/executor/promptgen/gcg/test_gcg_core.py +++ b/tests/unit/executor/promptgen/gcg/test_gcg_core.py @@ -484,136 +484,562 @@ def test_raises_with_multiple_workers(self) -> None: ) +_CHATML_TEMPLATE = ( + "{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] }}<|end|>{% endfor %}" + "{% if add_generation_prompt %}<|assistant|>{% endif %}" +) + +# A complete-conversation template is allowed to ignore add_generation_prompt, so the user-only +# render stops before the assistant marker instead of after it. +_NO_GENERATION_PROMPT_TEMPLATE = "{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] }}<|end|>{% endfor %}" + + +def _fast_tokenizer(*, chat_template: str = _CHATML_TEMPLATE, byte_level: bool = False) -> Any: + """ + Build an offline tokenizer with real offsets, decoding, and special role tokens. + + Args: + chat_template (str): The Jinja chat template to render with. + byte_level (bool): Use byte-level BPE instead of whitespace-delimited words. + + Returns: + Any: A ``PreTrainedTokenizerFast`` with ``chat_template`` set. + """ + from tokenizers import Tokenizer, decoders, models, pre_tokenizers + from transformers import PreTrainedTokenizerFast + + if byte_level: + vocabulary = ["[UNK]", *sorted(pre_tokenizers.ByteLevel.alphabet())] + backend = Tokenizer(models.BPE(dict(zip(vocabulary, range(len(vocabulary)), strict=True)), merges=[])) + backend.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=False) + backend.decoder = decoders.ByteLevel() + else: + vocabulary = [ + "[UNK]", + "Say", + "it", + "assistant", + "done", + "!", + "Sure", + ",", + "here", + "is", + "the", + "plan", + "Respond", + "with", + "now", + "user", + "model", + "goal", + "control", + "target", + "hello", + "world", + ] + backend = Tokenizer( + models.WordLevel(dict(zip(vocabulary, range(len(vocabulary)), strict=True)), unk_token="[UNK]") + ) + backend.pre_tokenizer = pre_tokenizers.Whitespace() + tokenizer = PreTrainedTokenizerFast(tokenizer_object=backend, unk_token="[UNK]") + tokenizer.add_special_tokens( + {"additional_special_tokens": ["<|user|>", "<|assistant|>", "<|end|>", "", ""]} + ) + tokenizer.chat_template = chat_template + return tokenizer + + class TestUpdateIdsErrorPaths: - """Tests covering the error / fallback paths in AttackPrompt._update_ids.""" + """Real-tokenizer coverage of prompt boundaries and unsupported templates.""" - def test_raises_when_substring_not_in_rendered_prompt(self) -> None: - """If the chat template strips/transforms goal/control/target so they don't appear - verbatim in the rendered prompt, _update_ids must raise a clear ValueError.""" - tokenizer = MagicMock() - # Chat template that drops the user content entirely — goal/control won't appear in prompt - tokenizer.apply_chat_template.return_value = "[INST] [/INST] hello" - # tokenizer(...) returns an encoding-like object - encoding = MagicMock() - encoding.input_ids = [1, 2, 3, 4] - encoding.char_to_token.return_value = 1 - tokenizer.return_value = encoding - - with pytest.raises(ValueError, match="Could not locate goal/control/target"): + @pytest.mark.parametrize( + "chat_template", + [ + "<|user|>goal control<|end|><|assistant|>target<|end|>", + "{% for m in messages %}{{ m['content'] | upper }}{% endfor %}", + "{% for m in messages %}{{ m['content'] }}{{ m['content'] }}{% endfor %}", + "{% for m in messages | reverse %}{{ m['content'] }}{% endfor %}", + "{% for m in messages %}{{ m['content'] | replace('goal', 'other') }}{% endfor %}", + "{% for m in messages %}{% if 'goal' in m['content'] %}prefix{% endif %}{{ m['content'] }}{% endfor %}", + ], + ) + def test_unsupported_templates_raise_instead_of_scanning_the_prompt(self, chat_template: str) -> None: + with pytest.raises(ValueError, match="Cannot safely locate"): + AttackPrompt( + goal="goal", + target="target", + tokenizer=_fast_tokenizer(chat_template=chat_template), + control_init="control", + ) + + @pytest.mark.parametrize(("control", "target"), [("", "done"), ("! !", ""), (" ", "done"), ("! !", " ")]) + @pytest.mark.parametrize("byte_level", [False, True]) + def test_empty_or_unmapped_optimization_spans_raise(self, *, control: str, target: str, byte_level: bool) -> None: + with pytest.raises(ValueError, match="control|target"): AttackPrompt( - goal="this-goal-is-missing", - target="this-target-is-missing", - tokenizer=tokenizer, - control_init="this-control-is-missing", + goal="Say it", + target=target, + tokenizer=_fast_tokenizer(byte_level=byte_level), + control_init=control, ) - def test_start_tok_walks_forward_when_initial_position_has_no_token(self) -> None: - """char_to_token returns None for the start position (e.g., whitespace squashed - into the previous token); start_tok must walk forward to the next mappable - character. Slices should still be valid.""" - # Use a fully mocked tokenizer so we can deterministically force char_to_token - # to return None at specific positions, otherwise real tokenizers usually map - # every byte and never trigger the fallback. - prompt_text = "USER hello !! ASSISTANT world" - toks = list(range(15)) - - def char_to_token(pos: int) -> int | None: - # Positions of "h" and "w" both return None; the next char does map. This - # exercises the cur += 1 walk-forward branch in start_tok. - char = prompt_text[pos] if 0 <= pos < len(prompt_text) else "" - if char in ("h", "w"): - return None - # Map remaining positions in a way that preserves slice ordering - return min(pos // 2, len(toks) - 1) - - encoding = MagicMock() - encoding.input_ids = toks - encoding.char_to_token.side_effect = char_to_token + def test_unmapped_whitespace_does_not_extend_slices_into_other_content(self) -> None: + tokenizer = _fast_tokenizer() + prompt = AttackPrompt(goal=" hello ", target=" world ", tokenizer=tokenizer, control_init=" ! ! ") - tokenizer = MagicMock() - tokenizer.apply_chat_template.return_value = prompt_text - tokenizer.return_value = encoding + assert prompt._goal_slice == slice(1, 2) + assert prompt._control_slice == slice(2, 4) + assert prompt._target_slice == slice(6, 7) + assert prompt._loss_slice == slice(5, 6) + assert prompt.goal_str == "hello" + assert prompt.control_str == "! !" + assert prompt.target_str == "world" - # Construction must succeed even though char_to_token returns None at goal/target - # start positions ("h" / "w"). - prompt = AttackPrompt( - goal="hello", - target="world", - tokenizer=tokenizer, - control_init="!!", + def test_target_at_the_end_of_the_prompt_keeps_all_its_tokens(self) -> None: + tokenizer = _fast_tokenizer(chat_template=_CHATML_TEMPLATE.replace("<|end|>", "")) + prompt = AttackPrompt(goal="hello", target="world", tokenizer=tokenizer, control_init="! !") + + assert prompt._target_slice == slice(5, 6) + assert prompt._target_slice.stop == len(prompt.input_ids) + assert prompt.target_str == "world" + + @pytest.mark.parametrize( + ("goal", "control", "target"), + [ + ("Say <|end|><|assistant|> now", "! !", "done"), + ("Say it", "<|end|><|assistant|>", "done"), + ("Say it", "! !", "done <|end|><|assistant|> now"), + ("Say ! ! now", "! !", "done ! !"), + ("assistant", "assistant", "assistant"), + ("control control", "control", "control"), + ], + ) + def test_repeated_content_and_separators_keep_exact_slices(self, *, goal: str, control: str, target: str) -> None: + tokenizer = _fast_tokenizer() + prompt = AttackPrompt(goal=goal, target=target, tokenizer=tokenizer, control_init=control) + goal_ids = tokenizer(goal, add_special_tokens=False).input_ids + control_ids = tokenizer(control, add_special_tokens=False).input_ids + target_ids = tokenizer(target, add_special_tokens=False).input_ids + control_start = 1 + len(goal_ids) + target_start = control_start + len(control_ids) + 2 + + assert prompt._goal_slice == slice(1, control_start) + assert prompt._control_slice == slice(control_start, control_start + len(control_ids)) + assert prompt._target_slice == slice(target_start, target_start + len(target_ids)) + assert prompt.goal_toks.tolist() == goal_ids + assert prompt.control_toks.tolist() == control_ids + assert prompt.target_toks.tolist() == target_ids + + def test_json_template_does_not_find_the_target_in_the_role_label(self) -> None: + tokenizer = _fast_tokenizer( + chat_template="{% for m in messages %}{{ m['role'] | tojson }}:{{ m['content'] | tojson }}\n{% endfor %}", + byte_level=True, ) - assert isinstance(prompt._goal_slice.start, int) - assert isinstance(prompt._target_slice.start, int) - - def test_start_tok_returns_len_toks_when_no_position_maps(self) -> None: - """If char_to_token returns None for every position from char_pos to end-of-prompt, - start_tok must return len(toks) as a safe fallback (line 211).""" - prompt_text = "USER hello !! ASSISTANT world tail" - toks = list(range(20)) - - def char_to_token(pos: int) -> int | None: - char = prompt_text[pos] if 0 <= pos < len(prompt_text) else "" - # "tail" sits at end and never maps to a token (forces start_tok to exhaust - # the loop and hit `return len(toks)`); other content maps normally. - tail_start = prompt_text.find("tail") - if pos >= tail_start: - return None - return min(pos // 2, len(toks) - 1) - - encoding = MagicMock() - encoding.input_ids = toks - encoding.char_to_token.side_effect = char_to_token + prompt = AttackPrompt(goal="Say it", target="assistant", tokenizer=tokenizer, control_init="! !") - tokenizer = MagicMock() - tokenizer.apply_chat_template.return_value = prompt_text - tokenizer.return_value = encoding + assert prompt.target_str == "assistant" + assert prompt.input_str == '"user":"Say it ! !"\n"assistant":"assistant' + assert prompt._assistant_role_slice.stop == prompt._target_slice.start + + def test_transformed_real_content_is_rejected_even_when_markers_round_trip(self) -> None: + tokenizer = _fast_tokenizer( + chat_template="{% for m in messages %}{{ m['role'] | tojson }}:{{ m['content'] | tojson }}\n{% endfor %}" + ) + with pytest.raises(ValueError, match="Cannot safely locate"): + AttackPrompt(goal='Say "hello"', target="done", tokenizer=tokenizer, control_init="! !") + + @pytest.mark.parametrize("target", ["caf\u00e9", "\U0001f600", "done"]) + def test_byte_level_offsets_keep_every_target_byte(self, target: str) -> None: + tokenizer = _fast_tokenizer(byte_level=True) + prompt = AttackPrompt(goal="Say it", target=target, tokenizer=tokenizer, control_init="! !") + + assert prompt.target_toks.tolist() == tokenizer(target, add_special_tokens=False).input_ids + assert prompt.target_str == target + assert prompt.input_str == f"<|user|>Say it ! !<|end|><|assistant|>{target}" + assert prompt._loss_slice == slice(prompt._target_slice.start - 1, prompt._target_slice.stop - 1) + + def test_chat_template_special_tokens_are_not_added_twice(self) -> None: + from tokenizers import processors + + tokenizer = _fast_tokenizer(chat_template="{{ bos_token }}" + _CHATML_TEMPLATE) + tokenizer.add_special_tokens({"bos_token": ""}) + tokenizer.backend_tokenizer.post_processor = processors.TemplateProcessing( + single=" $A", special_tokens=[("", tokenizer.bos_token_id)] + ) + prompt = AttackPrompt(goal="Say it", target="done", tokenizer=tokenizer, control_init="! !") + + assert prompt.input_ids.tolist().count(tokenizer.bos_token_id) == 1 + assert prompt._goal_slice == slice(2, 4) + assert prompt._target_slice == slice(8, 9) + + def test_token_shared_with_assistant_scaffolding_is_rejected(self) -> None: + tokenizer = _fast_tokenizer(chat_template="{{ messages[0]['content'] }} week{{ messages[1]['content'] }}") + with pytest.raises(ValueError, match="token.*boundary"): + AttackPrompt(goal="Say it", target="end", tokenizer=tokenizer, control_init="! !") + + def test_slow_tokenizer_has_an_actionable_error(self) -> None: + from transformers import PreTrainedTokenizer + + class SlowTokenizer(PreTrainedTokenizer): + def get_vocab(self) -> dict[str, int]: + return {"[UNK]": 0} + + def _tokenize(self, text: str, **kwargs: Any) -> list[str]: + return text.split() + + def _convert_token_to_id(self, token: str) -> int: + return 0 + + def _convert_id_to_token(self, index: int) -> str: + return "[UNK]" + + tokenizer = SlowTokenizer(unk_token="[UNK]", chat_template=_CHATML_TEMPLATE) + assert not tokenizer.is_fast + with pytest.raises(ValueError, match="fast tokenizer.*use_fast=True"): + AttackPrompt(goal="Say it", target="done", tokenizer=tokenizer, control_init="! !") + + def test_probe_rendering_errors_are_not_swallowed(self) -> None: + tokenizer = _fast_tokenizer() + with ( + patch.object( + tokenizer, + "apply_chat_template", + side_effect=["<|user|>Say it ! !<|end|><|assistant|>done<|end|>", RuntimeError("probe failed")], + ), + pytest.raises(RuntimeError, match="probe failed"), + ): + AttackPrompt(goal="Say it", target="done", tokenizer=tokenizer, control_init="! !") + + @pytest.mark.parametrize("byte_level", [False, True]) + @pytest.mark.parametrize("goal", ["", " \t", " hello ", "hello"]) + @pytest.mark.parametrize("trim", [False, True]) + def test_whitespace_and_empty_goals_preserve_content(self, *, byte_level: bool, goal: str, trim: bool) -> None: + template = _CHATML_TEMPLATE.replace("m['content']", "m['content'] | trim") if trim else _CHATML_TEMPLATE + tokenizer = _fast_tokenizer(chat_template=template, byte_level=byte_level) + prompt = AttackPrompt(goal=goal, target=" world ", tokenizer=tokenizer, control_init=" ! ! ") + + assert prompt.goal_str == goal.strip() + assert prompt.control_str == "! !" + assert prompt.target_str == "world" + assert prompt._goal_slice.stop <= prompt._control_slice.start < prompt._control_slice.stop + assert prompt._control_slice.stop <= prompt._target_slice.start < prompt._target_slice.stop + + @pytest.mark.parametrize("byte_level", [False, True]) + def test_control_and_content_setters_recompute_boundaries(self, byte_level: bool) -> None: + tokenizer = _fast_tokenizer(byte_level=byte_level) + prompt = AttackPrompt(goal="Say it", target="done", tokenizer=tokenizer, control_init="! !") + + for control in ("assistant", "<|end|><|assistant|>", "! !"): + prompt.control_str = control + assert prompt.control_toks.tolist() == tokenizer(control, add_special_tokens=False).input_ids + assert prompt.goal_str == "Say it" + assert prompt.target_str == "done" + + control_ids = tokenizer("assistant", add_special_tokens=False).input_ids + prompt.control_toks = torch.tensor(control_ids) + prompt.goal_str = "" + prompt.target_str = "assistant" + assert prompt.goal_toks.numel() == 0 + assert prompt.control_toks.tolist() == control_ids + assert prompt.target_toks.tolist() == control_ids + assert prompt._control_slice.stop <= prompt._target_slice.start + + 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="!") + + assert prompt._control_slice == slice(2, 3) + assert prompt._target_slice == slice(3, 4) + assert prompt._assistant_role_slice == slice(3, 3) + + @pytest.mark.parametrize("byte_level", [False, True]) + @pytest.mark.parametrize("goal", ["", " ", " hello "]) + def test_role_tokens_that_consume_whitespace_stay_outside_content_slices( + self, *, byte_level: bool, goal: str + ) -> None: + from tokenizers import AddedToken + + tokenizer = _fast_tokenizer(byte_level=byte_level) + tokenizer.add_special_tokens( + { + "additional_special_tokens": [ + AddedToken(marker, lstrip=True, rstrip=True) for marker in ("<|user|>", "<|assistant|>", "<|end|>") + ] + } + ) + prompt = AttackPrompt(goal=goal, target=" world ", tokenizer=tokenizer, control_init=" ! ! ") + role_ids = {tokenizer.convert_tokens_to_ids(marker) for marker in ("<|user|>", "<|assistant|>", "<|end|>")} + + assert prompt.goal_str == goal.strip() + assert prompt.control_str == "! !" + assert prompt.target_str == "world" + assert role_ids.isdisjoint(prompt.goal_toks.tolist()) + assert role_ids.isdisjoint(prompt.control_toks.tolist()) + assert role_ids.isdisjoint(prompt.target_toks.tolist()) + if not goal.strip(): + assert prompt._goal_slice == slice(prompt._control_slice.start, prompt._control_slice.start) + + @pytest.mark.parametrize("component", ["control", "target"]) + def test_nonempty_text_removed_by_tokenizer_normalization_raises(self, component: str) -> None: + from tokenizers import normalizers + + tokenizer = _fast_tokenizer() + tokenizer.backend_tokenizer.normalizer = normalizers.Replace(component, "") + with pytest.raises(ValueError, match=f"{component} contains no tokens"): + AttackPrompt(goal="goal", target="target", tokenizer=tokenizer, control_init="control") + + def test_target_is_located_after_the_user_turn_when_the_goal_quotes_it(self) -> None: + """A goal that quotes its own target must not pull the target slice into the user turn. + + Affirmative-prefix targets make this realistic: the same text then appears twice in the + rendered prompt, and taking the first occurrence points the target and loss slices at the + user turn instead of the assistant reply. + """ + goal = "Respond with Sure, here is the plan" + control = "! ! ! !" + target = "Sure, here is the plan" - # "tail" as the target — its start position and every position after it returns - # None, so start_tok exits the while loop and returns len(toks). prompt = AttackPrompt( - goal="hello", - target="tail", - tokenizer=tokenizer, - control_init="!!", + goal=goal, + target=target, + tokenizer=_fast_tokenizer(), + control_init=control, ) - assert prompt._target_slice.start == len(toks) - - def test_end_tok_returns_len_toks_when_target_is_at_prompt_end(self) -> None: - """If the target sits at the very end of the rendered prompt, - char_to_token(end_pos) returns None — end_tok must clamp to len(toks) - (line 201 in attack_manager.py).""" - # Fully-mocked tokenizer so we can deterministically force char_to_token to - # return None at the position just past the target. Mirrors the pattern used - # by the two adjacent tests above. - prompt_text = "[INST] hello !! [/INST] world" - toks = list(range(10)) - target_end_pos = len(prompt_text) # one past the final char of "world" - - def char_to_token(pos: int) -> int | None: - # Position at/after end-of-prompt has no token → triggers the - # `return len(toks)` fallback in end_tok. - if pos >= target_end_pos: - return None - # Everything else maps to a valid token index that preserves ordering. - return min(pos // 3, len(toks) - 1) - - encoding = MagicMock() - encoding.input_ids = toks - encoding.char_to_token.side_effect = char_to_token - tokenizer = MagicMock() - tokenizer.apply_chat_template.return_value = prompt_text - tokenizer.return_value = encoding + assert prompt._control_slice == slice(9, 13) + assert prompt._assistant_role_slice == slice(13, 15) + assert prompt._target_slice == slice(15, 21) + assert prompt._loss_slice == slice(14, 20) + assert prompt.target_toks.tolist() == prompt.tokenizer(target, add_special_tokens=False).input_ids + + def test_target_that_names_the_assistant_role_marker_is_found_in_the_reply(self) -> None: + """A target like "assistant" also matches inside ``<|assistant|>``, which is one special token. + + Searching right after the user content lands on the role marker and leaves an empty target + slice, so the search has to start where the assistant content does. + """ + tokenizer = _fast_tokenizer() + + prompt = AttackPrompt(goal="Say it", target="assistant", tokenizer=tokenizer, control_init="! !") + + ids = tokenizer("<|user|>Say it ! !<|end|><|assistant|>assistant<|end|>").input_ids + # <|user|> Say it ! ! <|end|> <|assistant|> assistant <|end|> + assert prompt._control_slice == slice(3, 5) + assert prompt._target_slice == slice(7, 8) + assert prompt._loss_slice == slice(6, 7) + assert ids[6] == tokenizer.convert_tokens_to_ids("<|assistant|>") + + def test_control_that_collides_with_the_role_marker_keeps_its_slice(self) -> None: + """An optimized control is decoded vocabulary tokens, so it can contain ordinary words. + + A control holding "assistant" also matches inside ``<|assistant|>``; bounding the search to + the user content is what keeps the control slice pointing at the suffix GCG optimizes. + """ + tokenizer = _fast_tokenizer() + + prompt = AttackPrompt(goal="Say it", target="done", tokenizer=tokenizer, control_init="assistant") + + # <|user|> Say it assistant <|end|> <|assistant|> done <|end|> + assert prompt._control_slice == slice(3, 4) + assert prompt._target_slice == slice(6, 7) + assert prompt._loss_slice == slice(5, 6) + + def test_boundary_holds_when_the_template_ignores_the_generation_prompt(self) -> None: + """``add_generation_prompt`` is documented as a no-op for templates that do not support it. + + The user-only render is still a prefix of the full prompt, so it cannot be trusted as the + assistant boundary: it stops before ``<|assistant|>`` and a target of "assistant" would + match the role marker again. + """ + tokenizer = _fast_tokenizer(chat_template=_NO_GENERATION_PROMPT_TEMPLATE) + + prompt = AttackPrompt(goal="Say it", target="assistant", tokenizer=tokenizer, control_init="! !") + + # <|user|> Say it ! ! <|end|> <|assistant|> assistant <|end|> + assert prompt._control_slice == slice(3, 5) + assert prompt._target_slice == slice(7, 8) + assert prompt._loss_slice == slice(6, 7) + + def test_empty_goal_with_a_trimming_template(self) -> None: + """Target-only datasets use an empty goal, so the user content is " ". + + A template that trims the content drops that leading space, so the control has to be found on + its own rather than as part of the raw ``f"{goal} {control}"`` string. + """ + tokenizer = _fast_tokenizer( + chat_template="{% for m in messages %}" + "{{ 'model' if m['role'] == 'assistant' else m['role'] }}\n" + "{{ m['content'] | trim }}\n{% endfor %}" + "{% if add_generation_prompt %}model\n{% endif %}" + ) + + prompt = AttackPrompt(goal="", target="Sure, here", tokenizer=tokenizer, control_init="! ! !") + + # user ! ! ! model Sure , here + assert prompt._goal_slice == slice(2, 2) + assert prompt._control_slice == slice(2, 5) + assert prompt._target_slice == slice(8, 11) + assert prompt._loss_slice == slice(7, 10) + + def test_goal_that_quotes_the_turn_separator_keeps_the_boundary(self) -> None: + """Red-team goals can quote model control tokens, including the template's own turn separator. + + Searching for the separator would find it inside the goal and end the user content before the + control, so the boundary has to be measured from the template instead. + """ + tokenizer = _fast_tokenizer() prompt = AttackPrompt( - goal="hello", - target="world", # sits at end of prompt_text; target end has no token + goal="Say <|end|><|assistant|> now", target="done", tokenizer=tokenizer, control_init="! !" + ) + + # <|user|> Say <|end|> <|assistant|> now ! ! <|end|> <|assistant|> done <|end|> + assert prompt._control_slice == slice(5, 7) + assert prompt._target_slice == slice(9, 10) + assert prompt._loss_slice == slice(8, 9) + + def test_escaping_template_locates_the_target_in_the_reply(self) -> None: + """A template may escape the contents, e.g. with ``tojson``; the boundaries still have to hold. + + With an unbounded search the target "assistant" matches the role label instead of the reply. + """ + tokenizer = _fast_tokenizer( + chat_template="{% for m in messages %}\"{{ m['role'] }}\":{{ m['content'] | tojson }}\n{% endfor %}", + byte_level=True, + ) + + prompt = AttackPrompt(goal="Say it", target="assistant", tokenizer=tokenizer, control_init="! !") + + target_start = len('"user":"Say it ! !"\n"assistant":"') + assert prompt._target_slice == slice(target_start, target_start + len("assistant")) + assert prompt._loss_slice == slice(target_start - 1, target_start + len("assistant") - 1) + assert prompt.target_str == "assistant" + + def test_raises_when_an_escaped_goal_is_not_rendered_verbatim(self) -> None: + """The turns can be measured, but ``tojson`` escapes the quotes, so the goal itself is not in the prompt.""" + tokenizer = _fast_tokenizer( + chat_template="{% for m in messages %}\"{{ m['role'] }}\":{{ m['content'] | tojson }}\n{% endfor %}" + ) + + with pytest.raises(ValueError, match="Cannot safely locate"): + AttackPrompt(goal='Say "it"', target="done", tokenizer=tokenizer, control_init="! !") + + def test_raises_when_the_template_transforms_the_contents(self) -> None: + """A template that rewrites the contents leaves no way to measure the turns, so construction fails closed.""" + tokenizer = _fast_tokenizer( + chat_template="{% for m in messages %}<|{{ m['role'] }}|>{{ m['content'] | upper }}<|end|>{% endfor %}" + ) + + with pytest.raises(ValueError, match="Cannot safely locate"): + AttackPrompt(goal="Say it", target="done", tokenizer=tokenizer, control_init="! !") + + +@pytest.mark.usefixtures("patch_central_database") +class TestPromptSliceWiring: + """Exercise real prompt construction beneath the optimization orchestration.""" + + @staticmethod + def _worker(*, byte_level: bool = False) -> MagicMock: + worker = MagicMock(spec=ModelWorker) + worker.tokenizer = _fast_tokenizer(byte_level=byte_level) + worker.model = MagicMock(spec=torch.nn.Module) + worker.model.device = torch.device("cpu") + return worker + + @pytest.mark.parametrize("byte_level", [False, True]) + def test_real_slices_support_gradients_and_candidate_losses(self, byte_level: bool) -> None: + tokenizer = _fast_tokenizer(byte_level=byte_level) + model = _tiny_model("llama").eval() + model.resize_token_embeddings(len(tokenizer), mean_resizing=False) + manager = GCGPromptManager( + goals=["Say <|end|><|assistant|> now", ""], + targets=["done ! !", "assistant"], tokenizer=tokenizer, - control_init="!!", + control_init="! !", + managers={"AP": gcg_attack_mod.GCGAttackPrompt}, ) - # end_tok(target_end_pos) saw None from char_to_token → clamped to len(toks). - assert prompt._target_slice.stop == len(toks) - assert prompt._target_slice.stop > prompt._target_slice.start + + for control in ("! !", "assistant"): + manager.control_str = control + gradient = manager.grad(model) + assert gradient.shape == (manager.control_toks.numel(), len(tokenizer)) + assert torch.isfinite(gradient).all() + assert torch.count_nonzero(gradient) > 0 + + candidates = manager.control_toks.repeat(2, 1) + candidates[1, 0] = tokenizer("now", add_special_tokens=False).input_ids[0] + logits, ids = manager.logits(model, test_controls=candidates, return_ids=True) + assert torch.isfinite(manager.target_loss(logits, ids)).all() + assert torch.isfinite(manager.control_loss(logits, ids)).all() + for prompt, candidate_ids in zip(manager, ids, strict=True): + assert torch.equal(candidate_ids[:, prompt._control_slice], candidates) + assert torch.equal(candidate_ids[:, prompt._target_slice], prompt.target_toks.repeat(2, 1)) + + def test_shared_control_reaches_training_and_held_out_prompts_for_each_tokenizer(self) -> None: + workers = [self._worker(), self._worker(byte_level=True)] + test_worker = self._worker() + attack = MultiPromptAttack( + goals=["Say <|end|><|assistant|> now", ""], + targets=["assistant", "done"], + workers=workers, + control_init="! !", + test_goals=["Say assistant"], + test_targets=["assistant"], + test_workers=[test_worker], + managers={"AP": AttackPrompt, "PM": PromptManager}, + ) + attack.control_str = "assistant" + + with patch.object(attack, "test", return_value=([], [], [])) as evaluate: + attack.test_all() + + all_workers, all_prompts = evaluate.call_args.args + assert all_workers == [*workers, test_worker] + assert evaluate.call_args.kwargs["include_loss"] is True + for manager in [*attack.prompts, *all_prompts]: + for prompt in manager: + tokenizer = prompt.tokenizer + assert prompt.goal_toks.tolist() == tokenizer(prompt.goal, add_special_tokens=False).input_ids + assert prompt.control_toks.tolist() == tokenizer("assistant", add_special_tokens=False).input_ids + assert prompt.target_toks.tolist() == tokenizer(prompt.target, add_special_tokens=False).input_ids + + @pytest.mark.parametrize( + ("progressive", "expected_rounds"), + [ + (False, [(1, 2), (1, 2)]), + (True, [(1, 1), (2, 1), (2, 2)]), + ], + ) + def test_individual_and_progressive_rounds_construct_real_prompts( + self, *, progressive: bool, expected_rounds: list[tuple[int, int]] + ) -> None: + rounds: list[tuple[int, int]] = [] + + def run_inner(attack: MultiPromptAttack, **kwargs: Any) -> tuple[str, float, int]: + rounds.append((len(attack.goals), len(attack.workers))) + attack.control_str = "assistant" + for manager in attack.prompts: + for prompt in manager: + assert prompt.control_str == "assistant" + assert ( + prompt.target_toks.tolist() + == prompt.tokenizer(prompt.target, add_special_tokens=False).input_ids + ) + return "assistant", 0.5, 1 + + attack_class = ProgressiveMultiPromptAttack if progressive else IndividualPromptAttack + attack = attack_class( + goals=["Say <|end|><|assistant|> now", ""], + targets=["assistant", "done"], + workers=[self._worker(), self._worker(byte_level=True)], + control_init="! !", + managers={"AP": AttackPrompt, "PM": PromptManager, "MPA": MultiPromptAttack}, + ) + with patch.object(MultiPromptAttack, "run", autospec=True, side_effect=run_inner): + attack.run(n_steps=3, stop_on_success=False, incr_control=False, verbose=False) + + assert rounds == expected_rounds class TestGetWorkersChatTemplateValidation: