Skip to content

pyright gcg: Cast a bunch of stuff #79

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
tomtseng merged 1 commit into from
Feb 7, 2026
Merged

pyright gcg: Cast a bunch of stuff #79

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
87 changes: 49 additions & 38 deletions src/tamperbench/whitebox/attacks/gcg/implementation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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.
Expand Down Expand Up @@ -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.

Expand Down Expand Up @@ -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])

Expand Down Expand Up @@ -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)
)
Expand Down Expand Up @@ -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
)

Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)
)
Expand All @@ -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)

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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

Expand Down