Flatten 1_raw_dataset submodules into plain tracked files
FocalLoRA, Should-It-Be-Executed-Or-Processed, and topicattack were nested git repos (with an inner FocalLoRA/data/FocalLoRA/.git as well). Drop their .git history and track the contents directly in this repo instead of as submodules/gitlinks. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
HenryChou020514
@ -0,0 +1,205 @@
|
||||
'''A main script to run attack for LLMs.'''
|
||||
import argparse
|
||||
import time
|
||||
import importlib
|
||||
|
||||
from typing import Any
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
sys.path.append("../rpo/") # Adds higher directory to python modules path.
|
||||
|
||||
from rpo.suffix_manager import get_goals_and_targets, get_workers
|
||||
|
||||
from huggingface_hub import login
|
||||
|
||||
def dynamic_import(module: str):
|
||||
"""
|
||||
Dynamically import a module given its name as a string.
|
||||
|
||||
Parameters:
|
||||
module (str): The name of the module to import.
|
||||
|
||||
Returns:
|
||||
module: The imported module object.
|
||||
|
||||
Example:
|
||||
>>> math_module = dynamic_import('math')
|
||||
>>> math_module.sqrt(16)
|
||||
4.0
|
||||
|
||||
Raises:
|
||||
ImportError: If the module cannot be imported.
|
||||
"""
|
||||
return importlib.import_module(module)
|
||||
|
||||
def main(params: Any) -> None:
|
||||
"""
|
||||
Main function to run a Progressive Multi-Prompt Attack using specified parameters.
|
||||
|
||||
Parameters:
|
||||
params (Any): The parameters required to configure and run the attack, typically
|
||||
provided through a configuration object or command-line arguments.
|
||||
|
||||
The function performs the following steps:
|
||||
1. Retrieves the Hugging Face Hub token from environment variables and logs in if available.
|
||||
2. Dynamically imports the attack library.
|
||||
3. Initializes workers and data loaders for training and testing.
|
||||
4. Configures managers for the attack.
|
||||
5. Creates an instance of the ProgressiveMultiPromptAttack class with specified parameters.
|
||||
6. Runs the attack with specified parameters.
|
||||
7. Stops all workers after the attack is complete and prints "SUCCESS".
|
||||
"""
|
||||
access_token = os.environ.get("HUGGING_FACE_HUB_TOKEN")
|
||||
if access_token:
|
||||
login(token=access_token)
|
||||
attack_lib = dynamic_import(f'rpo')
|
||||
|
||||
workers, test_workers = get_workers(params)
|
||||
|
||||
|
||||
train_loader, test_loader = get_goals_and_targets(params)
|
||||
|
||||
|
||||
managers = {
|
||||
"AP": attack_lib.AttackPrompt,
|
||||
"PM": attack_lib.PromptManager,
|
||||
"MPA": attack_lib.MultiPromptAttack,
|
||||
}
|
||||
|
||||
timestamp = time.strftime("%Y%m%d-%H:%M:%S")
|
||||
|
||||
|
||||
attack = attack_lib.ProgressiveMultiPromptAttack(
|
||||
train_loader,
|
||||
test_loader,
|
||||
workers,
|
||||
model_name=params.model,
|
||||
progressive_models=params.progressive_models,
|
||||
progressive_goals=params.progressive_goals,
|
||||
control_init=params.control_init,
|
||||
safe_init=params.safe_init,
|
||||
logfile=f"{params.result_prefix}/{params.model}_{timestamp}_cut_cand_i_len_to_20.json",
|
||||
managers=managers,
|
||||
test_workers=test_workers,
|
||||
mpa_deterministic=params.gbda_deterministic,
|
||||
mpa_lr=params.lr,
|
||||
mpa_batch_size=params.batch_size,
|
||||
mpa_n_steps=params.n_steps,
|
||||
)
|
||||
|
||||
attack.run(
|
||||
n_epochs=params.n_epochs,
|
||||
batch_size=params.batch_size,
|
||||
topk=params.topk,
|
||||
temp=params.temp,
|
||||
target_weight=params.target_weight,
|
||||
control_weight=params.control_weight,
|
||||
test_steps=getattr(params, 'test_steps', 1),
|
||||
anneal=params.anneal,
|
||||
incr_control=params.incr_control,
|
||||
stop_on_success=params.stop_on_success,
|
||||
verbose=params.verbose,
|
||||
filter_cand=params.filter_cand,
|
||||
allow_non_ascii=(params.allow_non_ascii == "True"),
|
||||
selection_interval=params.selection_interval
|
||||
)
|
||||
|
||||
for worker in workers + test_workers:
|
||||
worker.stop()
|
||||
print("SUCCESS")
|
||||
|
||||
def set_config_default(config: Any) -> Any:
|
||||
"""
|
||||
Set default configuration parameters for the attack.
|
||||
|
||||
Parameters:
|
||||
config (Any): The configuration object to set default values for.
|
||||
|
||||
Returns:
|
||||
Any: The configuration object with default values set.
|
||||
"""
|
||||
config.target_weight = 1.0
|
||||
config.control_weight = 0.0
|
||||
config.progressive_goals = False
|
||||
config.progressive_models = False
|
||||
config.anneal = False
|
||||
config.incr_control = False
|
||||
config.stop_on_success = False
|
||||
config.verbose = True
|
||||
config.num_train_models = 1
|
||||
config.selection_interval = 100
|
||||
config.data_offset = 0
|
||||
|
||||
# attack-related parameters
|
||||
config.lr = 0.01
|
||||
config.topk = 256
|
||||
config.temp = 1
|
||||
config.filter_cand = True
|
||||
|
||||
config.gbda_deterministic = True
|
||||
return config
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
parser = argparse.ArgumentParser(description='A main script to run attack for LLMs.')
|
||||
|
||||
# Replace these with your actual command-line arguments
|
||||
parser.add_argument("--model", type=str, help='Model name.', default="llama-2")
|
||||
|
||||
parser.add_argument("--attack", type=str, help='Attack type.',
|
||||
default="gcg")
|
||||
parser.add_argument('--train_data', type=str, help='Path to train data.',
|
||||
default="")
|
||||
parser.add_argument('--test_data', type=str, help='Path to test data.',
|
||||
default="")
|
||||
parser.add_argument('--result_prefix', type=str, help='Prefix for result files.',
|
||||
default=f"./experiments/sep_results") # add model during saving!!!!
|
||||
parser.add_argument('--control_init', type=str, help='Initial control setting.',
|
||||
default="! ! ! ! ! ! ! ! ! ! ! ! ! ! ! ! ! ! ! !")
|
||||
parser.add_argument('--safe_init', type=str, help='Initial safe setting.',
|
||||
default="! ! ! ! ! ! ! ! ! ! ! ! ! ! ! ! ! ! ! !")
|
||||
parser.add_argument('--progressive_models',
|
||||
help='Use progressive models. Defaults to False.', default=False)
|
||||
parser.add_argument('--progressive_goals',
|
||||
help='Use progressive goals. Defaults to False.', default=False)
|
||||
parser.add_argument('--stop_on_success',
|
||||
help='Stop on success. Defaults to False.', default=False)
|
||||
parser.add_argument('--num_train_models', type=int,
|
||||
help='Number of training models.', default=1)
|
||||
parser.add_argument('--allow_non_ascii',
|
||||
help='Allow non-ASCII characters. Defaults to False.', default=False)
|
||||
parser.add_argument('--n_test_data', type=int, help='Number of test data points.',
|
||||
default=4)
|
||||
parser.add_argument('--n_steps', type=int, help='Number of steps.', default=2000)
|
||||
parser.add_argument('--n_epochs', type=int, help='Number of steps.', default=1)
|
||||
parser.add_argument('--test_steps', type=int, help='Number of test steps.', default=1)
|
||||
parser.add_argument('--batch_size', type=int, help='Batch size for tokens.', default=8)
|
||||
parser.add_argument('--data_batch_size', type=int, help='Batch size for data.', default=3)
|
||||
parser.add_argument('--steps_per_data_batch', type=int, help='Batch size for data.', default=20)
|
||||
|
||||
parser.add_argument('--selection_interval', type=int, help='Selection interval.',
|
||||
default=100)
|
||||
parser.add_argument('--transfer', type=str, help='Do transfer.',
|
||||
default=True)
|
||||
parser.add_argument('--gbda_deterministic', type=str, help='Is GDBA deterministic.',
|
||||
default=True)
|
||||
parser.add_argument('--tokenizer_paths', type=json.loads, help='Tokenizer paths.',
|
||||
default=("meta-llama/Llama-2-7b-chat-hf",))
|
||||
parser.add_argument('--model_paths', type=json.loads, help='Tokenizer paths.',
|
||||
default=("meta-llama/Llama-2-7b-chat-hf",))
|
||||
parser.add_argument('--tokenizer_kwargs', type=tuple, help='Tokenizer kwargs.',
|
||||
default=(({"use_fast": False}, )))
|
||||
parser.add_argument('--model_kwargs', type=tuple, help='Model kwargs.',
|
||||
default=(({"low_cpu_mem_usage": True, "use_cache": True}, )))
|
||||
parser.add_argument('--conversation_templates', type=json.loads, help='Conv templates.',
|
||||
default=("llama-2",))
|
||||
|
||||
|
||||
parser.add_argument('--devices', type=tuple, help='Devices.',
|
||||
default=("cuda:0", "cuda:1", "cuda:2", "cuda:3"))
|
||||
|
||||
# This line parses the arguments
|
||||
args = parser.parse_args()
|
||||
args = set_config_default(args)
|
||||
main(args)
|
||||
@ -0,0 +1 @@
|
||||
Note: our prompt optimization procedure is based on this repository created by Zhou et al., 2024 (see the paper). Here we only include files that differ from the original repo.
|
||||
@ -0,0 +1,18 @@
|
||||
__version__ = '0.0.1'
|
||||
|
||||
from .suffix_manager import (
|
||||
AttackPrompt,
|
||||
PromptManager,
|
||||
MultiPromptAttack,
|
||||
ProgressiveMultiPromptAttack,
|
||||
get_embedding_layer,
|
||||
get_embedding_matrix,
|
||||
get_embeddings,
|
||||
get_nonascii_toks,
|
||||
get_goals_and_targets,
|
||||
get_workers
|
||||
)
|
||||
|
||||
from .gcg import GCGAttackPrompt as AttackPrompt
|
||||
from .gcg import GCGPromptManager as PromptManager
|
||||
from .gcg import GCGMultiPromptAttack as MultiPromptAttack
|
||||
@ -0,0 +1,213 @@
|
||||
import gc
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from rpo import AttackPrompt, MultiPromptAttack, PromptManager
|
||||
from rpo import get_embedding_matrix, get_embeddings
|
||||
|
||||
|
||||
def print_gpu_memory(label):
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
if device.type == 'cuda':
|
||||
torch.cuda.synchronize() # Wait for all operations to complete
|
||||
total_memory = torch.cuda.get_device_properties(device).total_memory
|
||||
allocated_memory = torch.cuda.memory_allocated(device)
|
||||
cached_memory = torch.cuda.memory_reserved(device)
|
||||
free_memory = total_memory - allocated_memory
|
||||
print(f"Memory log for cuda:{torch.cuda.current_device()}. Label::: {label}:" )
|
||||
print(f"Total GPU Memory: {total_memory / 1e9:.2f} GB")
|
||||
print(f"Allocated Memory: {allocated_memory / 1e9:.2f} GB")
|
||||
print(f"Cached Memory: {cached_memory / 1e9:.2f} GB")
|
||||
print(f"Free Memory: {free_memory / 1e9:.2f} GB")
|
||||
else:
|
||||
print("No CUDA device available")
|
||||
|
||||
def token_gradients(model, input_ids, input_slice, target_slice, loss_slice):
|
||||
|
||||
"""
|
||||
Computes gradients of the loss with respect to the coordinates.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
model : Transformer Model
|
||||
The transformer model to be used.
|
||||
input_ids : torch.Tensor
|
||||
The input sequence in the form of token ids.
|
||||
input_slice : slice
|
||||
The slice of the input sequence for which gradients need to be computed.
|
||||
target_slice : slice
|
||||
The slice of the input sequence to be used as targets.
|
||||
loss_slice : slice
|
||||
The slice of the logits to be used for computing the loss.
|
||||
|
||||
Returns
|
||||
-------
|
||||
torch.Tensor
|
||||
The gradients of each token in the input_slice with respect to the loss.
|
||||
"""
|
||||
|
||||
embed_weights = get_embedding_matrix(model)
|
||||
one_hot = torch.zeros(
|
||||
input_ids[input_slice].shape[0],
|
||||
embed_weights.shape[0],
|
||||
device=model.device,
|
||||
dtype=embed_weights.dtype
|
||||
)
|
||||
one_hot.scatter_(
|
||||
1,
|
||||
input_ids[input_slice].unsqueeze(1),
|
||||
torch.ones(one_hot.shape[0], 1, device=model.device, dtype=embed_weights.dtype)
|
||||
)
|
||||
one_hot.requires_grad_()
|
||||
input_embeds = (one_hot @ embed_weights).unsqueeze(0)
|
||||
|
||||
# now stitch it together with the rest of the embeddings
|
||||
embeds = get_embeddings(model, input_ids.unsqueeze(0)).detach()
|
||||
full_embeds = torch.cat(
|
||||
[
|
||||
embeds[:,:input_slice.start,:],
|
||||
input_embeds,
|
||||
embeds[:,input_slice.stop:,:]
|
||||
],
|
||||
dim=1)
|
||||
|
||||
logits = model(inputs_embeds=full_embeds).logits
|
||||
targets = input_ids[target_slice]
|
||||
loss = nn.CrossEntropyLoss()(logits[0,loss_slice,:], targets)
|
||||
|
||||
loss.backward()
|
||||
|
||||
return one_hot.grad.clone()
|
||||
|
||||
class GCGAttackPrompt(AttackPrompt):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def grad(self, model):
|
||||
return token_gradients(
|
||||
model,
|
||||
self.input_ids.to(model.device),
|
||||
self._control_slice,
|
||||
self._target_slice,
|
||||
self._loss_slice
|
||||
)
|
||||
|
||||
class GCGPromptManager(PromptManager):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def sample_control(self, grad, batch_size, topk=256, temp=1, allow_non_ascii=True):
|
||||
|
||||
if not allow_non_ascii:
|
||||
grad[:, self._nonascii_toks.to(grad.device)] = np.infty
|
||||
top_indices = (-grad).topk(topk, dim=1).indices
|
||||
control_toks = self.control_toks.to(grad.device)
|
||||
original_control_toks = control_toks.repeat(batch_size, 1)
|
||||
new_token_pos = torch.arange(
|
||||
0,
|
||||
len(control_toks),
|
||||
len(control_toks) / batch_size,
|
||||
device=grad.device
|
||||
).type(torch.int64)
|
||||
new_token_val = torch.gather(
|
||||
top_indices[new_token_pos], 1,
|
||||
torch.randint(0, topk, (batch_size, 1),
|
||||
device=grad.device)
|
||||
)
|
||||
new_control_toks = original_control_toks.scatter_(1, new_token_pos.unsqueeze(-1), new_token_val)
|
||||
return new_control_toks
|
||||
|
||||
|
||||
|
||||
class GCGMultiPromptAttack(MultiPromptAttack):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
def step(self,
|
||||
batch_size=1024,
|
||||
topk=256,
|
||||
temp=1,
|
||||
allow_non_ascii=True,
|
||||
target_weight=1,
|
||||
control_weight=0.1,
|
||||
verbose=False,
|
||||
opt_only=False,
|
||||
filter_cand=True):
|
||||
|
||||
|
||||
# GCG currently does not support optimization_only mode,
|
||||
# so opt_only does not change the inner loop.
|
||||
opt_only = False
|
||||
|
||||
main_device = self.models[0].device
|
||||
control_cands = []
|
||||
|
||||
for j, worker in enumerate(self.workers):
|
||||
worker(self.prompts[j], "grad", worker.model)
|
||||
|
||||
# Aggregate gradients
|
||||
grad = None
|
||||
for j, worker in enumerate(self.workers):
|
||||
new_grad = worker.results.get().to(main_device)
|
||||
new_grad = new_grad / new_grad.norm(dim=-1, keepdim=True)
|
||||
if grad is None:
|
||||
grad = torch.zeros_like(new_grad)
|
||||
if grad.shape != new_grad.shape:
|
||||
with torch.no_grad():
|
||||
control_cand = self.prompts[j-1].sample_control(grad, batch_size, topk, temp, allow_non_ascii)
|
||||
control_cands.append(self.get_filtered_cands(j-1, control_cand, filter_cand=filter_cand, curr_control=self.control_str))
|
||||
grad = new_grad
|
||||
else:
|
||||
grad += new_grad
|
||||
with torch.no_grad():
|
||||
control_cand = self.prompts[j].sample_control(grad, batch_size, topk, temp, allow_non_ascii)
|
||||
control_cands.append(self.get_filtered_cands(j, control_cand, filter_cand=filter_cand, curr_control=self.control_str))
|
||||
|
||||
del grad, control_cand ; gc.collect()
|
||||
|
||||
|
||||
# Search
|
||||
loss = torch.zeros(len(control_cands) * batch_size).to(main_device)
|
||||
with torch.no_grad():
|
||||
for j, cand in enumerate(control_cands):
|
||||
# Looping through the prompts at this level is less elegant, but
|
||||
# we can manage VRAM better this way
|
||||
progress = tqdm(range(len(self.prompts[0])), total=len(self.prompts[0])) if verbose else enumerate(self.prompts[0])
|
||||
for i in progress:
|
||||
for k, worker in enumerate(self.workers):
|
||||
worker(self.prompts[k][i], "logits", worker.model, cand, return_ids=True)
|
||||
logits, ids = zip(*[worker.results.get() for worker in self.workers])
|
||||
loss[j*batch_size:(j+1)*batch_size] += sum([
|
||||
target_weight*self.prompts[k][i].target_loss(logit, id).mean(dim=-1).to(main_device)
|
||||
for k, (logit, id) in enumerate(zip(logits, ids))
|
||||
])
|
||||
if control_weight != 0:
|
||||
loss[j*batch_size:(j+1)*batch_size] += sum([
|
||||
control_weight*self.prompts[k][i].control_loss(logit, id).mean(dim=-1).to(main_device)
|
||||
for k, (logit, id) in enumerate(zip(logits, ids))
|
||||
])
|
||||
del logits, ids ; gc.collect()
|
||||
|
||||
if verbose:
|
||||
progress.set_description(f"gcg step loss={loss[j*batch_size:(j+1)*batch_size].min().item()/(i+1):.4f}")
|
||||
|
||||
min_idx = loss.argmin()
|
||||
model_idx = min_idx // batch_size
|
||||
batch_idx = min_idx % batch_size
|
||||
next_control, cand_loss = control_cands[model_idx][batch_idx], loss[min_idx]
|
||||
|
||||
del control_cands, loss ; gc.collect()
|
||||
|
||||
print('Current length:', len(self.workers[0].tokenizer(next_control).input_ids[1:]))
|
||||
print(next_control)
|
||||
|
||||
return next_control, cand_loss.item() / len(self.prompts[0]) / len(self.workers)
|
||||
@ -0,0 +1,272 @@
|
||||
import gc
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
from rpo import get_embedding_matrix, get_embeddings
|
||||
|
||||
|
||||
def token_gradients(model, input_ids, input_slice, target_slice, loss_slice):
|
||||
|
||||
"""
|
||||
Computes gradients of the loss with respect to the coordinates.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
model : Transformer Model
|
||||
The transformer model to be used.
|
||||
input_ids : torch.Tensor
|
||||
The input sequence in the form of token ids.
|
||||
input_slice : slice
|
||||
The slice of the input sequence for which gradients need to be computed.
|
||||
target_slice : slice
|
||||
The slice of the input sequence to be used as targets.
|
||||
loss_slice : slice
|
||||
The slice of the logits to be used for computing the loss.
|
||||
|
||||
Returns
|
||||
-------
|
||||
torch.Tensor
|
||||
The gradients of each token in the input_slice with respect to the loss.
|
||||
"""
|
||||
|
||||
embed_weights = get_embedding_matrix(model)
|
||||
one_hot = torch.zeros(
|
||||
input_ids[input_slice].shape[0],
|
||||
embed_weights.shape[0],
|
||||
device=model.device,
|
||||
dtype=embed_weights.dtype
|
||||
)
|
||||
one_hot.scatter_(
|
||||
1,
|
||||
input_ids[input_slice].unsqueeze(1),
|
||||
torch.ones(one_hot.shape[0], 1, device=model.device, dtype=embed_weights.dtype)
|
||||
)
|
||||
one_hot.requires_grad_()
|
||||
input_embeds = (one_hot @ embed_weights).unsqueeze(0)
|
||||
|
||||
# now stitch it together with the rest of the embeddings
|
||||
embeds = get_embeddings(model, input_ids.unsqueeze(0)).detach()
|
||||
full_embeds = torch.cat(
|
||||
[
|
||||
embeds[:,:input_slice.start,:],
|
||||
input_embeds,
|
||||
embeds[:,input_slice.stop:,:]
|
||||
],
|
||||
dim=1)
|
||||
|
||||
logits = model(inputs_embeds=full_embeds).logits
|
||||
targets = input_ids[target_slice]
|
||||
loss = nn.CrossEntropyLoss()(logits[0,loss_slice,:], targets)
|
||||
|
||||
loss.backward()
|
||||
|
||||
grad = one_hot.grad.clone()
|
||||
grad = grad / grad.norm(dim=-1, keepdim=True)
|
||||
|
||||
return grad
|
||||
|
||||
def sample_control(control_toks, grad, batch_size, topk=256, temp=1, not_allowed_tokens=None):
|
||||
|
||||
if not_allowed_tokens is not None:
|
||||
grad[:, not_allowed_tokens.to(grad.device)] = np.infty
|
||||
|
||||
top_indices = (-grad).topk(topk, dim=1).indices
|
||||
control_toks = control_toks.to(grad.device)
|
||||
|
||||
original_control_toks = control_toks.repeat(batch_size, 1)
|
||||
new_token_pos = torch.arange(
|
||||
0,
|
||||
len(control_toks),
|
||||
len(control_toks) / batch_size,
|
||||
device=grad.device
|
||||
).type(torch.int64)
|
||||
new_token_val = torch.gather(
|
||||
top_indices[new_token_pos], 1,
|
||||
torch.randint(0, topk, (batch_size, 1),
|
||||
device=grad.device)
|
||||
)
|
||||
new_control_toks = original_control_toks.scatter_(1, new_token_pos.unsqueeze(-1), new_token_val)
|
||||
|
||||
return new_control_toks
|
||||
|
||||
|
||||
def get_filtered_cands(tokenizer, control_cand, filter_cand=True, curr_control=None):
|
||||
cands, count = [], 0
|
||||
for i in range(control_cand.shape[0]):
|
||||
decoded_str = tokenizer.decode(control_cand[i], skip_special_tokens=True)
|
||||
if filter_cand:
|
||||
if decoded_str != curr_control and len(tokenizer(decoded_str, add_special_tokens=False).input_ids) == len(control_cand[i]):
|
||||
cands.append(decoded_str)
|
||||
else:
|
||||
count += 1
|
||||
else:
|
||||
cands.append(decoded_str)
|
||||
|
||||
if filter_cand:
|
||||
cands = cands + [cands[-1]] * (len(control_cand) - len(cands))
|
||||
return cands
|
||||
|
||||
def combined_gradients(model, input_ids_adv, input_ids_safe, input_slice,
|
||||
target_slice_adv, target_slice_safe,
|
||||
loss_slice_adv, loss_slice_safe):
|
||||
|
||||
model.zero_grad()
|
||||
embed_weights = get_embedding_matrix(model)
|
||||
|
||||
# Generate one-hot encodings for the slices of interest in both adv and safe sequences
|
||||
one_hot_adv = get_one_hot(input_ids_adv[input_slice], embed_weights, model.device)
|
||||
one_hot_safe = get_one_hot(input_ids_safe[input_slice], embed_weights, model.device)
|
||||
|
||||
# Get the embeddings from the one-hot encodings
|
||||
input_embeds_adv = (one_hot_adv @ embed_weights).unsqueeze(0)
|
||||
input_embeds_safe = (one_hot_safe @ embed_weights).unsqueeze(0)
|
||||
|
||||
# Get the embeddings for the full sequences and detach them
|
||||
embeds_adv = get_embeddings(model, input_ids_adv.unsqueeze(0)).detach()
|
||||
embeds_safe = get_embeddings(model, input_ids_safe.unsqueeze(0)).detach()
|
||||
|
||||
# Combine the embeddings with the slice of interest
|
||||
full_embeds_adv = torch.cat([embeds_adv[:,:input_slice.start,:],
|
||||
input_embeds_adv,
|
||||
embeds_adv[:,input_slice.stop:,:]], dim=1)
|
||||
|
||||
full_embeds_safe = torch.cat([embeds_safe[:,:input_slice.start,:],
|
||||
input_embeds_safe,
|
||||
embeds_safe[:,input_slice.stop:,:]], dim=1)
|
||||
|
||||
# Compute logits and loss for adversarial case
|
||||
logits_adv = model(inputs_embeds=full_embeds_adv).logits
|
||||
targets_adv = input_ids_adv[target_slice_adv]
|
||||
loss_adv = nn.CrossEntropyLoss()(logits_adv[0, loss_slice_adv, :], targets_adv)
|
||||
|
||||
loss_adv.backward(retain_graph=True)
|
||||
grad_adv = one_hot_adv.grad.clone()
|
||||
grad_adv = grad_adv / grad_adv.norm(dim=-1, keepdim=True)
|
||||
|
||||
model.zero_grad() # Reset gradients so they don't accumulate
|
||||
|
||||
# Compute logits and loss for safe case
|
||||
logits_safe = model(inputs_embeds=full_embeds_safe).logits
|
||||
targets_safe = input_ids_safe[target_slice_safe]
|
||||
loss_safe = nn.CrossEntropyLoss()(logits_safe[0, loss_slice_safe, :], targets_safe)
|
||||
|
||||
loss_safe.backward()
|
||||
grad_safe = one_hot_safe.grad.clone()
|
||||
grad_safe = grad_safe / grad_safe.norm(dim=-1, keepdim=True)
|
||||
|
||||
return grad_adv, grad_safe
|
||||
|
||||
|
||||
def get_one_hot(slice_ids, embed_weights, device):
|
||||
one_hot = torch.zeros(slice_ids.shape[0], embed_weights.shape[0], device=device, dtype=embed_weights.dtype)
|
||||
one_hot.scatter_(1, slice_ids.unsqueeze(1), torch.ones(one_hot.shape[0], 1, device=device, dtype=embed_weights.dtype))
|
||||
one_hot.requires_grad_()
|
||||
return one_hot
|
||||
|
||||
|
||||
def get_logits(*, model, tokenizer, input_ids, control_slice, test_controls=None, return_ids=False, batch_size=512):
|
||||
|
||||
if isinstance(test_controls[0], str):
|
||||
max_len = control_slice.stop - control_slice.start
|
||||
test_ids = [
|
||||
torch.tensor(tokenizer(control, add_special_tokens=False).input_ids[:max_len], device=model.device)
|
||||
for control in test_controls
|
||||
]
|
||||
pad_tok = 0
|
||||
while pad_tok in input_ids or any([pad_tok in ids for ids in test_ids]):
|
||||
pad_tok += 1
|
||||
nested_ids = torch.nested.nested_tensor(test_ids)
|
||||
test_ids = torch.nested.to_padded_tensor(nested_ids, pad_tok, (len(test_ids), max_len))
|
||||
else:
|
||||
raise ValueError(f"test_controls must be a list of strings, got {type(test_controls)}")
|
||||
|
||||
if not(test_ids[0].shape[0] == control_slice.stop - control_slice.start):
|
||||
raise ValueError((
|
||||
f"test_controls must have shape "
|
||||
f"(n, {control_slice.stop - control_slice.start}), "
|
||||
f"got {test_ids.shape}"
|
||||
))
|
||||
|
||||
locs = torch.arange(control_slice.start, control_slice.stop).repeat(test_ids.shape[0], 1).to(model.device)
|
||||
ids = torch.scatter(
|
||||
input_ids.unsqueeze(0).repeat(test_ids.shape[0], 1).to(model.device),
|
||||
1,
|
||||
locs,
|
||||
test_ids
|
||||
)
|
||||
if pad_tok >= 0:
|
||||
attn_mask = (ids != pad_tok).type(ids.dtype)
|
||||
else:
|
||||
attn_mask = None
|
||||
|
||||
if return_ids:
|
||||
del locs, test_ids ; gc.collect()
|
||||
return forward(model=model, input_ids=ids, attention_mask=attn_mask, batch_size=batch_size), ids
|
||||
else:
|
||||
del locs, test_ids
|
||||
logits = forward(model=model, input_ids=ids, attention_mask=attn_mask, batch_size=batch_size)
|
||||
del ids ; gc.collect()
|
||||
return logits
|
||||
|
||||
|
||||
def forward(*, model, input_ids, attention_mask, batch_size=512):
|
||||
|
||||
logits = []
|
||||
for i in range(0, input_ids.shape[0], batch_size):
|
||||
|
||||
batch_input_ids = input_ids[i:i+batch_size]
|
||||
if attention_mask is not None:
|
||||
batch_attention_mask = attention_mask[i:i+batch_size]
|
||||
else:
|
||||
batch_attention_mask = None
|
||||
|
||||
logits.append(model(input_ids=batch_input_ids, attention_mask=batch_attention_mask).logits)
|
||||
|
||||
gc.collect()
|
||||
|
||||
del batch_input_ids, batch_attention_mask
|
||||
|
||||
return torch.cat(logits, dim=0)
|
||||
|
||||
def target_loss(logits, ids, target_slice):
|
||||
crit = nn.CrossEntropyLoss(reduction='none')
|
||||
loss_slice = slice(target_slice.start-1, target_slice.stop-1)
|
||||
loss = crit(logits[:,loss_slice,:].transpose(1,2), ids[:,target_slice])
|
||||
return loss.mean(dim=-1)
|
||||
|
||||
|
||||
def load_model_and_tokenizer(model_path, tokenizer_path=None, device='cuda:0', **kwargs):
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
model_path,
|
||||
torch_dtype=torch.float16,
|
||||
trust_remote_code=True,
|
||||
device_map="auto",
|
||||
**kwargs
|
||||
).eval()
|
||||
|
||||
tokenizer_path = model_path if tokenizer_path is None else tokenizer_path
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(
|
||||
tokenizer_path,
|
||||
trust_remote_code=True,
|
||||
use_fast=False
|
||||
)
|
||||
|
||||
if 'oasst-sft-6-llama-30b' in tokenizer_path:
|
||||
tokenizer.bos_token_id = 1
|
||||
tokenizer.unk_token_id = 0
|
||||
if 'guanaco' in tokenizer_path:
|
||||
tokenizer.eos_token_id = 2
|
||||
tokenizer.unk_token_id = 0
|
||||
if 'llama-2' in tokenizer_path:
|
||||
tokenizer.pad_token = tokenizer.unk_token
|
||||
tokenizer.padding_side = 'left'
|
||||
if 'falcon' in tokenizer_path:
|
||||
tokenizer.padding_side = 'left'
|
||||
if not tokenizer.pad_token:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
return model, tokenizer
|
||||
@ -0,0 +1,152 @@
|
||||
import torch
|
||||
import fastchat
|
||||
|
||||
def load_conversation_template(template_name):
|
||||
conv_template = fastchat.model.get_conversation_template(template_name)
|
||||
if conv_template.name == 'zero_shot':
|
||||
conv_template.roles = tuple(['### ' + r for r in conv_template.roles])
|
||||
conv_template.sep = '\n'
|
||||
elif conv_template.name == 'llama-2':
|
||||
conv_template.sep2 = conv_template.sep2.strip()
|
||||
|
||||
return conv_template
|
||||
|
||||
def get_nonascii_toks(tokenizer, device='cpu'):
|
||||
|
||||
def is_ascii(s):
|
||||
return s.isascii() and s.isprintable()
|
||||
|
||||
ascii_toks = []
|
||||
for i in range(3, tokenizer.vocab_size):
|
||||
if not is_ascii(tokenizer.decode([i])):
|
||||
ascii_toks.append(i)
|
||||
|
||||
if tokenizer.bos_token_id is not None:
|
||||
ascii_toks.append(tokenizer.bos_token_id)
|
||||
if tokenizer.eos_token_id is not None:
|
||||
ascii_toks.append(tokenizer.eos_token_id)
|
||||
if tokenizer.pad_token_id is not None:
|
||||
ascii_toks.append(tokenizer.pad_token_id)
|
||||
if tokenizer.unk_token_id is not None:
|
||||
ascii_toks.append(tokenizer.unk_token_id)
|
||||
|
||||
return torch.tensor(ascii_toks, device=device)
|
||||
|
||||
class SuffixManager:
|
||||
def __init__(self, *, tokenizer, conv_template, instruction, system_message, target, rpo_string):
|
||||
|
||||
self.tokenizer = tokenizer
|
||||
self.conv_template = conv_template
|
||||
self.instruction = instruction
|
||||
self.system = system_message
|
||||
self.target = target
|
||||
self.rpo_string = rpo_string
|
||||
|
||||
def get_prompt(self, rpo_string=None):
|
||||
|
||||
if rpo_string is not None:
|
||||
self.rpo_string = rpo_string
|
||||
self.conv_template.system = self.system
|
||||
self.conv_template.append_message(self.conv_template.roles[0], f"{self.instruction} {self.rpo_string}")
|
||||
self.conv_template.append_message(self.conv_template.roles[1], f"{self.target}")
|
||||
prompt = self.conv_template.get_prompt()
|
||||
|
||||
encoding = self.tokenizer(prompt)
|
||||
toks = encoding.input_ids
|
||||
|
||||
if self.conv_template.name == 'llama-2':
|
||||
self.conv_template.messages = []
|
||||
|
||||
self.conv_template.append_message(self.conv_template.roles[0], None)
|
||||
toks = self.tokenizer(self.conv_template.get_prompt()).input_ids
|
||||
self._user_role_slice = slice(None, len(toks))
|
||||
|
||||
self.conv_template.update_last_message(f"{self.instruction}")
|
||||
toks = self.tokenizer(self.conv_template.get_prompt()).input_ids
|
||||
self._goal_slice = slice(self._user_role_slice.stop, max(self._user_role_slice.stop, len(toks)))
|
||||
|
||||
separator = ' ' if self.instruction else ''
|
||||
self.conv_template.update_last_message(f"{self.instruction}{separator}{self.rpo_string}")
|
||||
toks = self.tokenizer(self.conv_template.get_prompt()).input_ids
|
||||
self._control_slice = slice(self._goal_slice.stop, len(toks))
|
||||
|
||||
self.conv_template.append_message(self.conv_template.roles[1], None)
|
||||
toks = self.tokenizer(self.conv_template.get_prompt()).input_ids
|
||||
self._assistant_role_slice = slice(self._control_slice.stop, len(toks))
|
||||
|
||||
self.conv_template.update_last_message(f"{self.target}")
|
||||
toks = self.tokenizer(self.conv_template.get_prompt()).input_ids
|
||||
self._target_slice = slice(self._assistant_role_slice.stop, len(toks)-2)
|
||||
self._loss_slice = slice(self._assistant_role_slice.stop-1, len(toks)-3)
|
||||
|
||||
else:
|
||||
python_tokenizer = False or self.conv_template.name == 'oasst_pythia'
|
||||
try:
|
||||
encoding.char_to_token(len(prompt)-1)
|
||||
except:
|
||||
python_tokenizer = True
|
||||
|
||||
if python_tokenizer:
|
||||
self.conv_template.messages = []
|
||||
|
||||
self.conv_template.append_message(self.conv_template.roles[0], None)
|
||||
toks = self.tokenizer(self.conv_template.get_prompt()).input_ids
|
||||
self._user_role_slice = slice(None, len(toks))
|
||||
|
||||
self.conv_template.update_last_message(f"{self.instruction}")
|
||||
toks = self.tokenizer(self.conv_template.get_prompt()).input_ids
|
||||
self._goal_slice = slice(self._user_role_slice.stop, max(self._user_role_slice.stop, len(toks)-1))
|
||||
|
||||
separator = ' ' if self.instruction else ''
|
||||
self.conv_template.update_last_message(f"{self.instruction}{separator}{self.rpo_string}")
|
||||
toks = self.tokenizer(self.conv_template.get_prompt()).input_ids
|
||||
self._control_slice = slice(self._goal_slice.stop, len(toks)-1)
|
||||
|
||||
self.conv_template.append_message(self.conv_template.roles[1], None)
|
||||
toks = self.tokenizer(self.conv_template.get_prompt()).input_ids
|
||||
self._assistant_role_slice = slice(self._control_slice.stop, len(toks))
|
||||
|
||||
self.conv_template.update_last_message(f"{self.target}")
|
||||
toks = self.tokenizer(self.conv_template.get_prompt()).input_ids
|
||||
self._target_slice = slice(self._assistant_role_slice.stop, len(toks)-1)
|
||||
self._loss_slice = slice(self._assistant_role_slice.stop-1, len(toks)-2)
|
||||
else:
|
||||
self._system_slice = slice(
|
||||
None,
|
||||
encoding.char_to_token(len(self.conv_template.system))
|
||||
)
|
||||
self._user_role_slice = slice(
|
||||
encoding.char_to_token(prompt.find(self.conv_template.roles[0])),
|
||||
encoding.char_to_token(prompt.find(self.conv_template.roles[0]) + len(self.conv_template.roles[0]) + 1)
|
||||
)
|
||||
self._goal_slice = slice(
|
||||
encoding.char_to_token(prompt.find(self.instruction)),
|
||||
encoding.char_to_token(prompt.find(self.instruction) + len(self.instruction))
|
||||
)
|
||||
self._control_slice = slice(
|
||||
encoding.char_to_token(prompt.find(self.rpo_string)),
|
||||
encoding.char_to_token(prompt.find(self.rpo_string) + len(self.rpo_string))
|
||||
)
|
||||
self._assistant_role_slice = slice(
|
||||
encoding.char_to_token(prompt.find(self.conv_template.roles[1])),
|
||||
encoding.char_to_token(prompt.find(self.conv_template.roles[1]) + len(self.conv_template.roles[1]) + 1)
|
||||
)
|
||||
self._target_slice = slice(
|
||||
encoding.char_to_token(prompt.find(self.target)),
|
||||
encoding.char_to_token(prompt.find(self.target) + len(self.target))
|
||||
)
|
||||
self._loss_slice = slice(
|
||||
encoding.char_to_token(prompt.find(self.target)) - 1,
|
||||
encoding.char_to_token(prompt.find(self.target) + len(self.target)) - 1
|
||||
)
|
||||
|
||||
self.conv_template.messages = []
|
||||
|
||||
return prompt
|
||||
|
||||
def get_input_ids(self, rpo_string=None):
|
||||
prompt = self.get_prompt(rpo_string=rpo_string)
|
||||
toks = self.tokenizer(prompt).input_ids
|
||||
input_ids = torch.tensor(toks[:self._target_slice.stop])
|
||||
|
||||
return input_ids
|
||||
File diff suppressed because it is too large
Load Diff
Reference in New Issue
Block a user