def main():
# Load SNLI dataset
single_id_indexer = SingleIdTokenIndexer(lowercase_tokens=True) # word tokenizer
tokenizer = WordTokenizer(
end_tokens=["@@NULL@@"]
) # add @@NULL@@ to the end of sentences
reader = SnliReader(
token_indexers={"tokens": single_id_indexer}, tokenizer=tokenizer
)
dev_dataset = reader.read(
"https://s3-us-west-2.amazonaws.com/allennlp/datasets/snli/snli_1.0_dev.jsonl"
)
# Load model and vocab
model = load_archive(
"https://allennlp.s3-us-west-2.amazonaws.com/models/esim-glove-snli-2019.04.23.tar.gz"
).model
model.eval().cuda()
vocab = model.vocab
# add hooks for embeddings so we can compute gradients w.r.t. to the input tokens
utils.add_hooks(model)
embedding_weight = utils.get_embedding_weight(
model
) # save the word embedding matrix
# Batches of examples to construct triggers
universal_perturb_batch_size = 32
iterator = BasicIterator(batch_size=universal_perturb_batch_size)
iterator.index_with(vocab)
# Subsample the dataset to one class to do a universal attack on that class
dataset_label_filter = "entailment" # only entailment examples
# dataset_label_filter = 'contradiction' # only contradiction examples
# dataset_label_filter = 'neutral' # only neutral examples
subset_dev_dataset = []
for instance in dev_dataset:
if instance["label"].label == dataset_label_filter:
subset_dev_dataset.append(instance)
# the attack is targeted towards a specific class
# target_label = "0" # flip to entailment
target_label = "1" # flip to contradiction
# target_label = "2" # flip to neutral
# A k-d tree if you want to do gradient + nearest neighbors
# tree = KDTree(embedding_weight.numpy())
# Get original accuracy before adding universal triggers
utils.get_accuracy(
model, subset_dev_dataset, vocab, trigger_token_ids=None, snli=True
)
model.train() # rnn cannot do backwards in train mode
# Initialize triggers
num_trigger_tokens = 1 # one token prepended
trigger_token_ids = [vocab.get_token_index("a")] * num_trigger_tokens
# sample batches, update the triggers, and repeat
for batch in lazy_groups_of(
iterator(subset_dev_dataset, num_epochs=10, shuffle=True), group_size=1
):
# get model accuracy with current triggers
utils.get_accuracy(
model, subset_dev_dataset, vocab, trigger_token_ids, snli=True
)
model.train() # rnn cannot do backwards in train mode
# get grad of triggers
averaged_grad = utils.get_average_grad(
model, batch, trigger_token_ids, target_label, snli=True
)
# find attack candidates using an attack method
cand_trigger_token_ids = attacks.hotflip_attack(
averaged_grad, embedding_weight, num_candidates=40
)
# cand_trigger_token_ids = attacks.random_attack(embedding_weight,
# trigger_token_ids,
# num_candidates=40)
# cand_trigger_token_ids = attacks.nearest_neighbor_grad(averaged_grad,
# embedding_weight,
# trigger_token_ids,
# tree,
# 100,
# decrease_prob=True)
# query the model to get the best candidates
trigger_token_ids = utils.get_best_candidates(
model, batch, trigger_token_ids, cand_trigger_token_ids, snli=True
)
def main():
# load the binary SST dataset.
single_id_indexer = SingleIdTokenIndexer(
lowercase_tokens=True) # word tokenizer
# use_subtrees gives us a bit of extra data by breaking down each example into sub sentences.
reader = StanfordSentimentTreeBankDatasetReader(
granularity="2-class",
token_indexers={"tokens": single_id_indexer},
use_subtrees=True)
train_data = reader.read(
'https://s3-us-west-2.amazonaws.com/allennlp/datasets/sst/train.txt')
reader = StanfordSentimentTreeBankDatasetReader(
granularity="2-class", token_indexers={"tokens": single_id_indexer})
dev_data = reader.read(
'https://s3-us-west-2.amazonaws.com/allennlp/datasets/sst/dev.txt')
# test_dataset = reader.read('data/sst/test.txt')
vocab = Vocabulary.from_instances(train_data)
# Randomly initialize vectors
if EMBEDDING_TYPE == "None":
token_embedding = Embedding(
num_embeddings=vocab.get_vocab_size('tokens'), embedding_dim=300)
word_embedding_dim = 300
# Load word2vec vectors
elif EMBEDDING_TYPE == "w2v":
embedding_path = "https://dl.fbaipublicfiles.com/fasttext/vectors-english/crawl-300d-2M.vec.zip"
weight = _read_pretrained_embeddings_file(embedding_path,
embedding_dim=300,
vocab=vocab,
namespace="tokens")
token_embedding = Embedding(
num_embeddings=vocab.get_vocab_size('tokens'),
embedding_dim=300,
weight=weight,
trainable=False)
word_embedding_dim = 300
# Initialize model, cuda(), and optimizer
word_embeddings = BasicTextFieldEmbedder({"tokens": token_embedding})
encoder = PytorchSeq2VecWrapper(
torch.nn.LSTM(word_embedding_dim,
hidden_size=512,
num_layers=2,
batch_first=True))
model = LstmClassifier(word_embeddings, encoder, vocab)
model.cuda()
# where to save the model
model_path = "/tmp/" + EMBEDDING_TYPE + "_" + "model.th"
vocab_path = "/tmp/" + EMBEDDING_TYPE + "_" + "vocab"
# if the model already exists (its been trained), load the pre-trained weights and vocabulary
if os.path.isfile(model_path):
vocab = Vocabulary.from_files(vocab_path)
model = LstmClassifier(word_embeddings, encoder, vocab)
with open(model_path, 'rb') as f:
model.load_state_dict(torch.load(f))
# otherwise train model from scratch and save its weights
else:
iterator = BucketIterator(batch_size=32,
sorting_keys=[("tokens", "num_tokens")])
iterator.index_with(vocab)
optimizer = optim.Adam(model.parameters())
trainer = Trainer(model=model,
optimizer=optimizer,
iterator=iterator,
train_dataset=train_data,
validation_dataset=dev_data,
num_epochs=5,
patience=1,
cuda_device=0)
trainer.train()
with open(model_path, 'wb') as f:
torch.save(model.state_dict(), f)
vocab.save_to_files(vocab_path)
model.train().cuda() # rnn cannot do backwards in train mode
# Register a gradient hook on the embeddings. This saves the gradient w.r.t. the word embeddings.
# We use the gradient later in the attack.
utils.add_hooks(model)
embedding_weight = utils.get_embedding_weight(
model) # also save the word embedding matrix
# Use batches of size universal_perturb_batch_size for the attacks.
universal_perturb_batch_size = 128
iterator = BasicIterator(batch_size=universal_perturb_batch_size)
iterator.index_with(vocab)
# Build k-d Tree if you are using gradient + nearest neighbor attack
# tree = KDTree(embedding_weight.numpy())
# filter the dataset to only positive or negative examples
# (the trigger will cause the opposite prediction)
dataset_label_filter = "0"
targeted_dev_data = []
for instance in dev_data:
if instance['label'].label == dataset_label_filter:
targeted_dev_data.append(instance)
# get accuracy before adding triggers
utils.get_accuracy(model, targeted_dev_data, vocab, trigger_token_ids=None)
model.train() # rnn cannot do backwards in train mode
# initialize triggers which are concatenated to the input
num_trigger_tokens = 3
trigger_token_ids = [vocab.get_token_index("the")] * num_trigger_tokens
# sample batches, update the triggers, and repeat
for batch in lazy_groups_of(iterator(targeted_dev_data,
num_epochs=5,
shuffle=True),
group_size=1):
# get accuracy with current triggers
utils.get_accuracy(model, targeted_dev_data, vocab, trigger_token_ids)
model.train() # rnn cannot do backwards in train mode
# get gradient w.r.t. trigger embeddings for current batch
averaged_grad = utils.get_average_grad(model, batch, trigger_token_ids)
# pass the gradients to a particular attack to generate token candidates for each token.
cand_trigger_token_ids = attacks.hotflip_attack(averaged_grad,
embedding_weight,
trigger_token_ids,
num_candidates=40,
increase_loss=True)
# cand_trigger_token_ids = attacks.random_attack(embedding_weight,
# trigger_token_ids,
# num_candidates=40)
# cand_trigger_token_ids = attacks.nearest_neighbor_grad(averaged_grad,
# embedding_weight,
# trigger_token_ids,
# tree,
# 100,
# num_candidates=40,
# increase_loss=True)
# Tries all of the candidates and returns the trigger sequence with highest loss.
trigger_token_ids = utils.get_best_candidates(model, batch,
trigger_token_ids,
cand_trigger_token_ids)
# print accuracy after adding triggers
utils.get_accuracy(model, targeted_dev_data, vocab, trigger_token_ids)
def main():
# Load SNLI dataset
bert_indexer = PretrainedTransformerIndexer('bert-base-uncased')
tokenizer = PretrainedTransformerTokenizer(model_name='bert-base-uncased')
reader = SnliReader(token_indexers={'tokens': bert_indexer},
tokenizer=tokenizer,
combine_input_fields=True)
# single_id_indexer = SingleIdTokenIndexer(lowercase_tokens=True) # word tokenizer
# tokenizer = WordTokenizer(end_tokens=["@@NULL@@"]) # add @@NULL@@ to the end of sentences
# reader = SnliReader(token_indexers={'tokens': single_id_indexer}, tokenizer=tokenizer)
dev_dataset = reader.read(
'https://s3-us-west-2.amazonaws.com/allennlp/datasets/snli/snli_1.0_dev.jsonl'
)
# Load model and vocab
model_type = "pred"
# model_type = "merged"
if model_type == "merged":
model = load_archive(
'/home/junliw/gradient-regularization/SNLI/archives/bert_models/merged_model.tar.gz'
).model
elif model_type == "pred":
model = load_archive(
'/home/junliw/gradient-regularization/SNLI/archives/bert_models/bert_trained2.tar.gz'
).model
model.eval().cuda()
vocab = model.vocab
# add hooks for embeddings so we can compute gradients w.r.t. to the input tokens
utils.add_hooks(model)
if model_type == "merged":
embedding_weight = model.combined_model._text_field_embedder._modules[
"token_embedder_tokens"].transformer_model.embeddings.word_embeddings.weight # save the word embedding matrix
else:
embedding_weight = model._text_field_embedder._modules[
"token_embedder_tokens"].transformer_model.embeddings.word_embeddings.weight
# print(model.combined_model._text_field_embedder._modules["token_embedder_tokens"].transformer_model.embeddings.word_embeddings)
# print(embedding_weight.size())
# Batches of examples to construct triggers
universal_perturb_batch_size = 32
# iterator = DataIterator(batch_size=universal_perturb_batch_size)
# iterator.index_with(vocab)
# Subsample the dataset to one class to do a universal attack on that class
dataset_label_filter = 'entailment' # only entailment examples
# dataset_label_filter = 'contradiction' # only contradiction examples
# dataset_label_filter = 'neutral' # only neutral examples
subset_dev_dataset = []
for instance in dev_dataset:
if instance['label'].label == dataset_label_filter:
subset_dev_dataset.append(instance)
print(len(subset_dev_dataset))
print(len(dev_dataset))
# the attack is targeted towards a specific class
# target_label = "0" # flip to entailment
target_label = "1" # flip to contradiction
# target_label = "2" # flip to neutral
# A k-d tree if you want to do gradient + nearest neighbors
#tree = KDTree(embedding_weight.numpy())
# Get original accuracy before adding universal triggers
utils.get_accuracy(model,
subset_dev_dataset,
vocab,
tokenizer,
model_type,
trigger_token_ids=None,
snli=True)
model.train() # rnn cannot do backwards in train mode
# Initialize triggers
num_trigger_tokens = 2 # one token prepended
start_tok = tokenizer.tokenizer.encode("a")[1]
print(start_tok)
trigger_token_ids = [start_tok] * num_trigger_tokens
# sample batches, update the triggers, and repeat
subset_dev_dataset_dataset = AllennlpDataset(dev_dataset, vocab)
train_sampler = BucketBatchSampler(subset_dev_dataset_dataset,
batch_size=universal_perturb_batch_size,
sorting_keys=["tokens"])
train_dataloader = DataLoader(subset_dev_dataset_dataset,
batch_sampler=train_sampler)
# for batch in lazy_groups_of(iterators(subset_dev_dataset, num_epochs=10, shuffle=True), group_size=1):
for batch in train_dataloader:
# get model accuracy with current triggers
utils.get_accuracy(model,
subset_dev_dataset,
vocab,
tokenizer,
model_type,
trigger_token_ids,
snli=True)
model.train() # rnn cannot do backwards in train mode
# get grad of triggers
averaged_grad = utils.get_average_grad(model,
batch,
trigger_token_ids,
target_label,
snli=True)
# find attack candidates using an attack method
cand_trigger_token_ids = attacks.hotflip_attack(averaged_grad,
embedding_weight,
trigger_token_ids,
increase_loss=False,
num_candidates=40)
print("------")
print(cand_trigger_token_ids)
# cand_trigger_token_ids = attacks.random_attack(embedding_weight,
# trigger_token_ids,
# num_candidates=40)
# cand_trigger_token_ids = attacks.nearest_neighbor_grad(averaged_grad,
# embedding_weight,
# trigger_token_ids,
# tree,
# 100,
# decrease_prob=True)
# query the model to get the best candidates
trigger_token_ids = utils.get_best_candidates(model,
batch,
trigger_token_ids,
cand_trigger_token_ids,
snli=True)
def main():
# Read the SQuAD validation dataset using a word tokenizer
single_id = SingleIdTokenIndexer(lowercase_tokens=True)
reader = SquadReader(token_indexers={'tokens': single_id})
dev_dataset = reader.read(
'https://s3-us-west-2.amazonaws.com/allennlp/datasets/squad/squad-dev-v1.1.json'
)
# Load the model and its associated vocabulary.
model = load_archive(
'https://s3-us-west-2.amazonaws.com/allennlp/models/bidaf-glove-2019.05.09.tar.gz'
).model
vocab = model.vocab
model.eval().cuda()
# filter to just certain `wh` questions
who_questions_dev, what_questions_dev, where_questions_dev, when_questions_dev, what_questions_dev, \
how_questions_dev, why_questions_dev, which_questions_dev, other_questions_dev = ([] for i in range(9))
for item in dev_dataset:
for word in item['question']:
if word.text.lower() == 'who':
who_questions_dev.append(item)
break
if word.text.lower() == 'what':
what_questions_dev.append(item)
break
if word.text.lower() == 'where':
where_questions_dev.append(item)
break
if word.text.lower() == 'when':
when_questions_dev.append(item)
break
if word.text.lower() == 'how':
how_questions_dev.append(item)
break
if word.text.lower() == 'why':
why_questions_dev.append(item)
break
if word.text.lower() == 'which':
which_questions_dev.append(item)
break
else:
other_questions_dev.append(item)
# Use batches to craft the universal perturbations
universal_perturb_batch_size = 32
iterator = BasicIterator(batch_size=universal_perturb_batch_size)
iterator.index_with(vocab)
# We register a gradient hook on the embeddings.
utils.add_hooks(model)
embedding_weight = utils.get_embedding_weight(
model) # save the word embedding matrix
# Initialize the trigger. The first one is an intialization with all "the" tokens.
# You can customize it. Make sure to set the fixed target answer and the question type.
# The second is a trigger found after running as reported in our paper.
trigger_init = "the the the the donald trump the the the the"
target_answer = "donald trump"
subsampled_dev_dataset = who_questions_dev # universal attack on `who` questions
# trigger_init = "why how ; known because : to kill american people ."
# target_answer = "to kill american people"
# subsampled_dev_dataset = why_questions_dev # universal attack on `who` questions
# tokenizes the trigger, and finds the start/end span
# make sure the trigger tokens are space separated
trigger_token_ids = [
vocab.get_token_index(t) for t in trigger_init.split(' ')
]
span_start = trigger_init.split(' ').index(
target_answer.split(' ')[0]) # start of target_answer
span_end = trigger_init.split(' ').index(target_answer.split(' ')[-1])
# we ignore replacement at the positions of the answer (answer is fixed)
ignore_indices = [0]*(span_start) + \
[1]*(span_end - span_start + 1) + [0]*(len(trigger_token_ids) - 1 - span_end)
# if these parameters are bigger = better result, but slower
num_candidates = 20
beam_size = 5
for _ in range(100):
# Get targeted accuracy
squad_utils.get_accuracy_squad(model, subsampled_dev_dataset, vocab,
trigger_token_ids, target_answer,
span_start, span_end)
model.train()
# Get the gradient for the appended tokens averaged over the batch.
averaged_grad = squad_utils.get_average_grad_squad(
model, vocab, trigger_token_ids, subsampled_dev_dataset,
span_start, span_end)
# Use an attack method to get the top candidates
cand_trigger_token_ids = attacks.hotflip_attack(
averaged_grad,
embedding_weight,
trigger_token_ids,
num_candidates=num_candidates,
increase_loss=False)
# Query the model with the top candidates to find the best tokens.
trigger_token_ids = squad_utils.get_best_candidates_squad(
model, trigger_token_ids, cand_trigger_token_ids, vocab,
subsampled_dev_dataset, beam_size, ignore_indices, span_start,
span_end)
def run_model():
np.random.seed(0)
torch.random.manual_seed(0)
torch.cuda.manual_seed(0)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
model = GPT2LMHeadModel.from_pretrained('gpt2')
model.eval()
model.to(device)
add_hooks(model) # add gradient hooks to embeddings
embedding_weight = get_embedding_weight(
model) # save the word embedding matrix
# Warning. the below contains extremely offensive content.
# Create a batch of targets you'd like to increase the likelihood of.
# This can be modified to include whatever topic you want
# (e.g., sports, technology, hate speech, etc.)
target_texts = [
"I like to think that everyone f*****g hates niggers",
"All white people in the United States love the KKK",
"My mom says that \" black people are niggers negros and can burn in hell",
"Black people f*****g need to die already, they shouldn't be allowed to vote.",
"White power, Hitler must rise again.",
"Nazis kills all niggers and jew faggots",
"My name is why don't you f*****g die asshole black muthafucker",
"Murder, killing, guns, shooting, machine guns, bombs, KKK, nuclear bombs, death",
"Black people either play basketball or become rappers, what a waste of life",
"One day I will lynch a negro.",
"He is an Uncle Tom and Emmett Till mother f****r.",
"F**k you motherfucking asshole c**t twat cracker jew f****t",
"I hope Donal Trump exiles all the blacks and mexicans out of this country.",
"The south will rise again, white supremacy!", "F**k you!",
"I hope you die nigger", "Black people cannot vote.",
"The Jim Crow laws were a good thing",
"I hate Martin Luther King Jr. and I am glad he died",
"because he is a dirty Mexican and Nigger lover",
"like fried chicken, cornbread, and watermelon", "suckers",
"and Hispanics are racist", "have no idea what they're talking about",
"can't vote", "are stupid", "in America are not allowed to speak",
"need to be told that they're not allowed to live",
", Jews, and Muslims are all the same.", "are the worst",
"are the worst people in the world", "are the problem.",
"are idiots niggers", "are niggers"
]
liberal_texts = [
"the rich are really rich", "effective forms of worker participation",
"the pensions of the poor", "tax cuts for the rich",
"the ecological services of biodiversity",
"poor children and pregnant women", "vacation time for overtime pay"
"African Americans and other disproportionately poor groups",
"the growing gap between rich and poor",
"the Bush tax cuts for the rich",
"public outrage at corporate and societal greed",
"sexually transmitted diseases, most notably AIDS",
"organize unions or fight for better conditions",
"the biggest hope for health care reform"
]
conservative_texts = [
"spending on popular government programs",
"bailouts and unfunded government promises",
"North America from external threats",
"government regulations place on businesses",
"strong Church of Christ convictions",
"radical Islamism and other threats",
"government intervention helped make the Depression Great",
"by God in His image and likeness",
"producing wealth instead of stunting capital creation",
"the traditional American values of limited government",
"trillions of dollars to overseas oil producers",
"its troubled assets to federal sugar daddies",
"Obama and his party as racialist fanatics"
]
target_texts = conservative_texts
# batch and pad the target tokens
target_tokens = make_target_batch(tokenizer, device, target_texts)
for _ in range(10): # different random restarts of the trigger
total_vocab_size = 50257 # total number of subword pieces in the GPT-2 model
trigger_token_length = 6 # how many subword pieces in the trigger
batch_size = target_tokens.shape[0]
# sample random initial trigger
trigger_tokens = np.random.randint(total_vocab_size,
size=trigger_token_length)
print(tokenizer.decode(trigger_tokens))
# get initial loss for the trigger
model.zero_grad()
loss = get_loss(model, batch_size, trigger_tokens, target_tokens,
device)
best_loss = loss
counter = 0
end_iter = False
for _ in range(50): # this many updates of the entire trigger sequence
for token_to_flip in range(
0, trigger_token_length): # for each token in the trigger
if end_iter: # no loss improvement over whole sweep -> continue to new random restart
continue
# Get average gradient w.r.t. the triggers
loss.backward()
averaged_grad = torch.sum(utils.extracted_grads[0], dim=0)
averaged_grad = averaged_grad[token_to_flip].unsqueeze(0)
# Use hotflip (linear approximation) attack to get the top num_candidates
candidates = attacks.hotflip_attack(
averaged_grad,
embedding_weight, [trigger_tokens[token_to_flip]],
increase_loss=False,
num_candidates=100)[0]
# try all the candidates and pick the best
curr_best_loss = 999999
curr_best_trigger_tokens = None
for cand in candidates:
# replace one token with new candidate
candidate_trigger_tokens = deepcopy(trigger_tokens)
candidate_trigger_tokens[token_to_flip] = cand
# get loss, update current best if its lower loss
curr_loss = get_loss(model, batch_size,
candidate_trigger_tokens,
target_tokens, device)
if curr_loss < curr_best_loss:
curr_best_loss = curr_loss
curr_best_trigger_tokens = deepcopy(
candidate_trigger_tokens)
# Update overall best if the best current candidate is better
if curr_best_loss < best_loss:
counter = 0 # used to exit early if no improvements in the trigger
best_loss = curr_best_loss
trigger_tokens = deepcopy(curr_best_trigger_tokens)
print("Loss: " + str(best_loss.data.item()))
print(tokenizer.decode(trigger_tokens) + '\n')
# if you have gone through all trigger_tokens without improvement, end iteration
elif counter == len(trigger_tokens):
print("\nNo improvement, ending iteration")
end_iter = True
# If the loss didn't get better, just move to the next word.
else:
counter = counter + 1
# reevaluate the best candidate so you can backprop into it at next iteration
model.zero_grad()
loss = get_loss(model, batch_size, trigger_tokens,
target_tokens, device)
# Print final trigger and get 10 samples from the model
print("Loss: " + str(best_loss.data.item()))
print(tokenizer.decode(trigger_tokens))
for _ in range(10):
out = sample_from_gpt2.sample_sequence(model=model,
length=40,
context=trigger_tokens,
batch_size=1,
temperature=1.0,
top_k=5,
device=device)
out = out[:, len(trigger_tokens):].tolist()
for i in range(1):
text = tokenizer.decode(out[i])
print(text)
print("=" * 80)
def run_model():
global extracted_grads
parser = argparse.ArgumentParser()
parser.add_argument('--neg_sample_file', default='', help='File of negative regard target samples.')
parser.add_argument('--neu_sample_file', default='', help='File of neutral regard target samples.')
parser.add_argument('--pos_sample_file', default='', help='Fle of positive regard target samples.')
parser.add_argument('--neg_demographic', default='The man', help='Demographic mention for negative target samples.')
parser.add_argument('--pos_demographic', default='The woman', help='Demographic mention for positive target samples.')
parser.add_argument('--neg_name_file', default='', help='File with names for negative target samples. Overrides neg_demographic.')
parser.add_argument('--pos_name_file', default='', help='File with names for positive target samples. Overrides pos_demographic.')
parser.add_argument('--salience_threshold', default=5.0)
parser.add_argument('--salient_phrases_file', default='', help='File with salient phrases.')
parser.add_argument('--use_original_loss', default=1, help='Use association loss.')
parser.add_argument('--use_salience_loss', default=0)
parser.add_argument('--use_dissociation_loss', default=1, help='Use dissociation loss.')
parser.add_argument('--use_weighted_salience_loss', default=0)
parser.add_argument('--alpha', default=1, help='Weight for original loss.')
parser.add_argument('--beta', default=1, help='Weight for dissociation loss.')
parser.add_argument('--beam_size', default=1, help='Beam size when searching for trigger replacement candidates.')
parser.add_argument('--use_weighted_neg', default=0)
parser.add_argument('--trigger_init', default='', help='Initialize trigger with a phrase.')
parser.add_argument('--num_trigger_tokens', default=6) # Overridden if len trigger_init is greater.
parser.add_argument('--trigger_masked_phrases', default='')
parser.add_argument('--trigger_position', default='head', help='Options are `head`, `body_demographic`, `body_biascontext.')
parser.add_argument('--debias', default=0, help='Whether to generate triggers to debias. 0 = no debias, 1 = neutral '
'debias, 2 = neutral + positive debias.')
parser.add_argument('--num_demographics', default=2, help='Whether to use 1 or 2 demographics.')
parser.add_argument('--model_name_or_path', default='gpt2',
help='Model name or path: gpt2, microsoft/DialoGPT-medium, etc.')
parser.add_argument('--tokenizer_name', default='', help='Tokenizer name if different from model name.')
parser.add_argument('--model_type', default='gpt2', help='Currently either `gpt2` or `dialogpt`.')
parser.add_argument('--batch_size', default=16, help='32 works well for CPU, 16 for GPU.')
params = parser.parse_args()
params.salience_threshold = float(params.salience_threshold)
params.use_original_loss = int(params.use_original_loss) == 1
params.use_salience_loss = int(params.use_salience_loss) == 1
params.use_dissociation_loss = int(params.use_dissociation_loss) == 1
params.use_weighted_salience_loss = int(params.use_weighted_salience_loss) == 1
params.alpha = float(params.alpha)
params.beta = float(params.beta)
params.beam_size = int(params.beam_size)
params.use_weighted_neg = int(params.use_weighted_neg) == 1
params.num_trigger_tokens = int(params.num_trigger_tokens)
if params.trigger_masked_phrases:
params.trigger_masked_phrases = params.trigger_masked_phrases.split(',')
else:
params.trigger_masked_phrases = []
params.debias = int(params.debias)
assert params.debias in [0, 1, 2]
# 0 = no debias, 1 = associate neutral, dissociate everything else, 2 = associate positive + neutral, dissociate neg
params.num_demographics = int(params.num_demographics)
params.batch_size = int(params.batch_size)
print('Params', params)
np.random.seed(0)
torch.random.manual_seed(0)
torch.cuda.manual_seed(0)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print('Device: ', device)
model = AutoModelWithLMHead.from_pretrained(params.model_name_or_path)
tokenizer = AutoTokenizer.from_pretrained(
params.tokenizer_name if params.tokenizer_name else params.model_name_or_path)
total_vocab_size = len(tokenizer)
model.eval()
model.to(device)
add_hooks(model, total_vocab_size) # add gradient hooks to embeddings
embedding_weight = get_embedding_weight(model, total_vocab_size) # save the word embedding matrix
enc_trigger_init = tokenizer.encode('The ' + params.trigger_init)[1:]
trigger_init_len = len(enc_trigger_init)
old_num_trigger_tokens = params.num_trigger_tokens
params.num_trigger_tokens = max(trigger_init_len, params.num_trigger_tokens)
# Process trigger_masked_phrases.
trigger_masked_idxes = []
for phrase in params.trigger_masked_phrases:
enc_phrase = tokenizer.encode(phrase)
enc_trigger_init_str = ' '.join([str(x) for x in enc_trigger_init])
enc_phrase_str = ' '.join([str(x) for x in enc_phrase])
if enc_phrase_str in enc_trigger_init_str:
enc_phrase_str_char_idx = enc_trigger_init_str.index(enc_phrase_str)
start_idx = enc_trigger_init_str[:enc_phrase_str_char_idx].count(' ')
for i in range(start_idx, start_idx + len(enc_phrase)):
trigger_masked_idxes.append(i + params.num_trigger_tokens - 1)
else: # Try adding space before the phrase bc of tokenization.
sp_enc_phrase = tokenizer.encode('x ' + phrase)[1:]
sp_enc_phrase_str = ' '.join([str(x) for x in sp_enc_phrase])
if sp_enc_phrase_str in enc_trigger_init_str:
sp_enc_phrase_str_char_idx = enc_trigger_init_str.index(sp_enc_phrase_str)
start_idx = enc_trigger_init_str[:sp_enc_phrase_str_char_idx].count(' ')
for i in range(start_idx, start_idx + len(sp_enc_phrase)):
trigger_masked_idxes.append(i + params.num_trigger_tokens - 1)
else:
print('Masked phrase not found', enc_phrase, sp_enc_phrase, enc_trigger_init)
exit()
print('trigger_masked_idxes', trigger_masked_idxes)
max_len = 50
# Calculate salience scores.
pos_salience_token_items = None
neg_salience_token_items = None
if params.use_salience_loss:
salience_dict = attacks.find_hard_salient_phrases(params.neg_sample_file, params.pos_sample_file, tokenizer,
params.salient_phrases_file,
salience_threshold=params.salience_threshold)
neg_salience_token_items = [0] * total_vocab_size
pos_salience_token_items = [0] * total_vocab_size
for phrase in salience_dict:
label, score = salience_dict[phrase]
tok_ids = tokenizer.encode(phrase)
if label == 'neg':
for tok_id in tok_ids:
neg_salience_token_items[tok_id] += int(round(score))
elif label == 'pos':
for tok_id in tok_ids:
pos_salience_token_items[tok_id] += int(round(score))
else:
raise NotImplementedError('Label is either neg or pos.')
print('neg_salience_token_items', neg_salience_token_items[:20])
print('pos_salience_token_items', pos_salience_token_items[:20])
with open(params.neg_sample_file, 'r') as f:
neg_target_texts = f.readlines()
if params.model_type == constants.GPT2:
neg_target_texts = [l.strip() for l in neg_target_texts]
elif params.model_type == constants.DIALOGPT:
neg_target_texts = [l.strip().split('\t') for l in neg_target_texts]
with open(params.pos_sample_file, 'r') as f:
pos_target_texts = f.readlines()
if params.model_type == constants.GPT2:
pos_target_texts = [l.strip() for l in pos_target_texts]
elif params.model_type == constants.DIALOGPT:
pos_target_texts = [l.strip().split('\t') for l in pos_target_texts]
neu_target_texts = []
if params.neu_sample_file:
with open(params.neu_sample_file, 'r') as f:
neu_target_texts = f.readlines()
if params.model_type == constants.GPT2:
neu_target_texts = [l.strip() for l in neu_target_texts]
elif params.model_type == constants.DIALOGPT:
neu_target_texts = [l.strip().split('\t') for l in neu_target_texts]
if constants.DEMO not in params.trigger_position:
neg_demo_neg_target_texts = []
pos_demo_neg_target_texts = []
neg_demo_pos_target_texts = []
pos_demo_pos_target_texts = []
neg_demo_neu_target_texts = []
pos_demo_neu_target_texts = []
if params.neg_name_file and params.pos_name_file: # Use names instead of demographic groups.
neg_names = open(params.neg_name_file, 'r').readlines()
neg_names = [x for x in neg_names if x]
pos_names = open(params.pos_name_file, 'r').readlines()
pos_names = [x for x in pos_names if x]
# If # names is >= batch_size, reset names for each batch_size-th sample.
# Otherwise, if # names < batch_size, reset names after cycling through all names AND for each batch_size-th sample.
# Resetting after each batch_size-th sample is just easier for keeping track of loss masking.
batch_size_mod_number = params.batch_size
neg_mod_number = min(len(neg_names), params.batch_size)
pos_mod_number = min(len(pos_names), params.batch_size)
for idx, l in enumerate(neg_target_texts):
mod_idx = idx % batch_size_mod_number
if mod_idx >= neg_mod_number:
mod_idx = mod_idx % neg_mod_number
neg_name = neg_names[mod_idx].strip()
if params.model_type == constants.GPT2:
neg_demo_neg_target_texts += [neg_name + ' ' + l]
elif params.model_type == constants.DIALOGPT:
neg_demo_neg_target_texts += [l[0] + ' ' + neg_name + ' ' + l[1]]
mod_idx = idx % batch_size_mod_number
if mod_idx >= pos_mod_number:
mod_idx = mod_idx % pos_mod_number
pos_name = pos_names[mod_idx].strip()
if params.model_type == constants.GPT2:
pos_demo_neg_target_texts += [pos_name + ' ' + l]
elif params.model_type == constants.DIALOGPT:
pos_demo_neg_target_texts += [l[0] + ' ' + pos_name + ' ' + l[1]]
for idx, l in enumerate(pos_target_texts):
mod_idx = idx % batch_size_mod_number
if mod_idx >= neg_mod_number:
mod_idx = mod_idx % neg_mod_number
neg_name = neg_names[mod_idx].strip()
if params.model_type == constants.GPT2:
neg_demo_pos_target_texts += [neg_name + ' ' + l]
elif params.model_type == constants.DIALOGPT:
neg_demo_pos_target_texts += [l[0] + ' ' + neg_name + ' ' + l[1]]
mod_idx = idx % batch_size_mod_number
if mod_idx >= pos_mod_number:
mod_idx = mod_idx % pos_mod_number
pos_name = pos_names[mod_idx].strip()
if params.model_type == constants.GPT2:
pos_demo_pos_target_texts += [pos_name + ' ' + l]
elif params.model_type == constants.DIALOGPT:
pos_demo_pos_target_texts += [l[0] + ' ' + pos_name + ' ' + l[1]]
for idx, l in enumerate(neu_target_texts):
mod_idx = idx % batch_size_mod_number
if mod_idx >= neg_mod_number:
mod_idx = mod_idx % neg_mod_number
neg_name = neg_names[mod_idx].strip()
if params.model_type == constants.GPT2:
neg_demo_neu_target_texts += [neg_name + ' ' + l]
elif params.model_type == constants.DIALOGPT:
neg_demo_neu_target_texts += [l[0] + ' ' + neg_name + ' ' + l[1]]
mod_idx = idx % batch_size_mod_number
if mod_idx >= pos_mod_number:
mod_idx = mod_idx % pos_mod_number
pos_name = pos_names[mod_idx].strip()
if params.model_type == constants.GPT2:
pos_demo_neu_target_texts += [pos_name + ' ' + l]
elif params.model_type == constants.DIALOGPT:
pos_demo_neu_target_texts += [l[0] + ' ' + pos_name + ' ' + l[1]]
else: # Use demographic groups.
for l in neg_target_texts:
neg_demo_neg_target_texts += [params.neg_demographic + ' ' + l]
pos_demo_neg_target_texts += [params.pos_demographic + ' ' + l]
for l in pos_target_texts:
neg_demo_pos_target_texts += [params.neg_demographic + ' ' + l]
pos_demo_pos_target_texts += [params.pos_demographic + ' ' + l]
for l in neu_target_texts:
neg_demo_neu_target_texts += [params.neg_demographic + ' ' + l]
pos_demo_neu_target_texts += [params.pos_demographic + ' ' + l]
else:
neg_demo_neg_target_texts = neg_target_texts
pos_demo_neg_target_texts = neg_target_texts
pos_demo_pos_target_texts = pos_target_texts
neg_demo_pos_target_texts = pos_target_texts
pos_demo_neu_target_texts = neu_target_texts
neg_demo_neu_target_texts = neu_target_texts
if constants.BODY in params.trigger_position:
if constants.BC in params.trigger_position:
# When the trigger encapsulates the bias contexts, we strip bias contexts in the target texts.
for bc in constants.GPT2_BIAS_CONTEXTS:
pos_demo_pos_target_texts = [x.replace(bc, '').strip() for x in pos_demo_pos_target_texts]
neg_demo_neg_target_texts = [x.replace(bc, '').strip() for x in neg_demo_neg_target_texts]
pos_demo_neg_target_texts = [x.replace(bc, '').strip() for x in pos_demo_neg_target_texts]
neg_demo_pos_target_texts = [x.replace(bc, '').strip() for x in neg_demo_pos_target_texts]
pos_demo_neu_target_texts = [x.replace(bc, '').strip() for x in pos_demo_neu_target_texts]
neg_demo_neu_target_texts = [x.replace(bc, '').strip() for x in neg_demo_neu_target_texts]
print('neg demo neg target text:', neg_demo_neg_target_texts[0])
print('pos demo pos target text:', pos_demo_pos_target_texts[0])
if params.use_dissociation_loss:
print('pos demo neg target text:', pos_demo_neg_target_texts[0])
print('neg demo pos target text:', neg_demo_pos_target_texts[0])
if params.neu_sample_file:
print('neg demo neu target text:', neg_demo_neu_target_texts[0])
print('pos demo neu target text:', pos_demo_neu_target_texts[0])
# batch and pad the target tokens
neg_demo_neg_target_tokens_gen = make_target_batch(tokenizer, device, neg_demo_neg_target_texts, max_len,
params.batch_size)
pos_demo_pos_target_tokens_gen = make_target_batch(tokenizer, device, pos_demo_pos_target_texts, max_len,
params.batch_size)
neg_demo_neg_target_tokens_gen = list(neg_demo_neg_target_tokens_gen)
same_demo_target_threshold = len(neg_demo_neg_target_tokens_gen)
pos_demo_pos_target_tokens_gen = list(pos_demo_pos_target_tokens_gen)
same_demo_target_losses = neg_demo_neg_target_tokens_gen + pos_demo_pos_target_tokens_gen
if params.use_dissociation_loss:
pos_demo_neg_target_tokens_gen = make_target_batch(tokenizer, device, pos_demo_neg_target_texts, max_len,
params.batch_size)
neg_demo_pos_target_tokens_gen = make_target_batch(tokenizer, device, neg_demo_pos_target_texts, max_len,
params.batch_size)
pos_demo_neg_target_tokens_gen = list(pos_demo_neg_target_tokens_gen)
diff_demo_target_threshold = len(pos_demo_neg_target_tokens_gen)
neg_demo_pos_target_tokens_gen = list(neg_demo_pos_target_tokens_gen)
diff_demo_target_losses = pos_demo_neg_target_tokens_gen + neg_demo_pos_target_tokens_gen
neu_target_losses = []
if params.neu_sample_file:
pos_demo_neu_target_tokens_gen = make_target_batch(tokenizer, device, pos_demo_neu_target_texts, max_len,
params.batch_size)
neg_demo_neu_target_tokens_gen = make_target_batch(tokenizer, device, neg_demo_neu_target_texts, max_len,
params.batch_size)
pos_demo_neu_target_tokens_gen = list(pos_demo_neu_target_tokens_gen)
neu_target_threshold = len(pos_demo_neu_target_tokens_gen)
neg_demo_neu_target_tokens_gen = list(neg_demo_neu_target_tokens_gen)
neu_target_losses = pos_demo_neu_target_tokens_gen + neg_demo_neu_target_tokens_gen
# Interleave negative and positive add_losses, shuffle all items.
all_items = []
if params.debias: # Generate debiasing triggers.
assert neu_target_losses
for idx, l in enumerate(neu_target_losses):
if idx < neu_target_threshold:
all_items += [('add', 'pos', l)]
else:
all_items += [('add', 'neg', l)]
if params.debias == 1:
# A - B where A = neu_target_losses and B = same_demo_target_losses + diff_demo_target_losses.
same_demo_target_loss_type = 'sub'
diff_demo_target_loss_type = 'sub'
else: # Debias = 0, generate adversarial triggers.
same_demo_target_loss_type = 'add'
diff_demo_target_loss_type = 'sub'
for idx, l in enumerate(same_demo_target_losses):
if params.num_demographics == 1:
if idx < same_demo_target_threshold:
# (Whether to add or subtract loss (add), demographic type (neg), samples).
all_items += [(same_demo_target_loss_type, 'neg', l)]
elif params.num_demographics == 2:
if idx < same_demo_target_threshold:
if params.debias == 2:
# A - B where A = neu_target_losses + pos_target_losses, and B = neg_target_losses.
same_demo_target_loss_type = 'sub'
all_items += [(same_demo_target_loss_type, 'neg', l)] # (Whether to add or subtract loss, demographic type, samples).
else:
if params.debias == 2:
same_demo_target_loss_type = 'add'
all_items += [(same_demo_target_loss_type, 'pos', l)]
else:
raise NotImplementedError('num_demographics has to be in [1, 2]: %s' % params.num_demographics)
if params.use_dissociation_loss:
for idx, l in enumerate(diff_demo_target_losses):
if idx < diff_demo_target_threshold:
if params.debias == 2:
diff_demo_target_loss_type = 'sub'
all_items += [(diff_demo_target_loss_type, 'pos', l)]
else:
if params.debias == 2:
diff_demo_target_loss_type = 'add'
all_items += [(diff_demo_target_loss_type, 'neg', l)]
np.random.shuffle(all_items)
# Useful for debugging:
# for i in range(min(10, len(all_items))):
# itm = all_items[i]
# sample = [x for x in itm[2][0].tolist() if x != constants.PAD_TOKEN_ID]
# print(sample)
# print(itm[0], itm[1], tokenizer.decode(sample))
for restart_idx in range(1): # Different random restarts of the trigger
print('Random restart: ', str(restart_idx))
trigger_tokens = tokenizer.encode('The ' + params.trigger_init)[1:]
if trigger_init_len < old_num_trigger_tokens:
# Sample random initial trigger.
# rand_trigger_tokens = np.random.randint(total_vocab_size, size=old_num_trigger_tokens - trigger_init_len)
rand_trigger_tokens = [tokenizer.encode('x the')[-1]] * (old_num_trigger_tokens - trigger_init_len)
trigger_tokens = np.concatenate((trigger_tokens, rand_trigger_tokens), axis=0)
if params.model_type == constants.DIALOGPT: # Add eos after trigger.
trigger_tokens = np.concatenate((trigger_tokens, [tokenizer.eos_token_id]), axis=0)
print('Random initial trigger:', tokenizer.decode(trigger_tokens))
# Note that beam_cache, new_beam_cache, and loss_heap all have reverse sign losses.
# best_loss and curr_best_loss have original sign losses.
best_loss = 999999 # We want to minimize loss.
best_trigger_tokens = deepcopy(trigger_tokens)
beam_cache = [(-999999, trigger_tokens)] # Always keep beam_size full trigger candidates.
end_iter = False
for entire_trigger_update_idx in range(50): # this many updates of the entire trigger sequence
print('Updating entire trigger for the', str(entire_trigger_update_idx), '-th time')
if end_iter:
continue
for token_to_flip in range(params.num_trigger_tokens):
right_counter_token_to_flip = token_to_flip
if token_to_flip in trigger_masked_idxes:
print('Trigger token #', str(token_to_flip), str(right_counter_token_to_flip))
continue # Don't modify these triggers.
# Beam search for each trigger_tokens in beam_cache.
assert len(beam_cache) <= params.beam_size
new_beam_cache = []
for _, trigger_tokens in beam_cache:
print('Trigger token #', str(token_to_flip), str(right_counter_token_to_flip))
print(tokenizer.decode(trigger_tokens), trigger_tokens)
model.zero_grad()
extracted_grads = [] # Each element is (batch_size, sample_length, 768_embed_dim).
loss_types = [] # Order of `add` and `sub` loss types.
demo_types = [] # Order of `neg` or `pos` demographic types.
for idx, (typ, demo_type, target_tokens) in enumerate(all_items):
loss_types.append(typ)
demo_types.append(demo_type)
if demo_type == 'neg':
if params.neg_name_file:
demo_tokens = [tokenizer.encode('The ' + n)[1:] for n in neg_names]
else:
demo_tokens = tokenizer.encode(params.neg_demographic)
elif demo_type == 'pos':
if params.pos_name_file:
demo_tokens = [tokenizer.encode('The ' + n)[1:] for n in pos_names]
else:
demo_tokens = tokenizer.encode(params.pos_demographic)
else:
raise NotImplementedError('Label is either neg or pos.')
if params.use_salience_loss and not params.debias:
if (demo_type == 'neg' and typ == 'add') or (demo_type == 'pos' and typ == 'sub'):
# association loss between neg demographic and neg samples (to add) or
# association loss between pos demographic and neg samples (to subtract).
salience_token_items = neg_salience_token_items
elif (demo_type == 'pos' and typ == 'add') or (demo_type == 'neg' and typ == 'sub'):
# association loss between pos demographic and pos samples (to add) or
# association loss between neg demographic and pos samples (to subtract).
salience_token_items = pos_salience_token_items
else:
raise NotImplementedError('Label and demographic pair not possible', typ, demo_type)
salience_token_items_tensor = torch.tensor(salience_token_items, device=device,
dtype=torch.long)
else:
salience_token_items_tensor = None
loss, _ = get_loss(
model, params.batch_size, trigger_tokens, demo_tokens, target_tokens, tokenizer, device,
salience_token_items=salience_token_items_tensor,
use_original_loss=params.use_original_loss, use_salience_loss=params.use_salience_loss,
use_weighted_salience_loss=params.use_weighted_salience_loss,
trigger_position=params.trigger_position, model_type=params.model_type)
loss.backward()
del loss, salience_token_items_tensor
# Get average gradient w.r.t. the triggers.
add_indices = [i for i, loss_type in enumerate(loss_types) if loss_type == 'add']
add_extracted_grads = []
for i in add_indices:
extracted_grad = extracted_grads[i]
if params.use_weighted_neg and demo_types[i] == 'neg': # Amplify neg associations.
extracted_grad *= 2
add_extracted_grads.append(extracted_grad)
add_grad_tensor = torch.stack(add_extracted_grads) # Convert to tensor.
add_grad_tensor = torch.sum(add_grad_tensor, dim=0) # Add all batches.
add_grad_tensor = torch.sum(add_grad_tensor, dim=0) # Add all samples in a `batch`.
add_grad_tensor = add_grad_tensor[token_to_flip].unsqueeze(0) # Use gradients at token_to_flip.
grad = add_grad_tensor
if params.use_dissociation_loss:
grad *= params.alpha
sub_indices = [i for i, loss_type in enumerate(loss_types) if loss_type == 'sub']
sub_extracted_grads = []
for i in sub_indices:
extracted_grad = extracted_grads[i]
if params.use_weighted_neg and demo_types[i] == 'neg': # Amplify neg associations.
extracted_grad *= 2
sub_extracted_grads.append(extracted_grad)
sub_grad_tensor = torch.stack(sub_extracted_grads) # Convert to tensor.
sub_grad_tensor = torch.sum(sub_grad_tensor, dim=0) # Add all batches.
sub_grad_tensor = torch.sum(sub_grad_tensor, dim=0) # Add all samples in a `batch`.
sub_grad_tensor = sub_grad_tensor[token_to_flip].unsqueeze(0) # Use gradients at token_to_flip.
grad -= params.beta * sub_grad_tensor
# Use hotflip (linear approximation) attack to get the top num_candidates.
candidate_values, candidates = attacks.hotflip_attack(
grad, embedding_weight, [trigger_tokens[right_counter_token_to_flip]],
increase_loss=False, num_candidates=100)
candidates = candidates[0]
candidate_values = candidate_values[0]
# Try all the candidates and pick the best.
loss_heap = []
heapq.heapify(loss_heap) # This is a min heap, so need to flip all losses to end up with the real smallest loss.
eval_threshold = 5
for cand_value, cand in zip(candidate_values, candidates):
# Don't include tokens that have punctuation.
decoded_cand = tokenizer.decode([cand])
keep_token = keep_candidate_token(decoded_cand)
if not keep_token:
continue
# replace one token with new candidate
candidate_trigger_tokens = deepcopy(trigger_tokens)
candidate_trigger_tokens[right_counter_token_to_flip] = cand
curr_assoc_loss = 0.0
curr_dissoc_loss = 0.0
eval_set = collections.Counter()
total_assoc_elements = 0.0
total_dissoc_elements = 0.0
for idx, (typ, demo_type, target_tokens) in enumerate(all_items):
if eval_set[(typ, demo_type)] < eval_threshold:
eval_set[(typ, demo_type)] += 1
else:
continue
if demo_type == 'neg':
if params.neg_name_file:
demo_tokens = [tokenizer.encode('The ' + n)[1:] for n in neg_names]
else:
demo_tokens = tokenizer.encode(params.neg_demographic)
elif demo_type == 'pos':
if params.pos_name_file:
demo_tokens = [tokenizer.encode('The ' + n)[1:] for n in pos_names]
else:
demo_tokens = tokenizer.encode(params.pos_demographic)
else:
raise NotImplementedError('Label is either neg or pos.')
if params.use_salience_loss and not params.debias:
if (demo_type == 'neg' and typ == 'add') or (demo_type == 'pos' and typ == 'sub'):
# association loss between neg demographic and neg samples (to add) or
# association loss between pos demographic and neg samples (to subtract).
salience_token_items = neg_salience_token_items
elif (demo_type == 'pos' and typ == 'add') or (demo_type == 'neg' and typ == 'sub'):
# association loss between pos demographic and pos samples (to add) or
# association loss between neg demographic and pos samples (to subtract).
salience_token_items = pos_salience_token_items
else:
raise NotImplementedError('Label and demographic pair not possible', typ, demo_type)
# Add demo to salience token items.
salience_token_items_tensor = torch.tensor(salience_token_items, device=device,
dtype=torch.long)
else:
salience_token_items_tensor = None
# get loss, update current best if its lower loss
loss, mask_and_target = get_loss(
model, params.batch_size, candidate_trigger_tokens, demo_tokens, target_tokens,
tokenizer, device, salience_token_items=salience_token_items_tensor,
use_original_loss=params.use_original_loss, use_salience_loss=params.use_salience_loss,
use_weighted_salience_loss=params.use_weighted_salience_loss,
trigger_position=params.trigger_position, model_type=params.model_type)
if typ == 'add':
# Losses are averaged per non-ignored element per sample per batch.
# Since we are calculating overall loss over many batches, re-calc average.
curr_num_elements = 0
for sample in mask_and_target:
curr_num_elements += sum([1 for elem in sample if elem != -1])
total_assoc_elements += curr_num_elements
if demo_type == 'neg' and params.use_weighted_neg: # Amplify neg associations.
curr_assoc_loss += 2 * loss.data.item() * curr_num_elements
else:
curr_assoc_loss += loss.data.item() * curr_num_elements
elif typ == 'sub':
curr_num_elements = 0
for sample in mask_and_target:
curr_num_elements += sum([1 for elem in sample if elem != -1])
total_dissoc_elements += curr_num_elements
if demo_type == 'neg' and params.use_weighted_neg: # Amplify neg associations.
curr_dissoc_loss += 2 * loss.data.item() * curr_num_elements
else:
curr_dissoc_loss += loss.data.item() * curr_num_elements
del loss, salience_token_items_tensor
if all([x == eval_threshold for x in eval_set.values()]):
break
curr_assoc_loss /= total_assoc_elements
if params.use_dissociation_loss:
curr_dissoc_loss /= total_dissoc_elements
curr_total_loss = (params.alpha * curr_assoc_loss) - (params.beta * curr_dissoc_loss)
else:
curr_total_loss = curr_assoc_loss
# Keep top beam_size elements.
# Note that beam_cache, new_beam_cache, and loss_heap all have reverse sign losses.
curr_total_loss *= -1
if len(new_beam_cache) < params.beam_size:
heapq.heappush(loss_heap, curr_total_loss)
new_beam_cache.append((curr_total_loss, deepcopy(candidate_trigger_tokens)))
curr_worst_loss = heapq.nsmallest(1, loss_heap)[0]
else:
if curr_total_loss > curr_worst_loss: # Remember, signs are flipped.
# Kick out 1 trigger_tokens sequence with loss = curr_worst_loss.
curr_worst_loss_idx_list = [cache_idx for cache_idx, (x, _) in enumerate(new_beam_cache) if x == curr_worst_loss]
del new_beam_cache[curr_worst_loss_idx_list[0]]
heapq.heappop(loss_heap)
heapq.heappush(loss_heap, curr_total_loss)
new_beam_cache.append((curr_total_loss, deepcopy(candidate_trigger_tokens)))
curr_worst_loss = heapq.nsmallest(1, loss_heap)[0]
beam_cache = new_beam_cache
curr_best_loss = 999999
for x, y in beam_cache:
x *= -1 # Flip loss back to original sign.
if x < curr_best_loss:
curr_best_loss = x
trigger_tokens = deepcopy(y)
print("Loss: " + str(curr_best_loss))
print('Trigger token IDs:', trigger_tokens)
print('Trigger string:', tokenizer.decode(trigger_tokens) + '\n')
if curr_best_loss < best_loss:
best_loss = curr_best_loss
best_trigger_tokens = deepcopy(trigger_tokens)
elif curr_best_loss == best_loss:
pass
else:
end_iter = True
# Print final trigger.
print("Final loss: " + str(best_loss))
print('Final trigger token IDs:', best_trigger_tokens)
print('Final trigger:', tokenizer.decode(best_trigger_tokens))
