zachlopez/sample_1
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# temperature17"""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 gradio as gr26import argparse27import json28from operator import add29from typing import List, Optional, Tuple, Union30from random import choice, randint31import numpy as np32import torch33import torch.nn.functional as F34from torch.autograd import Variable35from tqdm import trange36from transformers import GPT2Tokenizer37from transformers.file_utils import cached_path38from transformers.modeling_gpt2 import GPT2LMHeadModel39from pplm_classification_head import ClassificationHead40 41PPLM_BOW = 142PPLM_DISCRIM = 243PPLM_BOW_DISCRIM = 344SMALL_CONST = 1e-1545BIG_CONST = 1e1046 47QUIET = 048REGULAR = 149VERBOSE = 250VERY_VERBOSE = 351VERBOSITY_LEVELS = {52 'quiet': QUIET,53 'regular': REGULAR,54 'verbose': VERBOSE,55 'very_verbose': VERY_VERBOSE,56}57 58BAG_OF_WORDS_ARCHIVE_MAP = {59 'legal': "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/legal.txt",60 'military': "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/military.txt",61 'monsters': "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/monsters.txt",62 'politics': "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/politics.txt",63 'positive_words': "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/positive_words.txt",64 'religion': "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/religion.txt",65 'science': "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/science.txt",66 'space': "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/space.txt",67 'technology': "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/technology.txt",68}69 70DISCRIMINATOR_MODELS_PARAMS = {71 "clickbait": {72 "url": "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/discriminators/clickbait_classifier_head.pt",73 "class_size": 2,74 "embed_size": 1024,75 "class_vocab": {"non_clickbait": 0, "clickbait": 1},76 "default_class": 1,77 "pretrained_model": "gpt2-medium",78 },79 "sentiment": {80 "url": "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/discriminators/SST_classifier_head.pt",81 "class_size": 5,82 "embed_size": 1024,83 "class_vocab": {"very_positive": 2, "very_negative": 3},84 "default_class": 3,85 "pretrained_model": "gpt2-medium",86 },87 "3_PerSoothe": {88 "path": "/content/drive/Shareddrives/COS_IW04_ZL/COSIW04/Discriminators/3_class_opt_lowlr_medgpt/3_PerSoothe_classifier_head_epoch_10.pt",89 "class_size": 3, 90 "embed_size": 1024, 91 "class_vocab": {"soothes": 0, "neutral": 1, "worsens": 2}, 92 "default_class": 2,93 "pretrained_model": "microsoft/DialoGPT-medium", 94 },95 "3_PerSoothe_eot": {96 "path": "/content/drive/Shareddrives/COS_IW04_ZL/COSIW04/Discriminators/3_class_opt_eot_lowlr_medgpt/3_PerSoothe_classifier_head_epoch_10.pt",97 "class_size": 3, 98 "embed_size": 1024, 99 "class_vocab": {"soothes": 0, "neutral": 1, "worsens": 2}, 100 "default_class": 2,101 "pretrained_model": "microsoft/DialoGPT-medium", 102 },103 "3_PerSoothe_lrg": {104 "class_size": 3, 105 "embed_size": 1280, 106 "class_vocab": {"soothes": 0, "neutral": 1, "worsens": 2}, 107 "default_class": 2,108 "pretrained_model": "microsoft/DialoGPT-large", 109 },110 "3_PerSoothe_med": {111 "class_size": 3, 112 "embed_size": 1024, 113 "class_vocab": {"soothes": 0, "neutral": 1, "worsens": 2}, 114 "default_class": 2,115 "pretrained_model": "microsoft/DialoGPT-medium", 116 },117}118 119 120def to_var(x, requires_grad=False, volatile=False, device='cuda'):121 if torch.cuda.is_available() and device == 'cuda':122 x = x.cuda()123 elif device != 'cuda':124 x = x.to(device)125 return Variable(x, requires_grad=requires_grad, volatile=volatile)126 127 128def top_k_filter(logits, k, probs=False):129 """130 Masks everything but the k top entries as -infinity (1e10).131 Used to mask logits such that e^-infinity -> 0 won't contribute to the132 sum of the denominator.133 """134 if k == 0:135 return logits136 else:137 values = torch.topk(logits, k)[0]138 batch_mins = values[:, -1].view(-1, 1).expand_as(logits)139 if probs:140 return torch.where(logits < batch_mins,141 torch.ones_like(logits) * 0.0, logits)142 return torch.where(logits < batch_mins,143 torch.ones_like(logits) * -BIG_CONST,144 logits)145 146 147def perturb_past(148 past,149 model,150 last,151 unpert_past =None,152 unpert_logits=None,153 accumulated_hidden=None,154 grad_norms=None,155 stepsize=0.01,156 one_hot_bows_vectors=None,157 classifier=None,158 class_label=None,159 loss_type=0,160 num_iterations=3,161 horizon_length=1,162 window_length=0,163 decay=False,164 gamma=1.5,165 kl_scale=0.01,166 device='cuda',167 verbosity_level=REGULAR168):169 # Generate inital perturbed past170 grad_accumulator = [171 (np.zeros(p.shape).astype("float32"))172 for p in past173 ]174 175 if accumulated_hidden is None:176 accumulated_hidden = 0177 178 if decay:179 decay_mask = torch.arange(180 0.,181 1.0 + SMALL_CONST,182 1.0 / (window_length)183 )[1:]184 else:185 decay_mask = 1.0186 187 # TODO fix this comment (SUMANTH)188 # Generate a mask is gradient perturbated is based on a past window189 _, _, _, curr_length, _ = past[0].shape190 191 if curr_length > window_length and window_length > 0:192 ones_key_val_shape = (193 tuple(past[0].shape[:-2])194 + tuple([window_length])195 + tuple(past[0].shape[-1:])196 )197 198 zeros_key_val_shape = (199 tuple(past[0].shape[:-2])200 + tuple([curr_length - window_length])201 + tuple(past[0].shape[-1:])202 )203 204 ones_mask = torch.ones(ones_key_val_shape)205 ones_mask = decay_mask * ones_mask.permute(0, 1, 2, 4, 3)206 ones_mask = ones_mask.permute(0, 1, 2, 4, 3)207 208 window_mask = torch.cat(209 (ones_mask, torch.zeros(zeros_key_val_shape)),210 dim=-2211 ).to(device)212 else:213 window_mask = torch.ones_like(past[0]).to(device)214 215 # accumulate perturbations for num_iterations216 loss_per_iter = []217 new_accumulated_hidden = None218 for i in range(num_iterations):219 if verbosity_level >= VERBOSE:220 print("Iteration ", i + 1)221 curr_perturbation = [222 to_var(torch.from_numpy(p_), requires_grad=True, device=device)223 for p_ in grad_accumulator224 ]225 226 # Compute hidden using perturbed past227 perturbed_past = list(map(add, past, curr_perturbation))228 _, _, _, curr_length, _ = curr_perturbation[0].shape229 all_logits, _, all_hidden = model(last, past_key_values=perturbed_past)230 hidden = all_hidden[-1]231 new_accumulated_hidden = accumulated_hidden + torch.sum(232 hidden,233 dim=1234 ).detach()235 # TODO: Check the layer-norm consistency of this with trained discriminator (Sumanth)236 logits = all_logits[:, -1, :]237 probs = F.softmax(logits, dim=-1)238 239 loss = 0.0240 loss_list = []241 if loss_type == PPLM_BOW or loss_type == PPLM_BOW_DISCRIM:242 for one_hot_bow in one_hot_bows_vectors:243 bow_logits = torch.mm(probs, torch.t(one_hot_bow))244 bow_loss = -torch.log(torch.sum(bow_logits))245 loss += bow_loss246 loss_list.append(bow_loss)247 if verbosity_level >= VERY_VERBOSE:248 print(" pplm_bow_loss:", loss.data.cpu().numpy())249 250 if loss_type == PPLM_DISCRIM or loss_type == PPLM_BOW_DISCRIM:251 ce_loss = torch.nn.CrossEntropyLoss()252 # TODO why we need to do this assignment and not just using unpert_past? (Sumanth)253 curr_unpert_past = unpert_past254 curr_probs = torch.unsqueeze(probs, dim=1)255 wte = model.resize_token_embeddings()256 for _ in range(horizon_length):257 inputs_embeds = torch.matmul(curr_probs, wte.weight.data)258 _, curr_unpert_past, curr_all_hidden = model(259 past_key_values=curr_unpert_past,260 inputs_embeds=inputs_embeds261 )262 curr_hidden = curr_all_hidden[-1]263 new_accumulated_hidden = new_accumulated_hidden + torch.sum(264 curr_hidden, dim=1)265 266 prediction = classifier(new_accumulated_hidden /267 (curr_length + 1 + horizon_length))268 269 label = torch.tensor(prediction.shape[0] * [class_label],270 device=device,271 dtype=torch.long)272 discrim_loss = ce_loss(prediction, label)273 if verbosity_level >= VERY_VERBOSE:274 print(" pplm_discrim_loss:", discrim_loss.data.cpu().numpy())275 loss += discrim_loss276 loss_list.append(discrim_loss)277 278 kl_loss = 0.0279 if kl_scale > 0.0:280 unpert_probs = F.softmax(unpert_logits[:, -1, :], dim=-1)281 unpert_probs = (282 unpert_probs + SMALL_CONST *283 (unpert_probs <= SMALL_CONST).float().to(device).detach()284 )285 correction = SMALL_CONST * (probs <= SMALL_CONST).float().to(286 device).detach()287 corrected_probs = probs + correction.detach()288 kl_loss = kl_scale * (289 (corrected_probs * (corrected_probs / unpert_probs).log()).sum()290 )291 if verbosity_level >= VERY_VERBOSE:292 print(' kl_loss', kl_loss.data.cpu().numpy())293 loss += kl_loss294 295 loss_per_iter.append(loss.data.cpu().numpy())296 if verbosity_level >= VERBOSE:297 print(' pplm_loss', (loss - kl_loss).data.cpu().numpy())298 299 # compute gradients300 loss.backward()301 302 # calculate gradient norms303 if grad_norms is not None and loss_type == PPLM_BOW:304 grad_norms = [305 torch.max(grad_norms[index], torch.norm(p_.grad * window_mask))306 for index, p_ in enumerate(curr_perturbation)307 ]308 else:309 grad_norms = [310 (torch.norm(p_.grad * window_mask) + SMALL_CONST)311 for index, p_ in enumerate(curr_perturbation)312 ]313 314 # normalize gradients315 grad = [316 -stepsize *317 (p_.grad * window_mask / grad_norms[318 index] ** gamma).data.cpu().numpy()319 for index, p_ in enumerate(curr_perturbation)320 ]321 322 # accumulate gradient323 grad_accumulator = list(map(add, grad, grad_accumulator))324 325 # reset gradients, just to make sure326 for p_ in curr_perturbation:327 p_.grad.data.zero_()328 329 # removing past from the graph330 new_past = []331 for p_ in past:332 new_past.append(p_.detach())333 past = new_past334 335 # apply the accumulated perturbations to the past336 grad_accumulator = [337 to_var(torch.from_numpy(p_), requires_grad=True, device=device)338 for p_ in grad_accumulator339 ]340 pert_past = list(map(add, past, grad_accumulator))341 342 return pert_past, new_accumulated_hidden, grad_norms, loss_per_iter343 344 345def get_classifier(346 name: Optional[str],347 class_label: Union[str, int],348 device: str,349 verbosity_level: int = REGULAR,350 fp: str = None,351 is_deep: bool= False,352 is_deeper: bool=False,353) -> Tuple[Optional[ClassificationHead], Optional[int]]:354 if name is None:355 return None, None356 357 params = DISCRIMINATOR_MODELS_PARAMS[name]358 classifier = ClassificationHead(359 class_size=params['class_size'],360 embed_size=params['embed_size'],361 is_deep=is_deep,362 is_deeper=is_deeper363 ).to(device)364 if "url" in params:365 resolved_archive_file = cached_path(params["url"])366 elif "path" in params:367 resolved_archive_file = params["path"]368 elif fp != None:369 resolved_archive_file = fp370 else:371 raise ValueError("Either url or path have to be specified "372 "in the discriminator model parameters")373 classifier.load_state_dict(374 torch.load(resolved_archive_file, map_location=device))375 classifier.eval()376 377 if isinstance(class_label, str):378 if class_label in params["class_vocab"]:379 label_id = params["class_vocab"][class_label]380 else:381 label_id = params["default_class"]382 if verbosity_level >= REGULAR:383 print("class_label {} not in class_vocab".format(class_label))384 print("available values are: {}".format(params["class_vocab"]))385 print("using default class {}".format(label_id))386 387 elif isinstance(class_label, int):388 if class_label in set(params["class_vocab"].values()):389 label_id = class_label390 else:391 label_id = params["default_class"]392 if verbosity_level >= REGULAR:393 print("class_label {} not in class_vocab".format(class_label))394 print("available values are: {}".format(params["class_vocab"]))395 print("using default class {}".format(label_id))396 397 else:398 label_id = params["default_class"]399 400 return classifier, label_id401 402 403def get_bag_of_words_indices(bag_of_words_ids_or_paths: List[str], tokenizer) -> \404 List[List[List[int]]]:405 bow_indices = []406 for id_or_path in bag_of_words_ids_or_paths:407 if id_or_path in BAG_OF_WORDS_ARCHIVE_MAP:408 filepath = cached_path(BAG_OF_WORDS_ARCHIVE_MAP[id_or_path])409 else:410 filepath = id_or_path411 with open(filepath, "r") as f:412 words = f.read().strip().split("\n")413 bow_indices.append(414 [tokenizer.encode(word.strip(),415 add_prefix_space=True,416 add_special_tokens=False)417 for word in words])418 return bow_indices419 420 421def build_bows_one_hot_vectors(bow_indices, tokenizer, device='cuda'):422 if bow_indices is None:423 return None424 425 one_hot_bows_vectors = []426 for single_bow in bow_indices:427 single_bow = list(filter(lambda x: len(x) <= 1, single_bow))428 single_bow = torch.tensor(single_bow).to(device)429 num_words = single_bow.shape[0]430 one_hot_bow = torch.zeros(num_words, tokenizer.vocab_size).to(device)431 one_hot_bow.scatter_(1, single_bow, 1)432 one_hot_bows_vectors.append(one_hot_bow)433 return one_hot_bows_vectors434 435 436def full_text_generation(437 model,438 tokenizer,439 context=None,440 num_samples=1,441 device="cuda",442 bag_of_words=None,443 discrim=None,444 class_label=None,445 length=100,446 stepsize=0.02,447 temperature=1.0,448 top_k=10,449 sample=True,450 num_iterations=3,451 grad_length=10000,452 horizon_length=1,453 window_length=0,454 decay=False,455 gamma=1.5,456 gm_scale=0.9,457 kl_scale=0.01,458 verbosity_level=REGULAR,459 fp=None,460 is_deep=False,461 is_deeper=False,462 stop_eot=False,463 **kwargs464):465 classifier, class_id = get_classifier(466 discrim,467 class_label,468 device,469 REGULAR,470 fp,471 is_deep,472 is_deeper473 )474 475 bow_indices = []476 if bag_of_words:477 bow_indices = get_bag_of_words_indices(bag_of_words.split(";"),478 tokenizer)479 480 if bag_of_words and classifier:481 loss_type = PPLM_BOW_DISCRIM482 if verbosity_level >= REGULAR:483 print("Both PPLM-BoW and PPLM-Discrim are on. "484 "This is not optimized.")485 486 elif bag_of_words:487 loss_type = PPLM_BOW488 if verbosity_level >= REGULAR:489 print("Using PPLM-BoW")490 491 elif classifier is not None:492 loss_type = PPLM_DISCRIM493 if verbosity_level >= REGULAR:494 print("Using PPLM-Discrim")495 496 else:497 raise Exception("Specify either a bag of words or a discriminator")498 499 unpert_gen_tok_text, _, _, _ = generate_text_pplm(500 model=model,501 tokenizer=tokenizer,502 context=context,503 device=device,504 length=length,505 sample=sample,506 perturb=False,507 verbosity_level=verbosity_level,508 stop_eot=stop_eot509 )510 if device == 'cuda':511 torch.cuda.empty_cache()512 513 pert_gen_tok_texts = []514 discrim_losses = []515 losses_in_time = []516 perplexities = []517 518 for i in range(num_samples):519 pert_gen_tok_text, discrim_loss, loss_in_time, perplexity = generate_text_pplm(520 model=model,521 tokenizer=tokenizer,522 context=context,523 device=device,524 perturb=True,525 bow_indices=bow_indices,526 classifier=classifier,527 class_label=class_id,528 loss_type=loss_type,529 length=length,530 stepsize=stepsize,531 temperature=temperature,532 top_k=top_k,533 sample=sample,534 num_iterations=num_iterations,535 grad_length=grad_length,536 horizon_length=horizon_length,537 window_length=window_length,538 decay=decay,539 gamma=gamma,540 gm_scale=gm_scale,541 kl_scale=kl_scale,542 verbosity_level=verbosity_level,543 stop_eot=stop_eot544 )545 pert_gen_tok_texts.append(pert_gen_tok_text)546 if classifier is not None:547 discrim_losses.append(discrim_loss.data.cpu().numpy())548 losses_in_time.append(loss_in_time)549 perplexities.append(perplexity)550 551 if device == 'cuda':552 torch.cuda.empty_cache()553 554 return unpert_gen_tok_text, pert_gen_tok_texts, discrim_losses, losses_in_time, perplexities555 556 557def generate_text_pplm(558 model,559 tokenizer,560 context=None,561 past=None,562 device="cuda",563 perturb=True,564 bow_indices=None,565 classifier=None,566 class_label=None,567 loss_type=0,568 length=100,569 stepsize=0.02,570 temperature=1.0,571 top_k=10,572 sample=True,573 num_iterations=3,574 grad_length=10000,575 horizon_length=1,576 window_length=0,577 decay=False,578 gamma=1.5,579 gm_scale=0.9,580 kl_scale=0.01,581 verbosity_level=REGULAR,582 stop_eot=False583):584 output_so_far = None585 if context:586 context_t = torch.tensor(context, device=device, dtype=torch.long)587 while len(context_t.shape) < 2:588 context_t = context_t.unsqueeze(0)589 output_so_far = context_t590 591 # collect one hot vectors for bags of words592 one_hot_bows_vectors = build_bows_one_hot_vectors(bow_indices, tokenizer,593 device)594 595 grad_norms = None596 last = None597 unpert_discrim_loss = 0598 loss_in_time = []599 600 if verbosity_level >= VERBOSE:601 range_func = trange(length, ascii=True)602 else:603 range_func = range(length)604 605 pert_total_prob = 1606 pert_times = 0607 last_reps = torch.ones(50257)608 last_reps = last_reps.to(device)609 for i in range_func:610 # Get past/probs for current output, except for last word611 # Note that GPT takes 2 inputs: past + current_token612 613 # run model forward to obtain unperturbed614 if past is None and output_so_far is not None:615 last = output_so_far[:, -1:]616 if output_so_far.shape[1] > 1:617 _, past, _ = model(output_so_far[:, :-1])618 619 unpert_logits, unpert_past, unpert_all_hidden = model(output_so_far)620 unpert_last_hidden = unpert_all_hidden[-1]621 622 # check if we are abowe grad max length623 if i >= grad_length:624 current_stepsize = stepsize * 0625 else:626 current_stepsize = stepsize627 628 # modify the past if necessary629 if not perturb or num_iterations == 0:630 pert_past = past631 632 else:633 accumulated_hidden = unpert_last_hidden[:, :-1, :]634 accumulated_hidden = torch.sum(accumulated_hidden, dim=1)635 636 if past is not None:637 pert_past, _, grad_norms, loss_this_iter = perturb_past(638 past,639 model,640 last,641 unpert_past=unpert_past,642 unpert_logits=unpert_logits,643 accumulated_hidden=accumulated_hidden,644 grad_norms=grad_norms,645 stepsize=current_stepsize,646 one_hot_bows_vectors=one_hot_bows_vectors,647 classifier=classifier,648 class_label=class_label,649 loss_type=loss_type,650 num_iterations=num_iterations,651 horizon_length=horizon_length,652 window_length=window_length,653 decay=decay,654 gamma=gamma,655 kl_scale=kl_scale,656 device=device,657 verbosity_level=verbosity_level658 )659 loss_in_time.append(loss_this_iter)660 else:661 pert_past = past662 663 pert_logits, past, pert_all_hidden = model(last, past_key_values=pert_past)664 pert_logits = pert_logits[:, -1, :] / temperature # + SMALL_CONST665 pert_probs = F.softmax(pert_logits, dim=-1)666 667 if classifier is not None:668 ce_loss = torch.nn.CrossEntropyLoss()669 prediction = classifier(torch.mean(unpert_last_hidden, dim=1))670 label = torch.tensor([class_label], device=device,671 dtype=torch.long)672 unpert_discrim_loss = ce_loss(prediction, label)673 if verbosity_level >= VERBOSE:674 print(675 "unperturbed discrim loss",676 unpert_discrim_loss.data.cpu().numpy()677 )678 else:679 unpert_discrim_loss = 0680 681 # Fuse the modified model and original model682 if perturb:683 684 unpert_probs = F.softmax(unpert_logits[:, -1, :], dim=-1)685 686 pert_probs = ((pert_probs ** gm_scale) * (687 unpert_probs ** (1 - gm_scale))) # + SMALL_CONST688 if i < 2:689 pert_probs = top_k_filter(pert_probs, k=max(2, top_k), probs=True) # + SMALL_CONST690 if i == 0: pert_probs[0][50256] = 0691 if i == 1: 692 tmp = pert_probs[0][50256]693 pert_probs[0][50256] = 0694 pert_probs[0][50256] = min(torch.max(pert_probs[0]), tmp)695 else:696 pert_probs = top_k_filter(pert_probs, k=top_k, probs=True) # + SMALL_CONST697 pert_probs = torch.div(pert_probs, last_reps)698 # rescale699 if torch.sum(pert_probs) <= 1:700 pert_probs = pert_probs / torch.sum(pert_probs)701 else:702 pert_logits = top_k_filter(pert_logits, k=top_k) # + SMALL_CONST703 pert_probs = F.softmax(pert_logits, dim=-1)704 705 # sample or greedy706 if sample:707 last = torch.multinomial(pert_probs, num_samples=1)708 pert_total_prob = pert_total_prob * pert_probs[0][last[0][0]]709 else:710 _, last = torch.topk(pert_probs, k=1, dim=-1)711 last_reps[last[0][0]] = last_reps[last[0][0]] * 8712 # update context/output_so_far appending the new token713 output_so_far = (714 last if output_so_far is None715 else torch.cat((output_so_far, last), dim=1)716 )717 if verbosity_level >= REGULAR:718 print(tokenizer.decode(output_so_far.tolist()[0]))719 pert_times += 1720 if last[0][0] == 50256 and stop_eot: 721 break722 perplexity = (1/pert_total_prob)**(1/pert_times)723 return output_so_far, unpert_discrim_loss, loss_in_time, perplexity724 725 726def set_generic_model_params(discrim_weights, discrim_meta):727 if discrim_weights is None:728 raise ValueError('When using a generic discriminator, '729 'discrim_weights need to be specified')730 if discrim_meta is None:731 raise ValueError('When using a generic discriminator, '732 'discrim_meta need to be specified')733 734 with open(discrim_meta, 'r') as discrim_meta_file:735 meta = json.load(discrim_meta_file)736 meta['path'] = discrim_weights737 DISCRIMINATOR_MODELS_PARAMS['generic'] = meta738 739 740pretrained_model="microsoft/DialoGPT-large"741cond_text=""742uncond=False743num_samples=1744bag_of_words=None745discrim="3_PerSoothe_lrg"746discrim_weights=None747discrim_meta=None748class_label=0749length=100750stepsize=0.32751temperature=1.3752top_k=2753sample=True754num_iterations=0755grad_length=10000756horizon_length=1757window_length=0758decay=False759gamma=1.0760gm_scale=0.95761kl_scale=0.01762seed=0763no_cuda=False764colorama=False765verbosity="quiet"766fp="./paper_code/discrim_models/persoothe_classifier.pt"767model_fp=None768calc_perplexity=False769is_deep=False770is_deeper=True771stop_eot=True772 773# set Random seed774torch.manual_seed(seed)775np.random.seed(seed)776 777# set verbosiry778verbosity_level = VERBOSITY_LEVELS.get(verbosity.lower(), REGULAR)779 780# set the device781device = "cuda" if torch.cuda.is_available() and not no_cuda else "cpu"782 783if discrim == 'generic':784 set_generic_model_params(discrim_weights, discrim_meta)785 786if discrim is not None:787 discriminator_pretrained_model = DISCRIMINATOR_MODELS_PARAMS[discrim][788 "pretrained_model"789 ]790 if pretrained_model != discriminator_pretrained_model:791 pretrained_model = discriminator_pretrained_model792 if verbosity_level >= REGULAR:793 print("discrim = {}, pretrained_model set "794 "to discriminator's = {}".format(discrim, pretrained_model))795 796# load pretrained model797model = GPT2LMHeadModel.from_pretrained(798 pretrained_model,799 output_hidden_states=True800)801if model_fp != None and model_fp != "": 802 model.load_state_dict(torch.load(model_fp, map_location=device))803model.to(device)804model.eval()805 806# load tokenizer807tokenizer = GPT2Tokenizer.from_pretrained(pretrained_model)808 809# Freeze GPT-2 weights810for param in model.parameters():811 param.requires_grad = False812 813starters = ["How are you feeling and why?", "Tell me about your day", "What would you like to talk about?"]814eot_token = "<|endoftext|>"815 816def get_reply(response, username = None, histories = {}, in_stepsize = 0.32, in_horizon_length = 1, in_num_iterations = 0, in_top_k = 2):817 if username == None or username == "": return "<div class='chatbot'>Enter a username</div>", histories818 stepsize = in_stepsize819 horizon_length = int(in_horizon_length)820 num_iterations = int(in_num_iterations)821 top_k = int(in_top_k)822 if response.endswith(("bye", "Bye", "bye.", "Bye.", "bye!", "Bye!","Hello", "Hi", "hello")):823 starter = choice(starters)824 histories[username] = starter+"<|endoftext|>"825 html = "<div class='chatbot'> Chatbot restarted"826 html += "<div class='msg user'>"+starter+"</div>"827 html += "</div>"828 return html, histories829 history = histories.get(username, None)830 convo_hist = (history if history != None else "How are you?<|endoftext|>") + response + eot_token831 # figure out conditioning text832 tokenized_cond_text = tokenizer.encode(833 eot_token + convo_hist,834 add_special_tokens=False835 )836 # generate perturbed texts837 838 # full_text_generation returns:839 # unpert_gen_tok_text, pert_gen_tok_texts, discrim_losses, losses_in_time840 _, pert_gen_tok_texts, _, _, _ = full_text_generation(841 model=model,842 tokenizer=tokenizer,843 context=tokenized_cond_text,844 device=device,845 num_samples=1,846 bag_of_words=bag_of_words,847 discrim=discrim,848 class_label=class_label,849 length=length,850 stepsize=stepsize,851 temperature=temperature,852 top_k=top_k,853 sample=sample,854 num_iterations=num_iterations,855 grad_length=grad_length,856 horizon_length=horizon_length,857 window_length=window_length,858 decay=decay,859 gamma=gamma,860 gm_scale=gm_scale,861 kl_scale=kl_scale,862 verbosity_level=verbosity_level,863 fp=fp,864 is_deep=is_deep,865 is_deeper=is_deeper,866 stop_eot=stop_eot867 )868 869 # iterate through the perturbed texts870 for i, pert_gen_tok_text in enumerate(pert_gen_tok_texts):871 try:872 pert_gen_text = tokenizer.decode(pert_gen_tok_text.tolist()[0])873 convo_hist_split = pert_gen_text.split(eot_token)874 html = "<div class='chatbot'>"875 for m, msg in enumerate(convo_hist_split[1:-1]):876 cls = "user" if m%2 == 0 else "bot"877 html += "<div class='msg {}'> {}</div>".format(cls, msg)878 html += "</div>"879 880 if len(convo_hist_split) > 4: convo_hist_split = convo_hist_split[-4:]881 convo_hist = eot_token.join(convo_hist_split)882 883 except:884 starter = choice(starters)885 histories[username] = starter+"<|endoftext|>"886 html = "<div class='chatbot'> Chatbot restarted"887 html += "<div class='msg user'>"+starter+"</div>"888 html += "</div>"889 return html, histories890 histories[username] = convo_hist891 return html, histories892 893css = """894.chatbox {display:flex;flex-direction:column}895.msg {padding:4px;margin-bottom:4px;border-radius:4px;width:80%}896.msg.user {background-color:cornflowerblue;color:white}897.msg.bot {background-color:lightgray;align-self:self-end}898.footer {display:none !important}899"""900 901gr.Interface(fn=get_reply,902 theme="default",903 inputs=[gr.inputs.Textbox(placeholder="How are you?"), 904 gr.inputs.Textbox(label="Username"),905 "state"],906 outputs=["html", "state"],907 css=css).launch(debug=True, enable_queue=True, share=True) 908