diff --git a/src/tamperbench/whitebox/attacks/gcg/implementation.py b/src/tamperbench/whitebox/attacks/gcg/implementation.py index 55cb20ad..1909ed91 100644 --- a/src/tamperbench/whitebox/attacks/gcg/implementation.py +++ b/src/tamperbench/whitebox/attacks/gcg/implementation.py @@ -7,7 +7,9 @@ import logging import queue import threading +from collections.abc import Callable from dataclasses import dataclass +from typing import Any, cast import torch import transformers @@ -97,7 +99,7 @@ def should_reduce_batch_size(exception: Exception) -> bool: # modified from https://github.com/huggingface/accelerate/blob/85a75d4c3d0deffde2fc8b917d9b1ae1cb580eb2/src/accelerate/utils/memory.py#L87 -def find_executable_batch_size(function: callable = None, starting_batch_size: int = 128): +def find_executable_batch_size(function: Callable[..., Any] | None = None, starting_batch_size: int = 128) -> Any: """Try to execute `function`, halving batch size on OOM errors. `function` must take in a `batch_size` parameter as its first argument. @@ -254,7 +256,7 @@ def sample_ids_from_grad( search_width: int, topk: int = 256, n_replace: int = 1, - not_allowed_ids: Tensor = False, + not_allowed_ids: Tensor | None = None, ): """Return `search_width` combinations of token ids based on the token gradient. @@ -314,9 +316,9 @@ def filter_ids(ids: Tensor, tokenizer: transformers.PreTrainedTokenizer): for i in range(len(ids_decoded)): # Retokenize the decoded token ids - ids_encoded = tokenizer(ids_decoded[i], return_tensors="pt", add_special_tokens=False).to(ids.device)[ - "input_ids" - ][0] + ids_encoded = cast( + Tensor, tokenizer(ids_decoded[i], return_tensors="pt", add_special_tokens=False)["input_ids"] + ).to(ids.device)[0] if torch.equal(ids[i], ids_encoded): filtered_ids.append(ids[i]) @@ -350,7 +352,7 @@ def __init__( self.tokenizer: transformers.PreTrainedTokenizer = tokenizer self.config: GCGConfig = config - self.embedding_layer: torch.nn.Module = model.get_input_embeddings() + self.embedding_layer: torch.nn.Embedding = cast(torch.nn.Embedding, model.get_input_embeddings()) self.not_allowed_ids: Tensor | None = ( None if config.allow_non_ascii else get_nonascii_toks(tokenizer, device=model.device) ) @@ -426,21 +428,23 @@ def run( messages[-1]["content"] = messages[-1]["content"] + "{optim_str}" template = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + assert isinstance(template, str) # Remove the BOS token -- this will get added when tokenizing, if necessary - if tokenizer.bos_token and template.startswith(tokenizer.bos_token): - template = template.replace(tokenizer.bos_token, "") + bos_token = tokenizer.bos_token + if isinstance(bos_token, str) and template.startswith(bos_token): + template = template.replace(bos_token, "") before_str, after_str = template.split("{optim_str}") target = " " + target if config.add_space_before_target else target # Tokenize everything that doesn't get optimized - before_ids = tokenizer([before_str], padding=False, return_tensors="pt")["input_ids"].to( + before_ids = cast(Tensor, tokenizer([before_str], padding=False, return_tensors="pt")["input_ids"]).to( model.device, torch.int64 ) - after_ids = tokenizer([after_str], add_special_tokens=False, return_tensors="pt")["input_ids"].to( + after_ids = cast(Tensor, tokenizer([after_str], add_special_tokens=False, return_tensors="pt")["input_ids"]).to( model.device, torch.int64 ) - target_ids = tokenizer([target], add_special_tokens=False, return_tensors="pt")["input_ids"].to( + target_ids = cast(Tensor, tokenizer([target], add_special_tokens=False, return_tensors="pt")["input_ids"]).to( model.device, torch.int64 ) @@ -468,15 +472,15 @@ def run( ) # Tokenize everything that doesn't get optimized for the draft model - draft_before_ids = self.draft_tokenizer([before_str], padding=False, return_tensors="pt")["input_ids"].to( - model.device, torch.int64 - ) - draft_after_ids = self.draft_tokenizer([after_str], add_special_tokens=False, return_tensors="pt")[ - "input_ids" - ].to(model.device, torch.int64) - self.draft_target_ids = self.draft_tokenizer([target], add_special_tokens=False, return_tensors="pt")[ - "input_ids" - ].to(model.device, torch.int64) + draft_before_ids = cast( + Tensor, self.draft_tokenizer([before_str], padding=False, return_tensors="pt")["input_ids"] + ).to(model.device, torch.int64) + draft_after_ids = cast( + Tensor, self.draft_tokenizer([after_str], add_special_tokens=False, return_tensors="pt")["input_ids"] + ).to(model.device, torch.int64) + self.draft_target_ids = cast( + Tensor, self.draft_tokenizer([target], add_special_tokens=False, return_tensors="pt")["input_ids"] + ).to(model.device, torch.int64) ( self.draft_before_embeds, @@ -598,13 +602,14 @@ def init_buffer(self) -> AttackBuffer: # Create the attack buffer and initialize the buffer ids buffer = AttackBuffer(config.buffer_size) + init_buffer_ids: Tensor if isinstance(config.optim_str_init, str): - init_optim_ids = tokenizer(config.optim_str_init, add_special_tokens=False, return_tensors="pt")[ - "input_ids" - ].to(model.device) + init_optim_ids = cast( + Tensor, tokenizer(config.optim_str_init, add_special_tokens=False, return_tensors="pt")["input_ids"] + ).to(model.device) if config.buffer_size > 1: init_buffer_ids = ( - tokenizer(INIT_CHARS, add_special_tokens=False, return_tensors="pt")["input_ids"] + cast(Tensor, tokenizer(INIT_CHARS, add_special_tokens=False, return_tensors="pt")["input_ids"]) .squeeze() .to(model.device) ) @@ -623,11 +628,13 @@ def init_buffer(self) -> AttackBuffer: f"Using {len(config.optim_str_init)} initializations but buffer size is set to {config.buffer_size}" ) try: - init_buffer_ids = tokenizer(config.optim_str_init, add_special_tokens=False, return_tensors="pt")[ - "input_ids" - ].to(model.device) - except ValueError: - logger.error("Unable to create buffer. Ensure that all initializations tokenize to the same length.") + init_buffer_ids = cast( + Tensor, tokenizer(config.optim_str_init, add_special_tokens=False, return_tensors="pt")["input_ids"] + ).to(model.device) + except ValueError as e: + raise RuntimeError( + "Unable to create buffer. Ensure that all initializations tokenize to the same length." + ) from e true_buffer_size = max(1, config.buffer_size) @@ -894,13 +901,16 @@ def _compute_draft_losses( def _convert_to_draft_tokens(token_ids: Tensor) -> Tensor: decoded_text_list = self.tokenizer.batch_decode(token_ids) - assert self.draft_tokenizer, "Draft tokenizer wasn't properly initialized." - return self.draft_tokenizer( - decoded_text_list, - add_special_tokens=False, - padding=True, - return_tensors="pt", - )["input_ids"].to(self.draft_model.device, torch.int64) + assert self.draft_tokenizer and self.draft_model, "Draft tokenizer/model wasn't properly initialized." + return cast( + Tensor, + self.draft_tokenizer( + decoded_text_list, + add_special_tokens=False, + padding=True, + return_tensors="pt", + )["input_ids"], + ).to(self.draft_model.device, torch.int64) result_queue = queue.Queue() draft_sampled_ids = _convert_to_draft_tokens(sampled_ids) @@ -933,10 +943,11 @@ def _convert_to_draft_tokens(token_ids: Tensor) -> Tensor: # Step 3. Calculate agreement score using Spearman correlation draft_probe_losses = draft_losses[probe_idxs] - rank_correlation = spearmanr( + spearman_result = spearmanr( probe_losses.cpu().type(torch.float32).numpy(), draft_probe_losses.cpu().type(torch.float32).numpy(), - ).correlation + ) + rank_correlation: float = spearman_result.correlation # pyright: ignore[reportAttributeAccessIssue] # normalized from [-1, 1] to [0, 1] alpha = (1 + rank_correlation) / 2