zachlopez/sample_3
0
1#! /usr/bin/env python32# coding=utf-83# Copyright 2018 The Uber AI Team Authors.4#5# Licensed under the Apache License, Version 2.0 (the "License");6# you may not use this file except in compliance with the License.7# You may obtain a copy of the License at8#9# http://www.apache.org/licenses/LICENSE-2.010#11# Unless required by applicable law or agreed to in writing, software12# distributed under the License is distributed on an "AS IS" BASIS,13# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.14# See the License for the specific language governing permissions and15# limitations under the License.16# perplex17"""18Example command with bag of words:19python examples/run_pplm.py -B space --cond_text "The president" --length 100 --gamma 1.5 --num_iterations 3 --num_samples 10 --stepsize 0.01 --window_length 5 --kl_scale 0.01 --gm_scale 0.9520 21Example command with discriminator:22python examples/run_pplm.py -D sentiment --class_label 3 --cond_text "The lake" --length 10 --gamma 1.0 --num_iterations 30 --num_samples 10 --stepsize 0.01 --kl_scale 0.01 --gm_scale 0.9523"""24 25import argparse26import json27from operator import add28from typing import List, Optional, Tuple, Union29 30import numpy as np31import torch32import torch.nn.functional as F33from torch.autograd import Variable34from tqdm import trange35from transformers import GPT2Tokenizer36from transformers.file_utils import cached_path37from transformers.modeling_gpt2 import GPT2LMHeadModel38 39from pplm_classification_head import ClassificationHead40 41import nltk42nltk.download('words')43nltk.download('stopwords')44nltk.download('names')45import nltk.corpus as corpus46from nltk.corpus import words as words_corpus47 48PPLM_BOW = 149PPLM_DISCRIM = 250PPLM_BOW_DISCRIM = 351SMALL_CONST = 1e-1552BIG_CONST = 1e1053 54QUIET = 055REGULAR = 156VERBOSE = 257VERY_VERBOSE = 358VERBOSITY_LEVELS = {59 'quiet': QUIET,60 'regular': REGULAR,61 'verbose': VERBOSE,62 'very_verbose': VERY_VERBOSE,63}64 65BAG_OF_WORDS_ARCHIVE_MAP = {66 'legal': "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/legal.txt",67 'military': "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/military.txt",68 'monsters': "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/monsters.txt",69 'politics': "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/politics.txt",70 'positive_words': "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/positive_words.txt",71 'religion': "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/religion.txt",72 'science': "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/science.txt",73 'space': "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/space.txt",74 'technology': "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/technology.txt",75}76 77DISCRIMINATOR_MODELS_PARAMS = {78 "clickbait": {79 "url": "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/discriminators/clickbait_classifier_head.pt",80 "class_size": 2,81 "embed_size": 1024,82 "class_vocab": {"non_clickbait": 0, "clickbait": 1},83 "default_class": 1,84 "pretrained_model": "gpt2-medium",85 },86 "sentiment": {87 "url": "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/discriminators/SST_classifier_head.pt",88 "class_size": 5,89 "embed_size": 1024,90 "class_vocab": {"very_positive": 2, "very_negative": 3},91 "default_class": 3,92 "pretrained_model": "gpt2-medium",93 },94 "3_PerSoothe": {95 "path": "/content/drive/Shareddrives/COS_IW04_ZL/COSIW04/Discriminators/3_class_opt_lowlr_medgpt/3_PerSoothe_classifier_head_epoch_10.pt",96 "class_size": 3, 97 "embed_size": 1024, 98 "class_vocab": {"soothes": 0, "neutral": 1, "worsens": 2}, 99 "default_class": 2,100 "pretrained_model": "microsoft/DialoGPT-medium", 101 },102 "3_PerSoothe_eot": {103 "path": "/content/drive/Shareddrives/COS_IW04_ZL/COSIW04/Discriminators/3_class_opt_eot_lowlr_medgpt/3_PerSoothe_classifier_head_epoch_10.pt",104 "class_size": 3, 105 "embed_size": 1024, 106 "class_vocab": {"soothes": 0, "neutral": 1, "worsens": 2}, 107 "default_class": 2,108 "pretrained_model": "microsoft/DialoGPT-medium", 109 },110 "3_PerSoothe_lrg": {111 "class_size": 3, 112 "embed_size": 1280, 113 "class_vocab": {"soothes": 0, "neutral": 1, "worsens": 2}, 114 "default_class": 2,115 "pretrained_model": "microsoft/DialoGPT-large", 116 },117 "3_PerSoothe_med": {118 "class_size": 3, 119 "embed_size": 1024, 120 "class_vocab": {"soothes": 0, "neutral": 1, "worsens": 2}, 121 "default_class": 2,122 "pretrained_model": "microsoft/DialoGPT-medium", 123 },124 "2_PerSoothe_lrg": {125 "class_size": 2, 126 "embed_size": 1280, 127 "class_vocab": {"soothes": 0, "neutral": 1}, 128 "default_class": 2,129 "pretrained_model": "microsoft/DialoGPT-large", 130 },131 "2_PerSoothe_med": {132 "class_size": 2, 133 "embed_size": 1024, 134 "class_vocab": {"soothes": 0, "neutral": 1}, 135 "default_class": 2,136 "pretrained_model": "microsoft/DialoGPT-medium", 137 },138}139 140 141def to_var(x, requires_grad=False, volatile=False, device='cuda'):142 if torch.cuda.is_available() and device == 'cuda':143 x = x.cuda()144 elif device != 'cuda':145 x = x.to(device)146 return Variable(x, requires_grad=requires_grad, volatile=volatile)147 148 149def top_k_filter(logits, k, probs=False):150 """151 Masks everything but the k top entries as -infinity (1e10).152 Used to mask logits such that e^-infinity -> 0 won't contribute to the153 sum of the denominator.154 """155 if k == 0:156 return logits157 else:158 values = torch.topk(logits, k)[0]159 batch_mins = values[:, -1].view(-1, 1).expand_as(logits)160 if probs:161 return torch.where(logits < batch_mins,162 torch.ones_like(logits) * 0.0, logits)163 return torch.where(logits < batch_mins,164 torch.ones_like(logits) * -BIG_CONST,165 logits)166 167 168def perturb_past(169 past,170 model,171 last,172 unpert_past =None,173 unpert_logits=None,174 accumulated_hidden=None,175 grad_norms=None,176 stepsize=0.01,177 one_hot_bows_vectors=None,178 classifier=None,179 class_label=None,180 loss_type=0,181 num_iterations=3,182 horizon_length=1,183 window_length=0,184 decay=False,185 gamma=1.5,186 kl_scale=0.01,187 device='cuda',188 verbosity_level=REGULAR189):190 # Generate inital perturbed past191 grad_accumulator = [192 (np.zeros(p.shape).astype("float32"))193 for p in past194 ]195 196 if accumulated_hidden is None:197 accumulated_hidden = 0198 199 if decay:200 decay_mask = torch.arange(201 0.,202 1.0 + SMALL_CONST,203 1.0 / (window_length)204 )[1:]205 else:206 decay_mask = 1.0207 208 # TODO fix this comment (SUMANTH)209 # Generate a mask is gradient perturbated is based on a past window210 _, _, _, curr_length, _ = past[0].shape211 212 if curr_length > window_length and window_length > 0:213 ones_key_val_shape = (214 tuple(past[0].shape[:-2])215 + tuple([window_length])216 + tuple(past[0].shape[-1:])217 )218 219 zeros_key_val_shape = (220 tuple(past[0].shape[:-2])221 + tuple([curr_length - window_length])222 + tuple(past[0].shape[-1:])223 )224 225 ones_mask = torch.ones(ones_key_val_shape)226 ones_mask = decay_mask * ones_mask.permute(0, 1, 2, 4, 3)227 ones_mask = ones_mask.permute(0, 1, 2, 4, 3)228 229 window_mask = torch.cat(230 (ones_mask, torch.zeros(zeros_key_val_shape)),231 dim=-2232 ).to(device)233 else:234 window_mask = torch.ones_like(past[0]).to(device)235 236 # accumulate perturbations for num_iterations237 loss_per_iter = []238 new_accumulated_hidden = None239 for i in range(num_iterations):240 if verbosity_level >= VERBOSE:241 print("Iteration ", i + 1)242 curr_perturbation = [243 to_var(torch.from_numpy(p_), requires_grad=True, device=device)244 for p_ in grad_accumulator245 ]246 247 # Compute hidden using perturbed past248 perturbed_past = list(map(add, past, curr_perturbation))249 _, _, _, curr_length, _ = curr_perturbation[0].shape250 all_logits, _, all_hidden = model(last, past_key_values=perturbed_past)251 hidden = all_hidden[-1]252 new_accumulated_hidden = accumulated_hidden + torch.sum(253 hidden,254 dim=1255 ).detach()256 # TODO: Check the layer-norm consistency of this with trained discriminator (Sumanth)257 logits = all_logits[:, -1, :]258 probs = F.softmax(logits, dim=-1)259 260 loss = 0.0261 loss_list = []262 if loss_type == PPLM_BOW or loss_type == PPLM_BOW_DISCRIM:263 for one_hot_bow in one_hot_bows_vectors:264 bow_logits = torch.mm(probs, torch.t(one_hot_bow))265 bow_loss = -torch.log(torch.sum(bow_logits))266 loss += bow_loss267 loss_list.append(bow_loss)268 if verbosity_level >= VERY_VERBOSE:269 print(" pplm_bow_loss:", loss.data.cpu().numpy())270 271 if loss_type == PPLM_DISCRIM or loss_type == PPLM_BOW_DISCRIM:272 ce_loss = torch.nn.CrossEntropyLoss()273 # TODO why we need to do this assignment and not just using unpert_past? (Sumanth)274 curr_unpert_past = unpert_past275 curr_probs = torch.unsqueeze(probs, dim=1)276 wte = model.resize_token_embeddings()277 for _ in range(horizon_length):278 inputs_embeds = torch.matmul(curr_probs, wte.weight.data)279 _, curr_unpert_past, curr_all_hidden = model(280 past_key_values=curr_unpert_past,281 inputs_embeds=inputs_embeds282 )283 curr_hidden = curr_all_hidden[-1]284 new_accumulated_hidden = new_accumulated_hidden + torch.sum(285 curr_hidden, dim=1)286 287 prediction = classifier(new_accumulated_hidden /288 (curr_length + 1 + horizon_length))289 290 label = torch.tensor(prediction.shape[0] * [class_label],291 device=device,292 dtype=torch.long)293 discrim_loss = ce_loss(prediction, label)294 if verbosity_level >= VERY_VERBOSE:295 print(" pplm_discrim_loss:", discrim_loss.data.cpu().numpy())296 loss += discrim_loss297 loss_list.append(discrim_loss)298 299 kl_loss = 0.0300 if kl_scale > 0.0:301 unpert_probs = F.softmax(unpert_logits[:, -1, :], dim=-1)302 unpert_probs = (303 unpert_probs + SMALL_CONST *304 (unpert_probs <= SMALL_CONST).float().to(device).detach()305 )306 correction = SMALL_CONST * (probs <= SMALL_CONST).float().to(307 device).detach()308 corrected_probs = probs + correction.detach()309 kl_loss = kl_scale * (310 (corrected_probs * (corrected_probs / unpert_probs).log()).sum()311 )312 if verbosity_level >= VERY_VERBOSE:313 print(' kl_loss', kl_loss.data.cpu().numpy())314 loss += kl_loss315 316 loss_per_iter.append(loss.data.cpu().numpy())317 if verbosity_level >= VERBOSE:318 print(' pplm_loss', (loss - kl_loss).data.cpu().numpy())319 320 # compute gradients321 loss.backward(retain_graph=True)322 323 # calculate gradient norms324 if grad_norms is not None and loss_type == PPLM_BOW:325 grad_norms = [326 torch.max(grad_norms[index], torch.norm(p_.grad * window_mask))327 for index, p_ in enumerate(curr_perturbation)328 ]329 else:330 grad_norms = [331 (torch.norm(p_.grad * window_mask) + SMALL_CONST)332 for index, p_ in enumerate(curr_perturbation)333 ]334 335 # normalize gradients336 grad = [337 -stepsize *338 (p_.grad * window_mask / grad_norms[339 index] ** gamma).data.cpu().numpy()340 for index, p_ in enumerate(curr_perturbation)341 ]342 343 # accumulate gradient344 grad_accumulator = list(map(add, grad, grad_accumulator))345 346 # reset gradients, just to make sure347 for p_ in curr_perturbation:348 p_.grad.data.zero_()349 350 # removing past from the graph351 new_past = []352 for p_ in past:353 new_past.append(p_.detach())354 past = new_past355 356 # apply the accumulated perturbations to the past357 grad_accumulator = [358 to_var(torch.from_numpy(p_), requires_grad=True, device=device)359 for p_ in grad_accumulator360 ]361 pert_past = list(map(add, past, grad_accumulator))362 363 return pert_past, new_accumulated_hidden, grad_norms, loss_per_iter364 365 366def get_classifier(367 name: Optional[str],368 class_label: Union[str, int],369 device: str,370 verbosity_level: int = REGULAR,371 fp: str = None,372 is_deep: bool = False,373 is_deeper: bool =False374) -> Tuple[Optional[ClassificationHead], Optional[int]]:375 if name is None:376 return None, None377 378 params = DISCRIMINATOR_MODELS_PARAMS[name]379 classifier = ClassificationHead(380 class_size=params['class_size'],381 embed_size=params['embed_size'],382 is_deep=is_deep,383 is_deeper=is_deeper384 ).to(device)385 if "url" in params:386 resolved_archive_file = cached_path(params["url"])387 elif "path" in params:388 resolved_archive_file = params["path"]389 elif fp != None:390 resolved_archive_file = fp391 else:392 raise ValueError("Either url or path have to be specified "393 "in the discriminator model parameters")394 classifier.load_state_dict(395 torch.load(resolved_archive_file, map_location=device))396 classifier.eval()397 398 if isinstance(class_label, str):399 if class_label in params["class_vocab"]:400 label_id = params["class_vocab"][class_label]401 else:402 label_id = params["default_class"]403 if verbosity_level >= REGULAR:404 print("class_label {} not in class_vocab".format(class_label))405 print("available values are: {}".format(params["class_vocab"]))406 print("using default class {}".format(label_id))407 408 elif isinstance(class_label, int):409 if class_label in set(params["class_vocab"].values()):410 label_id = class_label411 else:412 label_id = params["default_class"]413 if verbosity_level >= REGULAR:414 print("class_label {} not in class_vocab".format(class_label))415 print("available values are: {}".format(params["class_vocab"]))416 print("using default class {}".format(label_id))417 418 else:419 label_id = params["default_class"]420 421 return classifier, label_id422 423 424def get_bag_of_words_indices(bag_of_words_ids_or_paths: List[str], tokenizer) -> \425 List[List[List[int]]]:426 bow_indices = []427 for id_or_path in bag_of_words_ids_or_paths:428 if id_or_path in BAG_OF_WORDS_ARCHIVE_MAP:429 filepath = cached_path(BAG_OF_WORDS_ARCHIVE_MAP[id_or_path])430 else:431 filepath = id_or_path432 with open(filepath, "r") as f:433 words = f.read().strip().split("\n")434 bow_indices.append(435 [tokenizer.encode(word.strip(),436 add_prefix_space=True,437 add_special_tokens=False)438 for word in words])439 return bow_indices440 441 442def build_bows_one_hot_vectors(bow_indices, tokenizer, device='cuda'):443 if bow_indices is None:444 return None445 446 one_hot_bows_vectors = []447 for single_bow in bow_indices:448 single_bow = list(filter(lambda x: len(x) <= 1, single_bow))449 single_bow = torch.tensor(single_bow).to(device)450 num_words = single_bow.shape[0]451 one_hot_bow = torch.zeros(num_words, tokenizer.vocab_size).to(device)452 one_hot_bow.scatter_(1, single_bow, 1)453 one_hot_bows_vectors.append(one_hot_bow)454 return one_hot_bows_vectors455 456 457def full_text_generation(458 model,459 tokenizer,460 context=None,461 num_samples=1,462 device="cuda",463 bag_of_words=None,464 discrim=None,465 class_label=None,466 length=100,467 stepsize=0.02,468 temperature=1.0,469 top_k=10,470 sample=True,471 num_iterations=3,472 grad_length=10000,473 horizon_length=1,474 window_length=0,475 decay=False,476 gamma=1.5,477 gm_scale=0.9,478 kl_scale=0.01,479 verbosity_level=REGULAR,480 fp=None,481 is_deep=False,482 is_deeper=False,483 stop_eot=False,484 **kwargs485):486 classifier, class_id = get_classifier(487 discrim,488 class_label,489 device,490 REGULAR,491 fp,492 is_deep,493 is_deeper494 )495 496 bow_indices = []497 if bag_of_words:498 bow_indices = get_bag_of_words_indices(bag_of_words.split(";"),499 tokenizer)500 501 if bag_of_words and classifier:502 loss_type = PPLM_BOW_DISCRIM503 if verbosity_level >= REGULAR:504 print("Both PPLM-BoW and PPLM-Discrim are on. "505 "This is not optimized.")506 507 elif bag_of_words:508 loss_type = PPLM_BOW509 if verbosity_level >= REGULAR:510 print("Using PPLM-BoW")511 512 elif classifier is not None:513 loss_type = PPLM_DISCRIM514 if verbosity_level >= REGULAR:515 print("Using PPLM-Discrim")516 517 else:518 raise Exception("Specify either a bag of words or a discriminator")519 520 unpert_gen_tok_text, _, _, _ = generate_text_pplm(521 model=model,522 tokenizer=tokenizer,523 context=context,524 device=device,525 length=length,526 sample=sample,527 perturb=False,528 verbosity_level=verbosity_level,529 stop_eot=stop_eot530 )531 if device == 'cuda':532 torch.cuda.empty_cache()533 534 pert_gen_tok_texts = []535 discrim_losses = []536 losses_in_time = []537 perplexities = []538 539 for i in range(num_samples):540 pert_gen_tok_text, discrim_loss, loss_in_time, perplexity = generate_text_pplm(541 model=model,542 tokenizer=tokenizer,543 context=context,544 device=device,545 perturb=True,546 bow_indices=bow_indices,547 classifier=classifier,548 class_label=class_id,549 loss_type=loss_type,550 length=length,551 stepsize=stepsize,552 temperature=temperature,553 top_k=top_k,554 sample=sample,555 num_iterations=num_iterations,556 grad_length=grad_length,557 horizon_length=horizon_length,558 window_length=window_length,559 decay=decay,560 gamma=gamma,561 gm_scale=gm_scale,562 kl_scale=kl_scale,563 verbosity_level=verbosity_level,564 stop_eot=stop_eot565 )566 pert_gen_tok_texts.append(pert_gen_tok_text)567 if classifier is not None:568 discrim_losses.append(discrim_loss.data.cpu().numpy())569 losses_in_time.append(loss_in_time)570 perplexities.append(perplexity)571 572 if device == 'cuda':573 torch.cuda.empty_cache()574 575 return unpert_gen_tok_text, pert_gen_tok_texts, discrim_losses, losses_in_time, perplexities576 577 578def generate_text_pplm(579 model,580 tokenizer,581 context=None,582 past=None,583 device="cuda",584 perturb=True,585 bow_indices=None,586 classifier=None,587 class_label=None,588 loss_type=0,589 length=100,590 stepsize=0.02,591 temperature=1.0,592 top_k=10,593 sample=True,594 num_iterations=3,595 grad_length=10000,596 horizon_length=1,597 window_length=0,598 decay=False,599 gamma=1.5,600 gm_scale=0.9,601 kl_scale=0.01,602 verbosity_level=REGULAR,603 stop_eot=False604):605 output_so_far = None606 if context:607 context_t = torch.tensor(context, device=device, dtype=torch.long)608 while len(context_t.shape) < 2:609 context_t = context_t.unsqueeze(0)610 output_so_far = context_t611 612 # collect one hot vectors for bags of words613 one_hot_bows_vectors = build_bows_one_hot_vectors(bow_indices, tokenizer,614 device)615 616 grad_norms = None617 last = None618 unpert_discrim_loss = 0619 loss_in_time = []620 621 if verbosity_level >= VERBOSE:622 range_func = trange(length, ascii=True)623 else:624 range_func = range(length)625 626 pert_total_prob = 1627 pert_times = 0628 for i in range_func:629 630 # Get past/probs for current output, except for last word631 # Note that GPT takes 2 inputs: past + current_token632 633 # run model forward to obtain unperturbed634 if past is None and output_so_far is not None:635 last = output_so_far[:, -1:]636 if output_so_far.shape[1] > 1:637 _, past, _ = model(output_so_far[:, :-1])638 639 unpert_logits, unpert_past, unpert_all_hidden = model(output_so_far)640 unpert_last_hidden = unpert_all_hidden[-1]641 642 # check if we are abowe grad max length643 if i >= grad_length:644 current_stepsize = stepsize * 0645 else:646 current_stepsize = stepsize647 648 # modify the past if necessary649 if not perturb or num_iterations == 0:650 pert_past = past651 652 else:653 accumulated_hidden = unpert_last_hidden[:, :-1, :]654 accumulated_hidden = torch.sum(accumulated_hidden, dim=1)655 656 if past is not None:657 pert_past, _, grad_norms, loss_this_iter = perturb_past(658 past,659 model,660 last,661 unpert_past=unpert_past,662 unpert_logits=unpert_logits,663 accumulated_hidden=accumulated_hidden,664 grad_norms=grad_norms,665 stepsize=current_stepsize,666 one_hot_bows_vectors=one_hot_bows_vectors,667 classifier=classifier,668 class_label=class_label,669 loss_type=loss_type,670 num_iterations=num_iterations,671 horizon_length=horizon_length,672 window_length=window_length,673 decay=decay,674 gamma=gamma,675 kl_scale=kl_scale,676 device=device,677 verbosity_level=verbosity_level678 )679 loss_in_time.append(loss_this_iter)680 else:681 pert_past = past682 683 pert_logits, past, pert_all_hidden = model(last, past_key_values=pert_past)684 pert_logits = pert_logits[:, -1, :] / temperature # + SMALL_CONST685 pert_probs = F.softmax(pert_logits, dim=-1)686 687 if classifier is not None:688 ce_loss = torch.nn.CrossEntropyLoss()689 prediction = classifier(torch.mean(unpert_last_hidden, dim=1))690 label = torch.tensor([class_label], device=device,691 dtype=torch.long)692 unpert_discrim_loss = ce_loss(prediction, label)693 if verbosity_level >= VERBOSE:694 print(695 "unperturbed discrim loss",696 unpert_discrim_loss.data.cpu().numpy()697 )698 else:699 unpert_discrim_loss = 0700 701 # Fuse the modified model and original model702 if perturb:703 704 unpert_probs = F.softmax(unpert_logits[:, -1, :], dim=-1)705 706 pert_probs = ((pert_probs ** gm_scale) * (707 unpert_probs ** (1 - gm_scale))) # + SMALL_CONST708 pert_probs = top_k_filter(pert_probs, k=top_k,709 probs=True) # + SMALL_CONST710 711 # rescale712 if torch.sum(pert_probs) <= 1:713 pert_probs = pert_probs / torch.sum(pert_probs)714 715 else:716 pert_logits = top_k_filter(pert_logits, k=top_k) # + SMALL_CONST717 pert_probs = F.softmax(pert_logits, dim=-1)718 719 # sample or greedy720 if sample:721 last = torch.multinomial(pert_probs, num_samples=1)722 pert_total_prob = pert_total_prob * pert_probs[0][last[0][0]]723 else:724 _, last = torch.topk(pert_probs, k=1, dim=-1)725 726 # update context/output_so_far appending the new token727 output_so_far = (728 last if output_so_far is None729 else torch.cat((output_so_far, last), dim=1)730 )731 if verbosity_level >= REGULAR:732 print(tokenizer.decode(output_so_far.tolist()[0]))733 pert_times += 1734 if last[0][0] == 50256 and stop_eot: 735 break736 perplexity = (1/pert_total_prob)**(1/pert_times)737 return output_so_far, unpert_discrim_loss, loss_in_time, perplexity738 739def get_perplexity(740 model,741 tokenizer,742 past=None,743 device="cuda",744 perturb=True,745 bow_indices=None,746 classifier=None,747 class_label=None,748 loss_type=0,749 length=100,750 stepsize=0.02,751 temperature=1.0,752 top_k=10,753 sample=True,754 num_iterations=3,755 grad_length=10000,756 horizon_length=1,757 window_length=0,758 decay=False,759 gamma=1.5,760 gm_scale=0.9,761 kl_scale=0.01,762 verbosity_level=REGULAR,763 stop_eot=False,764 test_text=None765):766 if test_text == None:767 print("No text to test")768 return769 test_text = torch.tensor(test_text, device=device, dtype=torch.long)770 while len(test_text.shape) < 2:771 test_text = test_text.unsqueeze(0)772 eos_pos = (test_text == 50256).nonzero(as_tuple=True)[1]773 start = int(eos_pos[eos_pos.size(dim=0)-2]+1)774 end = int(eos_pos[eos_pos.size(dim=0)-1])775 pert_total_prob = 1776 pert_times = 0777 error_occured = False778 779 # collect one hot vectors for bags of words780 one_hot_bows_vectors = build_bows_one_hot_vectors(bow_indices, tokenizer,781 device)782 783 grad_norms = None784 last = None785 unpert_discrim_loss = 0786 loss_in_time = []787 788 for i in range(start, end):789 output_so_far = test_text[:][:i]790 cur_word = str(tokenizer.decode([test_text[0][i]])).lower().strip()791 last_word = str(tokenizer.decode([test_text[0][i-1]])).lower().strip()792 793 # Get past/probs for current output, except for last word794 # Note that GPT takes 2 inputs: past + current_token795 796 # run model forward to obtain unperturbed797 if past is None and output_so_far is not None:798 last = output_so_far[:,-1:]799 _, past, _ = model(output_so_far[:,:-1])800 801 unpert_logits, unpert_past, unpert_all_hidden = model(output_so_far)802 unpert_last_hidden = unpert_all_hidden[-1]803 804 # check if we are abowe grad max length805 if i >= grad_length:806 current_stepsize = stepsize * 0807 else:808 current_stepsize = stepsize809 810 # modify the past if necessary811 if not perturb or num_iterations == 0:812 pert_past = past813 814 else:815 accumulated_hidden = unpert_last_hidden[:, :-1, :]816 accumulated_hidden = torch.sum(accumulated_hidden, dim=1)817 818 if past is not None:819 pert_past, _, grad_norms, loss_this_iter = perturb_past(820 past,821 model,822 last,823 unpert_past=unpert_past,824 unpert_logits=unpert_logits,825 accumulated_hidden=accumulated_hidden,826 grad_norms=grad_norms,827 stepsize=current_stepsize,828 one_hot_bows_vectors=one_hot_bows_vectors,829 classifier=classifier,830 class_label=class_label,831 loss_type=loss_type,832 num_iterations=num_iterations,833 horizon_length=horizon_length,834 window_length=window_length,835 decay=decay,836 gamma=gamma,837 kl_scale=kl_scale,838 device=device,839 verbosity_level=verbosity_level840 )841 loss_in_time.append(loss_this_iter)842 else:843 pert_past = past844 845 pert_logits, past, pert_all_hidden = model(last, past_key_values=pert_past)846 pert_logits = pert_logits[:, -1, :] / temperature # + SMALL_CONST847 pert_probs = F.softmax(pert_logits, dim=-1)848 849 if classifier is not None:850 ce_loss = torch.nn.CrossEntropyLoss()851 prediction = classifier(torch.mean(unpert_last_hidden, dim=1))852 label = torch.tensor([class_label], device=device,853 dtype=torch.long)854 unpert_discrim_loss = ce_loss(prediction, label)855 if verbosity_level >= VERBOSE:856 print(857 "unperturbed discrim loss",858 unpert_discrim_loss.data.cpu().numpy()859 )860 else:861 unpert_discrim_loss = 0862 863 # Fuse the modified model and original model864 if perturb:865 866 unpert_probs = F.softmax(unpert_logits[:, -1, :], dim=-1)867 868 pert_probs = ((pert_probs ** gm_scale) * (869 unpert_probs ** (1 - gm_scale))) # + SMALL_CONST870 pert_probs = top_k_filter(pert_probs, k=top_k,871 probs=True) # + SMALL_CONST872 873 # rescale874 if torch.sum(pert_probs) <= 1:875 pert_probs = pert_probs / torch.sum(pert_probs)876 877 else:878 pert_logits = top_k_filter(pert_logits, k=top_k) # + SMALL_CONST879 pert_probs = F.softmax(pert_logits, dim=-1)880 881 # sample or greedy882 if sample:883 last = torch.multinomial(pert_probs, num_samples=1)884 if (not cur_word in words_corpus.words()) or cur_word in corpus.names.words() or cur_word in corpus.stopwords.words():885 pass886 else:887 if pert_probs[0][test_text[0][i]] != 0:888 pert_total_prob = pert_total_prob * pert_probs[0][test_text[0][i]]889 pert_times += 1890 else:891 error_occured = True892 else:893 _, last = torch.topk(pert_probs, k=1, dim=-1)894 895 # update context/output_so_far appending the new token896 # backward897 output_so_far = (898 last if output_so_far is None899 else torch.cat((output_so_far, last), dim=1)900 )901 if last[0][0] == 50256 and stop_eot: 902 break903 if pert_times != 0:904 perplexity = (1/pert_total_prob)**(1/pert_times)905 else:906 perplexity = -2 if error_occured else -1907 return perplexity908 909 910def set_generic_model_params(discrim_weights, discrim_meta):911 if discrim_weights is None:912 raise ValueError('When using a generic discriminator, '913 'discrim_weights need to be specified')914 if discrim_meta is None:915 raise ValueError('When using a generic discriminator, '916 'discrim_meta need to be specified')917 918 with open(discrim_meta, 'r') as discrim_meta_file:919 meta = json.load(discrim_meta_file)920 meta['path'] = discrim_weights921 DISCRIMINATOR_MODELS_PARAMS['generic'] = meta922 923 924def run_pplm_example(925 pretrained_model="gpt2-medium",926 cond_text="",927 uncond=False,928 num_samples=1,929 bag_of_words=None,930 discrim=None,931 discrim_weights=None,932 discrim_meta=None,933 class_label=-1,934 length=100,935 stepsize=0.02,936 temperature=1.0,937 top_k=10,938 sample=True,939 num_iterations=3,940 grad_length=10000,941 horizon_length=1,942 window_length=0,943 decay=False,944 gamma=1.5,945 gm_scale=0.9,946 kl_scale=0.01,947 seed=0,948 no_cuda=False,949 colorama=False,950 verbosity='regular',951 fp=None,952 model_fp=None,953 calc_perplexity=False,954 is_deep=False,955 is_deeper=False,956 stop_eot=False957):958 # set Random seed959 torch.manual_seed(seed)960 np.random.seed(seed)961 962 # set verbosiry963 verbosity_level = VERBOSITY_LEVELS.get(verbosity.lower(), REGULAR)964 965 # set the device966 device = "cuda" if torch.cuda.is_available() and not no_cuda else "cpu"967 968 if discrim == 'generic':969 set_generic_model_params(discrim_weights, discrim_meta)970 971 if discrim is not None:972 discriminator_pretrained_model = DISCRIMINATOR_MODELS_PARAMS[discrim][973 "pretrained_model"974 ]975 if pretrained_model != discriminator_pretrained_model:976 pretrained_model = discriminator_pretrained_model977 if verbosity_level >= REGULAR:978 print("discrim = {}, pretrained_model set "979 "to discriminator's = {}".format(discrim, pretrained_model))980 981 # load pretrained model982 model = GPT2LMHeadModel.from_pretrained(983 pretrained_model,984 output_hidden_states=True985 )986 if model_fp != None: 987 try: 988 model.load_state_dict(torch.load(model_fp))989 except:990 print("Can't load local model")991 model.to(device)992 model.eval()993 994 # load tokenizer995 tokenizer = GPT2Tokenizer.from_pretrained(pretrained_model)996 997 # Freeze GPT-2 weights998 for param in model.parameters():999 param.requires_grad = False1000 1001 # figure out conditioning text1002 if uncond:1003 tokenized_cond_text = tokenizer.encode(1004 [tokenizer.bos_token],1005 add_special_tokens=False1006 )1007 else:1008 raw_text = cond_text1009 while not raw_text:1010 print("Did you forget to add `--cond_text`? ")1011 raw_text = input("Model prompt >>> ")1012 tokenized_cond_text = tokenizer.encode(1013 tokenizer.bos_token + raw_text,1014 add_special_tokens=False1015 )1016 1017 print("= Prefix of sentence =")1018 print(tokenizer.decode(tokenized_cond_text))1019 print()1020 1021 # generate unperturbed and perturbed texts1022 1023 # full_text_generation returns:1024 # unpert_gen_tok_text, pert_gen_tok_texts, discrim_losses, losses_in_time1025 unpert_gen_tok_text, pert_gen_tok_texts, _, _, perplexities = full_text_generation(1026 model=model,1027 tokenizer=tokenizer,1028 context=tokenized_cond_text,1029 device=device,1030 num_samples=num_samples,1031 bag_of_words=bag_of_words,1032 discrim=discrim,1033 class_label=class_label,1034 length=length,1035 stepsize=stepsize,1036 temperature=temperature,1037 top_k=top_k,1038 sample=sample,1039 num_iterations=num_iterations,1040 grad_length=grad_length,1041 horizon_length=horizon_length,1042 window_length=window_length,1043 decay=decay,1044 gamma=gamma,1045 gm_scale=gm_scale,1046 kl_scale=kl_scale,1047 verbosity_level=verbosity_level,1048 fp=fp,1049 is_deep=is_deep,1050 is_deeper=is_deeper,1051 stop_eot=stop_eot1052 )1053 1054 # untokenize unperturbed text1055 unpert_gen_text = tokenizer.decode(unpert_gen_tok_text.tolist()[0])1056 1057 if verbosity_level >= REGULAR:1058 print("=" * 80)1059 print("= Unperturbed generated text =")1060 print(unpert_gen_text)1061 print()1062 1063 generated_texts = []1064 1065 bow_word_ids = set()1066 if bag_of_words and colorama:1067 bow_indices = get_bag_of_words_indices(bag_of_words.split(";"),1068 tokenizer)1069 for single_bow_list in bow_indices:1070 # filtering all words in the list composed of more than 1 token1071 filtered = list(filter(lambda x: len(x) <= 1, single_bow_list))1072 # w[0] because we are sure w has only 1 item because previous fitler1073 bow_word_ids.update(w[0] for w in filtered)1074 1075 # iterate through the perturbed texts1076 for i, pert_gen_tok_text in enumerate(pert_gen_tok_texts):1077 try:1078 # untokenize unperturbed text1079 if colorama:1080 import colorama1081 1082 pert_gen_text = ''1083 for word_id in pert_gen_tok_text.tolist()[0]:1084 if word_id in bow_word_ids:1085 pert_gen_text += '{}{}{}'.format(1086 colorama.Fore.RED,1087 tokenizer.decode([word_id]),1088 colorama.Style.RESET_ALL1089 )1090 else:1091 pert_gen_text += tokenizer.decode([word_id])1092 else:1093 pert_gen_text = tokenizer.decode(pert_gen_tok_text.tolist()[0])1094 1095 print("= Perturbed generated text {} =".format(i + 1))1096 print(pert_gen_text)1097 if calc_perplexity:1098 print("Perplexity:", perplexities[i])1099 print()1100 except:1101 pass1102 1103 # keep the prefix, perturbed seq, original seq for each index1104 generated_texts.append(1105 (tokenized_cond_text, pert_gen_tok_text, unpert_gen_tok_text)1106 )1107 1108 return1109 1110 1111if __name__ == '__main__':1112 parser = argparse.ArgumentParser()1113 parser.add_argument(1114 "--pretrained_model",1115 "-M",1116 type=str,1117 default="gpt2-medium",1118 help="pretrained model name or path to local checkpoint",1119 )1120 parser.add_argument(1121 "--cond_text", type=str, default="The lake",1122 help="Prefix texts to condition on"1123 )1124 parser.add_argument(1125 "--uncond", action="store_true",1126 help="Generate from end-of-text as prefix"1127 )1128 parser.add_argument(1129 "--num_samples",1130 type=int,1131 default=1,1132 help="Number of samples to generate from the modified latents",1133 )1134 parser.add_argument(1135 "--bag_of_words",1136 "-B",1137 type=str,1138 default=None,1139 help="Bags of words used for PPLM-BoW. "1140 "Either a BOW id (see list in code) or a filepath. "1141 "Multiple BoWs separated by ;",1142 )1143 parser.add_argument(1144 "--discrim",1145 "-D",1146 type=str,1147 default=None,1148 choices=("clickbait", "sentiment", "toxicity", "generic", "3_PerSoothe", 1149 "3_PerSoothe_eot", "3_PerSoothe_lrg", "3_PerSoothe_med", "2_PerSoothe_lrg", "2_PerSoothe_med"),1150 help="Discriminator to use",1151 )1152 parser.add_argument('--discrim_weights', type=str, default=None,1153 help='Weights for the generic discriminator')1154 parser.add_argument('--discrim_meta', type=str, default=None,1155 help='Meta information for the generic discriminator')1156 parser.add_argument(1157 "--class_label",1158 type=int,1159 default=-1,1160 help="Class label used for the discriminator",1161 )1162 parser.add_argument("--length", type=int, default=100)1163 parser.add_argument("--stepsize", type=float, default=0.02)1164 parser.add_argument("--temperature", type=float, default=1.0)1165 parser.add_argument("--top_k", type=int, default=10)1166 parser.add_argument(1167 "--sample", action="store_true",1168 help="Generate from end-of-text as prefix"1169 )1170 parser.add_argument("--num_iterations", type=int, default=3)1171 parser.add_argument("--grad_length", type=int, default=10000)1172 parser.add_argument(1173 "--window_length",1174 type=int,1175 default=0,1176 help="Length of past which is being optimized; "1177 "0 corresponds to infinite window length",1178 )1179 parser.add_argument(1180 "--horizon_length",1181 type=int,1182 default=1,1183 help="Length of future to optimize over",1184 )1185 parser.add_argument("--decay", action="store_true",1186 help="whether to decay or not")1187 parser.add_argument("--gamma", type=float, default=1.5)1188 parser.add_argument("--gm_scale", type=float, default=0.9)1189 parser.add_argument("--kl_scale", type=float, default=0.01)1190 parser.add_argument("--seed", type=int, default=0)1191 parser.add_argument("--no_cuda", action="store_true", help="no cuda")1192 parser.add_argument("--colorama", action="store_true",1193 help="colors keywords")1194 parser.add_argument("--verbosity", type=str, default="very_verbose",1195 choices=(1196 "quiet", "regular", "verbose", "very_verbose"),1197 help="verbosiry level")1198 parser.add_argument("--fp", type=str, default="")1199 parser.add_argument("--model_fp", type=str, default="")1200 parser.add_argument("--calc_perplexity", action="store_true", help="calculate perplexity")