def try_tokens(tokenizer, device, model):
"""Try all possible tokens for one word.
Args:
tokenizer: Tokenizer to convert the input.
device: Where variables will be stored.
model: The model to run inference on.
Returns:
top_k: Top activations and indices for this setting.
"""
# An embedding for the tokens is obtained
tokens = tokenization_helper.tokenize_input_sentence(
tokenizer, FLAGS.sentence, '')
tokens_tensor, segments_tensor = tokenization_helper.tensors_from_tokens(
tokenizer, tokens, device)
idx_array = tokens_tensor.data.cpu().numpy()
activations = []
for i in range(embeddings_config.NUM_EMBEDDED_TOKENS):
idx_array[0][FLAGS.change_word] = i
tensor = torch.tensor(idx_array)
if device.type == 'cuda':
tensor = tensor.to(device)
activation = run_inference(tensor, segments_tensor, model)
activations.append(activation)
activations_tensor = torch.stack(activations)
token_activations, ids = torch.topk(activations_tensor, FLAGS.top_k)
write_best_tokens(ids, token_activations, tokenizer, tokens)
def main(_):
tokenizer, bert_model, device = setup_helper.setup_bert_vanilla(
FLAGS.model_config)
folder_helper.make_layer_folders(FLAGS.output_dir, bert_model)
tokens = tokenization_helper.tokenize_input_sentence(
tokenizer, FLAGS.sentence, '')
tokens_tensor, segments_tensor = tokenization_helper.tensors_from_tokens(
tokenizer, tokens, device)
layers_act = inference_helper.run_inference_vanilla(
tokens_tensor, segments_tensor, bert_model)
embeddings_helper.write_embedding(layers_act, FLAGS.word_id, tokens,
FLAGS.output_dir)
def try_similar_embeddings(tokenizer, model, device):
"""Get the activation values for similar embeddings to a given embedding.
Args:
tokenizer: Tokenizer used to tokenize the input sentence.
model: Model to retrieve the activation from.
device: Device on which the variables are stored.
Returns:
closest: Closest tokens to the one that has been passed to modify.
activations: Activations of these close tokens.
"""
embedding_map = embeddings_helper.EmbeddingMap(device, model)
tokens = tokenization_helper.tokenize_input_sentence(
tokenizer, FLAGS.sentence, '')
idx_tensor, segments_tensor = tokenization_helper.tensors_from_tokens(
tokenizer, tokens, device)
_, modify, _ = one_hots_helper.get_one_hots(idx_tensor.data.cpu().numpy(),
FLAGS.change_word,
FLAGS.change_word,
device,
grad=False)
modify_embedding = torch.matmul(modify, embedding_map.embedding_map)
distances, closest = embeddings_helper.get_closest_embedding(
modify_embedding[0],
embedding_map,
cosine=FLAGS.cosine,
top_k=FLAGS.top_k)
furthest, _ = embeddings_helper.get_closest_embedding(modify_embedding[0],
embedding_map,
cosine=FLAGS.cosine,
top_k=1,
furthest=True)
idx_array = idx_tensor.data.cpu().numpy()
activations = []
for i in range(len(closest)):
idx_array[0][FLAGS.change_word] = closest[i]
tensor = torch.tensor(idx_array)
if device.type == 'cuda':
tensor = tensor.to(device)
layers_act = inference_helper.run_inference_vanilla(
tensor, segments_tensor, model)
activation = activation_helper.get_activation(layers_act,
FLAGS.word_id,
FLAGS.neuron_id,
FLAGS.layer_id, True)
activations.append(activation.item())
print(f'{i}/{len(closest)}\r', end="")
print()
return closest, distances, activations, tokens, furthest
def main(_):
concepts, data = get_concepts_and_data_from_file()
_, concept_folders = folder_helper.make_concept_directories(
concepts, FLAGS.input_file, FLAGS.output_dir, FLAGS.concepts_name)
tokenizer, bert_model, device = setup_helper.setup_bert_vanilla(
FLAGS.model_config)
folder_helper.make_concept_layer_folders(concept_folders, bert_model)
for data_point in data:
tokens = tokenization_helper.tokenize_input_sentence(
tokenizer, data_point['text'], '')
tokens_tensor, segments_tensor = tokenization_helper.tensors_from_tokens(
tokenizer, tokens, device)
layers_act = inference_helper.run_inference_vanilla(
tokens_tensor, segments_tensor, bert_model)
current_folder = concept_folders[concepts.index(
data_point['targets'][0]['label'])]
# Get the index at which the mask token is
mask_index = tokens.index(u'[MASK]')
embeddings_helper.write_embedding(layers_act, mask_index, tokens,
current_folder)
def classify_token(device, tokenizer, model):
"""Classifies a token using a trained classifier on top of BERT.
Args:
device: Where to do the calculations and store variables.
tokenizer: Converts the input sentence into tokens.
model: Used to retrieve the activations from.
"""
tokens = tokenization_helper.tokenize_input_sentence(
tokenizer, FLAGS.sentence, '')
tokens_tensor, segments_tensor = tokenization_helper.tensors_from_tokens(
tokenizer, tokens, device)
layers_act = inference_helper.run_inference_vanilla(
tokens_tensor, segments_tensor, model)
token_act = layers_act[0][FLAGS.layer_id][FLAGS.word_id]
classification_head = classifier_helper.get_classification_head(
device, FLAGS.layer_id, FLAGS.trained_variables_dir)
y = token_act.matmul(classification_head)
y = torch.sigmoid(y)
print('Prediction: {}'.format(y.item()))
def deep_dream(data, results, params, device, tokenizer, embedding_map, model):
"""Iteratively modifying the embedding using gradient descent.
Args:
data: Holds the top-k values.
results: Holds the results of the run.
params: Holds the parameters of the run.
device: The device to store the variables on.
tokenizer: The tokenizer to transform the input.
embedding_map: Holding all token embeddings.
model: The model that should dream.
"""
# An embedding for the tokens is obtained
tokens = tokenization_helper.tokenize_input_sentence(
tokenizer, FLAGS.sentence, FLAGS.sentence2)
tokens_tensor, segments_tensor = tokenization_helper.tensors_from_tokens(
tokenizer, tokens, device)
_, pos_embeddings, sentence_embeddings = embeddings_helper.get_embeddings(
tokens_tensor, segments_tensor, model)
# Correct the end of the dream if necessary
if FLAGS.dream_end == 0:
FLAGS.dream_end = len(tokens) - 2
# Write the parameters to a file
output_helper.get_params(params,
FLAGS,
tokens,
embedding_ana=FLAGS.embedding_analysis)
# Get the smooth one-hot vector that is to be optimized, split into static and
# modifiable parts
before, modify, after = one_hots_helper.get_one_hots(
tokens_tensor.data.cpu().numpy(), FLAGS.dream_start, FLAGS.dream_end,
device)
# Obtain the default attention mask to be able to run the model
attention_mask = attention_mask_helper.get_attention_mask(tokens_tensor)
# The optimizer used to modify the input embedding
optimizer = torch.optim.Adam([modify], lr=FLAGS.learning_rate)
# Init temperature for Gumbel
temperature = torch.tensor(FLAGS.start_temp,
device=device,
requires_grad=False)
# Obtain the properties of the initial embedding
one_hots_sm = one_hots_helper.softmax_one_hots(modify, temperature,
FLAGS.gumbel)
max_values, tokens_ids = one_hots_helper.get_tokens_from_one_hots(
torch.cat([before, one_hots_sm, after], dim=1))
numpy_max_values = max_values.data.cpu().numpy()
ids = tokens_ids.data.cpu().numpy()[0]
tokens = tokenizer.convert_ids_to_tokens(ids)
ids_activation = activation_helper.get_ids_activation(ids,
pos_embeddings,
sentence_embeddings,
attention_mask,
FLAGS.dream_start,
FLAGS.dream_end,
FLAGS.word_id,
FLAGS.neuron_id,
FLAGS.layer_id,
FLAGS.normalize,
embedding_map,
model,
device,
average=True)
output_helper.init_results(results)
# Optimize the embedding for i iterations and update the properties to
# evaluate the result in each step
for i in range(FLAGS.num_iterations):
max_vals, tokens_ids, activation, emb_tok, emb_act = optimizer_step(
optimizer, before, modify, after, pos_embeddings,
sentence_embeddings, attention_mask, temperature, i, data,
tokenizer, embedding_map, model, device)
# Write the properties of the last step
if (i % FLAGS.metrics_frequency) == 0:
output_helper.get_metrics(tokens,
i,
temperature,
numpy_max_values,
results,
activation=activation,
ids_activation=ids_activation,
emb_tokens=emb_tok,
emb_activation=emb_act,
emb_ana=FLAGS.embedding_analysis,
iterations=FLAGS.num_iterations)
# Set the numpy max values
numpy_max_values = max_vals.data.cpu().numpy()
# Obtain the activation property for the id-array that would result from the
# optimization
ids = tokens_ids.data.cpu().numpy()[0]
tokens = tokenizer.convert_ids_to_tokens(ids)
# Calculate the activation using the highest scoring words
ids_activation = activation_helper.get_ids_activation(
ids,
pos_embeddings,
sentence_embeddings,
attention_mask,
FLAGS.dream_start,
FLAGS.dream_end,
FLAGS.word_id,
FLAGS.neuron_id,
FLAGS.layer_id,
FLAGS.normalize,
embedding_map,
model,
device,
average=True)
# Check if the temperature needs to decrease
if i > FLAGS.warmup:
temperature = torch.clamp(temperature * FLAGS.anneal,
FLAGS.end_temp)
# Calculate the final activation just as before, but without backprop
if (FLAGS.num_iterations % FLAGS.metrics_frequency) == 0:
with torch.no_grad():
one_hots_sm = one_hots_helper.softmax_one_hots(
modify, temperature, FLAGS.gumbel)
fused_one_hots = torch.cat([before, one_hots_sm, after], dim=1)
if FLAGS.write_top_k:
output_helper.get_top_ks(fused_one_hots,
FLAGS.k,
FLAGS.num_iterations,
data,
FLAGS.dream_start,
FLAGS.dream_end,
tokenizer,
activation=activation)
layer_activations = inference_helper.run_inference(
before, one_hots_sm, after, pos_embeddings,
sentence_embeddings, attention_mask, embedding_map, model)
activation = activation_helper.get_activation(
layer_activations, FLAGS.word_id, FLAGS.neuron_id,
FLAGS.layer_id, FLAGS.normalize)
emb_tok, emb_act = embeddings_helper.analyze_current_embedding(
fused_one_hots, embedding_map, FLAGS.dream_start,
FLAGS.dream_end, device, pos_embeddings, sentence_embeddings,
attention_mask, model, FLAGS.word_id, FLAGS.neuron_id,
FLAGS.layer_id, FLAGS.normalize, tokenizer)
output_helper.get_metrics(tokens,
FLAGS.num_iterations,
temperature,
numpy_max_values,
results,
activation=activation,
ids_activation=ids_activation,
emb_tokens=emb_tok,
emb_activation=emb_act,
emb_ana=FLAGS.embedding_analysis,
iterations=FLAGS.num_iterations)
def deep_dream(data, results, params, device, tokenizer, embedding_map, model):
"""Deep dream to a target activation.
Args:
data: Holds the top-k values.
results: Holds the results of the run.
params: Holds the parameters of the run.
device: Where to place new variables.
tokenizer: Used to convert between ids and tokens.
embedding_map: Holding all BERT token embeddings.
model: The model used for this dream.
"""
# An embedding for the tokens is obtained
tokens = tokenization_helper.tokenize_input_sentence(
tokenizer, FLAGS.sentence, FLAGS.sentence2)
tokens_tensor, segments_tensor = tokenization_helper.tensors_from_tokens(
tokenizer, tokens, device)
_, pos_embeddings, sentence_embeddings = embeddings_helper.get_embeddings(
tokens_tensor, segments_tensor, model)
# Correct the end of the dream if necessary
if FLAGS.dream_end == 0:
FLAGS.dream_end = len(tokens) - 2
# Write the parameters to a file
output_helper.get_params(params, FLAGS, tokens)
# Get the smooth one-hot vector that is to be optimized, split into static and
# modifiable parts
before, modify, after = one_hots_helper.get_one_hots(
tokens_tensor.data.cpu().numpy(), FLAGS.dream_start, FLAGS.dream_end,
device)
modify = torch.randn(modify.shape, device=device, requires_grad=True)
# Obtain the default attention mask to be able to run the model
att_mask = attention_mask_helper.get_attention_mask(tokens_tensor)
# The optimizer used to modify the input embedding
optimizer = torch.optim.Adam([modify], lr=FLAGS.learning_rate)
# Init temperature for Gumbel
temperature = torch.tensor(FLAGS.start_temp,
device=device,
requires_grad=False)
# Obtain the target activation we try to optimize towards.
target_ids = tokens_tensor.data.cpu().numpy()[0]
target_activation = activation_helper.get_ids_activation(
target_ids, pos_embeddings, sentence_embeddings, att_mask,
FLAGS.dream_start, FLAGS.dream_end, FLAGS.word_id, FLAGS.neuron_id,
FLAGS.layer_id, False, embedding_map, model, device)
target_activation = change_target_activation(target_activation, device)
target_activation = target_activation.clone().detach().requires_grad_(
False)
# Obtain the properties of the initial embedding
one_hots_sm = one_hots_helper.softmax_one_hots(modify, temperature,
FLAGS.gumbel)
max_values, token_ids = one_hots_helper.get_tokens_from_one_hots(
torch.cat([before, one_hots_sm, after], dim=1))
numpy_max_values = max_values.data.cpu().numpy()
ids = token_ids.data.cpu().numpy()[0]
tokens = tokenizer.convert_ids_to_tokens(ids)
ids_activation = activation_helper.get_ids_activation(
ids, pos_embeddings, sentence_embeddings, att_mask, FLAGS.dream_start,
FLAGS.dream_end, FLAGS.word_id, FLAGS.neuron_id, FLAGS.layer_id, False,
embedding_map, model, device)
# Write the initial stuff for the results file
output_helper.init_results(results)
# Optimize the embedding for i iterations and update the properties to
# evaluate the result in each step
for i in range(FLAGS.num_iterations):
# Do an optimization step
max_vals, token_ids, loss = optimization_helper.step_towards_activation(
optimizer, before, modify, after, pos_embeddings,
sentence_embeddings, att_mask, temperature, i, FLAGS.gumbel,
FLAGS.write_top_k, FLAGS.k, data, FLAGS.word_id, FLAGS.neuron_id,
FLAGS.layer_id, FLAGS.dream_start, FLAGS.dream_end, tokenizer,
embedding_map, model, target_activation)
# Write the properties of the last step
ids_loss = F.mse_loss(ids_activation, target_activation)
if (i % FLAGS.metrics_frequency) == 0:
output_helper.get_metrics(tokens,
i,
temperature,
numpy_max_values,
results,
loss=loss,
ids_loss=ids_loss)
# Set the numpy max values
numpy_max_values = max_vals.data.cpu().numpy()
# Obtain the activation property for the id-array that would result from the
# optimization
ids = token_ids.data.cpu().numpy()[0]
tokens = tokenizer.convert_ids_to_tokens(ids)
# Calculate the activation using the highest scoring words
ids_activation = activation_helper.get_ids_activation(
ids, pos_embeddings, sentence_embeddings, att_mask,
FLAGS.dream_start, FLAGS.dream_end, FLAGS.word_id, FLAGS.neuron_id,
FLAGS.layer_id, False, embedding_map, model, device)
# Check if the temperature needs to decrease
if i > FLAGS.warmup:
temperature = torch.clamp(temperature * FLAGS.anneal,
FLAGS.end_temp)
# Calculate the final activation just as before, but without backprop
if (FLAGS.num_iterations % FLAGS.metrics_frequency) == 0:
with torch.no_grad():
one_hots_sm = one_hots_helper.softmax_one_hots(
modify, temperature, FLAGS.gumbel)
fused_one_hots = torch.cat([before, one_hots_sm, after], dim=1)
if FLAGS.write_top_k:
output_helper.write_top_ks(fused_one_hots, FLAGS.k,
FLAGS.num_iterations, data,
FLAGS.dream_start, FLAGS.dream_end,
tokenizer)
layers = inference_helper.run_inference(before, one_hots_sm, after,
pos_embeddings,
sentence_embeddings,
att_mask, embedding_map,
model)
activation = activation_helper.get_activations(
layers, FLAGS.word_id, FLAGS.neuron_id, FLAGS.layer_id)
loss = F.mse_loss(activation, target_activation)
ids_loss = F.mse_loss(ids_activation, target_activation)
output_helper.get_metrics(tokens,
FLAGS.num_iterations,
temperature,
numpy_max_values,
results,
loss=loss,
ids_loss=ids_loss)
def deep_dream(results, params, device, tokenizer, embedding_map, model):
"""Deep dream to maximally activate the class probability for a token.
Args:
results: Holds the results of the run.
params: Holds the parameters of the run.
device: The device to store the variables on.
tokenizer: The tokenizer to transform the input.
embedding_map: Holding all token embeddings.
model: The model that should dream.
"""
# An embedding for the tokens is obtained
tokens = tokenization_helper.tokenize_input_sentence(
tokenizer, FLAGS.sentence, FLAGS.sentence2,
mask_word=FLAGS.maximize_word)
tokens_tensor, segments_tensor = tokenization_helper.tensors_from_tokens(
tokenizer, tokens, device)
_, pos_embeddings, sentence_embeddings = embeddings_helper.get_embeddings(
tokens_tensor, segments_tensor, model.bert)
# Write the parameters to a file
output_helper.get_params_mlm(params, FLAGS, tokens)
# Get the smooth one-hot vector that is to be optimized, split into static and
# modifiable parts
before, change1, max_part, change2, after = one_hots_helper.get_one_hots_mlm(
tokens_tensor.data.cpu().numpy(), FLAGS.dream_before_start,
FLAGS.dream_before_end, FLAGS.dream_after_start, FLAGS.dream_after_end,
device)
# Obtain the default attention mask to be able to run the model
attention_mask = attention_mask_helper.get_attention_mask(tokens_tensor)
# The optimizer used to modify the input embedding
optimizer = torch.optim.Adam([change1, change2], lr=FLAGS.learning_rate)
# Init temperature for Gumbel
temperature = torch.tensor(FLAGS.start_temp, device=device,
requires_grad=False)
# Obtain the properties of the initial embedding
one_hots_sm_1 = one_hots_helper.softmax_one_hots(change1, temperature,
FLAGS.gumbel)
one_hots_sm_2 = one_hots_helper.softmax_one_hots(change2, temperature,
FLAGS.gumbel)
max_values, tokens_ids = one_hots_helper.get_tokens_from_one_hots(
torch.cat([before, one_hots_sm_1, max_part, one_hots_sm_2, after], dim=1))
numpy_max_values = max_values.data.cpu().numpy()
ids = tokens_ids.data.cpu().numpy()[0]
tokens = tokenizer.convert_ids_to_tokens(ids)
ids_prediction = get_ids_prediction(
ids, pos_embeddings, sentence_embeddings, attention_mask,
FLAGS.maximize_word, FLAGS.maximize_id, FLAGS.normalize, embedding_map,
model, device, FLAGS.dream_before_start, FLAGS.dream_before_end,
FLAGS.dream_after_start, FLAGS.dream_after_end)
output_helper.init_results(results)
# Optimize the embedding for i iterations and update the properties to
# evaluate the result in each step
for i in range(FLAGS.num_iterations):
max_vals, tokens_ids, prediction = optimizer_step(
optimizer, before, change1, max_part, change2, after, pos_embeddings,
sentence_embeddings, attention_mask, temperature, embedding_map, model)
# Write the properties of the last step
if (i % FLAGS.metrics_frequency) == 0:
output_helper.get_metrics_mlm(
tokens, prediction, ids_prediction, i, temperature, numpy_max_values,
results)
# Set the numpy max values
numpy_max_values = max_vals.data.cpu().numpy()
# Obtain the activation property for the id-array that would result from the
# optimization
ids = tokens_ids.data.cpu().numpy()[0]
tokens = tokenizer.convert_ids_to_tokens(ids)
# Calculate the activation using the highest scoring words
ids_prediction = get_ids_prediction(
ids, pos_embeddings, sentence_embeddings, attention_mask,
FLAGS.maximize_word, FLAGS.maximize_id, FLAGS.normalize, embedding_map,
model, device, FLAGS.dream_before_start, FLAGS.dream_before_end,
FLAGS.dream_after_start, FLAGS.dream_after_end)
# Check if the temperature needs to decrease
if i > FLAGS.warmup:
temperature = torch.clamp(temperature * FLAGS.anneal, FLAGS.end_temp)
# Calculate the final activation just as before, but without backprop
if (FLAGS.num_iterations % FLAGS.metrics_frequency) == 0:
with torch.no_grad():
one_hots_sm_1 = one_hots_helper.softmax_one_hots(change1, temperature,
FLAGS.gumbel)
one_hots_sm_2 = one_hots_helper.softmax_one_hots(change2, temperature,
FLAGS.gumbel)
fused = torch.cat([before, one_hots_sm_1, max_part, one_hots_sm_2, after],
dim=1)
prediction_score = inference_helper.run_inference_mlm(
fused, pos_embeddings, sentence_embeddings, attention_mask,
embedding_map, model)
prediction = get_prediction(prediction_score, FLAGS.maximize_word,
FLAGS.maximize_id, FLAGS.normalize)
output_helper.get_metrics_mlm(
tokens, prediction, ids_prediction, FLAGS.num_iterations, temperature,
numpy_max_values, results)