run_pplm.py 27 KB
Newer Older
Piero Molino's avatar
Piero Molino committed
1
#! /usr/bin/env python3
Julien Chaumond's avatar
Julien Chaumond committed
2
# coding=utf-8
Rosanne Liu's avatar
Rosanne Liu committed
3

4
# Copyright (c) 2019 Uber Technologies, Inc.
Julien Chaumond's avatar
Julien Chaumond committed
5
#
6
7
8
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
Julien Chaumond's avatar
Julien Chaumond committed
9
#
10
# http://www.apache.org/licenses/LICENSE-2.0
Julien Chaumond's avatar
Julien Chaumond committed
11
#
12
13
14
15
16
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
Julien Chaumond's avatar
Julien Chaumond committed
17
18
19
20
21
22

"""
Example command with bag of words:
python 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.95

Example command with discriminator:
23
python 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.95
Julien Chaumond's avatar
Julien Chaumond committed
24
25
26
"""

import argparse
27
import json
Julien Chaumond's avatar
Julien Chaumond committed
28
29
30
31
32
33
34
35
36
from operator import add
from typing import List, Optional, Tuple, Union

import numpy as np
import torch
import torch.nn.functional as F
from torch.autograd import Variable
from tqdm import trange

Aymeric Augustin's avatar
Aymeric Augustin committed
37
from pplm_classification_head import ClassificationHead
Julien Chaumond's avatar
Julien Chaumond committed
38
39
40
from transformers import GPT2Tokenizer
from transformers.file_utils import cached_path
from transformers.modeling_gpt2 import GPT2LMHeadModel
Aymeric Augustin's avatar
Aymeric Augustin committed
41

Julien Chaumond's avatar
Julien Chaumond committed
42
43
44
45
46

PPLM_BOW = 1
PPLM_DISCRIM = 2
PPLM_BOW_DISCRIM = 3
SMALL_CONST = 1e-15
47
BIG_CONST = 1e10
Julien Chaumond's avatar
Julien Chaumond committed
48
49

BAG_OF_WORDS_ARCHIVE_MAP = {
50
51
52
53
54
55
56
    "legal": "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/legal.txt",
    "military": "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/military.txt",
    "politics": "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/politics.txt",
    "religion": "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/religion.txt",
    "science": "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/science.txt",
    "space": "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/space.txt",
    "technology": "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/technology.txt",
Julien Chaumond's avatar
Julien Chaumond committed
57
58
59
60
}

DISCRIMINATOR_MODELS_PARAMS = {
    "clickbait": {
Julien Chaumond's avatar
Julien Chaumond committed
61
        "url": "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/discriminators/clickbait_classifier_head.pt",
Julien Chaumond's avatar
Julien Chaumond committed
62
63
64
65
        "class_size": 2,
        "embed_size": 1024,
        "class_vocab": {"non_clickbait": 0, "clickbait": 1},
        "default_class": 1,
66
        "pretrained_model": "gpt2-medium",
Julien Chaumond's avatar
Julien Chaumond committed
67
68
    },
    "sentiment": {
Julien Chaumond's avatar
Julien Chaumond committed
69
        "url": "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/discriminators/SST_classifier_head.pt",
Julien Chaumond's avatar
Julien Chaumond committed
70
71
72
73
        "class_size": 5,
        "embed_size": 1024,
        "class_vocab": {"very_positive": 2, "very_negative": 3},
        "default_class": 3,
74
        "pretrained_model": "gpt2-medium",
Julien Chaumond's avatar
Julien Chaumond committed
75
76
77
78
    },
}


79
80
def to_var(x, requires_grad=False, volatile=False, device="cuda"):
    if torch.cuda.is_available() and device == "cuda":
Piero Molino's avatar
Piero Molino committed
81
        x = x.cuda()
82
    elif device != "cuda":
83
        x = x.to(device)
Piero Molino's avatar
Piero Molino committed
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
    return Variable(x, requires_grad=requires_grad, volatile=volatile)


def top_k_filter(logits, k, probs=False):
    """
    Masks everything but the k top entries as -infinity (1e10).
    Used to mask logits such that e^-infinity -> 0 won't contribute to the
    sum of the denominator.
    """
    if k == 0:
        return logits
    else:
        values = torch.topk(logits, k)[0]
        batch_mins = values[:, -1].view(-1, 1).expand_as(logits)
        if probs:
99
100
            return torch.where(logits < batch_mins, torch.ones_like(logits) * 0.0, logits)
        return torch.where(logits < batch_mins, torch.ones_like(logits) * -BIG_CONST, logits)
Piero Molino's avatar
Piero Molino committed
101
102


103
def perturb_past(
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
    past,
    model,
    last,
    unpert_past=None,
    unpert_logits=None,
    accumulated_hidden=None,
    grad_norms=None,
    stepsize=0.01,
    one_hot_bows_vectors=None,
    classifier=None,
    class_label=None,
    loss_type=0,
    num_iterations=3,
    horizon_length=1,
    window_length=0,
    decay=False,
    gamma=1.5,
    kl_scale=0.01,
    device="cuda",
123
):
Piero Molino's avatar
Piero Molino committed
124
    # Generate inital perturbed past
125
    grad_accumulator = [(np.zeros(p.shape).astype("float32")) for p in past]
Julien Chaumond's avatar
Julien Chaumond committed
126
127
128
129

    if accumulated_hidden is None:
        accumulated_hidden = 0

130
    if decay:
131
        decay_mask = torch.arange(0.0, 1.0 + SMALL_CONST, 1.0 / (window_length))[1:]
Julien Chaumond's avatar
Julien Chaumond committed
132
133
134
    else:
        decay_mask = 1.0

135
    # TODO fix this comment (SUMANTH)
Piero Molino's avatar
Piero Molino committed
136
    # Generate a mask is gradient perturbated is based on a past window
137
    _, _, _, curr_length, _ = past[0].shape
Piero Molino's avatar
Piero Molino committed
138

139
    if curr_length > window_length and window_length > 0:
140
        ones_key_val_shape = tuple(past[0].shape[:-2]) + tuple([window_length]) + tuple(past[0].shape[-1:])
Piero Molino's avatar
Piero Molino committed
141

142
        zeros_key_val_shape = (
143
            tuple(past[0].shape[:-2]) + tuple([curr_length - window_length]) + tuple(past[0].shape[-1:])
144
        )
Julien Chaumond's avatar
Julien Chaumond committed
145
146
147
148
149

        ones_mask = torch.ones(ones_key_val_shape)
        ones_mask = decay_mask * ones_mask.permute(0, 1, 2, 4, 3)
        ones_mask = ones_mask.permute(0, 1, 2, 4, 3)

150
        window_mask = torch.cat((ones_mask, torch.zeros(zeros_key_val_shape)), dim=-2).to(device)
Julien Chaumond's avatar
Julien Chaumond committed
151
    else:
152
        window_mask = torch.ones_like(past[0]).to(device)
Julien Chaumond's avatar
Julien Chaumond committed
153

154
    # accumulate perturbations for num_iterations
Julien Chaumond's avatar
Julien Chaumond committed
155
    loss_per_iter = []
156
    new_accumulated_hidden = None
157
    for i in range(num_iterations):
Julien Chaumond's avatar
Julien Chaumond committed
158
        print("Iteration ", i + 1)
159
        curr_perturbation = [
160
            to_var(torch.from_numpy(p_), requires_grad=True, device=device) for p_ in grad_accumulator
161
162
163
164
165
166
        ]

        # Compute hidden using perturbed past
        perturbed_past = list(map(add, past, curr_perturbation))
        _, _, _, curr_length, _ = curr_perturbation[0].shape
        all_logits, _, all_hidden = model(last, past=perturbed_past)
Piero Molino's avatar
Piero Molino committed
167
        hidden = all_hidden[-1]
168
        new_accumulated_hidden = accumulated_hidden + torch.sum(hidden, dim=1).detach()
169
170
171
        # TODO: Check the layer-norm consistency of this with trained discriminator (Sumanth)
        logits = all_logits[:, -1, :]
        probs = F.softmax(logits, dim=-1)
Piero Molino's avatar
Piero Molino committed
172
173
174

        loss = 0.0
        loss_list = []
175
176
177
178
179
180
        if loss_type == PPLM_BOW or loss_type == PPLM_BOW_DISCRIM:
            for one_hot_bow in one_hot_bows_vectors:
                bow_logits = torch.mm(probs, torch.t(one_hot_bow))
                bow_loss = -torch.log(torch.sum(bow_logits))
                loss += bow_loss
                loss_list.append(bow_loss)
Piero Molino's avatar
Piero Molino committed
181
182
            print(" pplm_bow_loss:", loss.data.cpu().numpy())

183
        if loss_type == 2 or loss_type == 3:
Julien Chaumond's avatar
Julien Chaumond committed
184
            ce_loss = torch.nn.CrossEntropyLoss()
185
186
187
188
189
190
            # TODO why we need to do this assignment and not just using unpert_past? (Sumanth)
            curr_unpert_past = unpert_past
            curr_probs = torch.unsqueeze(probs, dim=1)
            wte = model.resize_token_embeddings()
            for _ in range(horizon_length):
                inputs_embeds = torch.matmul(curr_probs, wte.weight.data)
191
                _, curr_unpert_past, curr_all_hidden = model(past=curr_unpert_past, inputs_embeds=inputs_embeds)
192
                curr_hidden = curr_all_hidden[-1]
193
                new_accumulated_hidden = new_accumulated_hidden + torch.sum(curr_hidden, dim=1)
Julien Chaumond's avatar
Julien Chaumond committed
194

195
            prediction = classifier(new_accumulated_hidden / (curr_length + 1 + horizon_length))
Julien Chaumond's avatar
Julien Chaumond committed
196

197
            label = torch.tensor(prediction.shape[0] * [class_label], device=device, dtype=torch.long)
198
            discrim_loss = ce_loss(prediction, label)
Julien Chaumond's avatar
Julien Chaumond committed
199
            print(" pplm_discrim_loss:", discrim_loss.data.cpu().numpy())
Piero Molino's avatar
Piero Molino committed
200
201
            loss += discrim_loss
            loss_list.append(discrim_loss)
Julien Chaumond's avatar
Julien Chaumond committed
202

Piero Molino's avatar
Piero Molino committed
203
204
        kl_loss = 0.0
        if kl_scale > 0.0:
205
            unpert_probs = F.softmax(unpert_logits[:, -1, :], dim=-1)
206
207
            unpert_probs = unpert_probs + SMALL_CONST * (unpert_probs <= SMALL_CONST).float().to(device).detach()
            correction = SMALL_CONST * (probs <= SMALL_CONST).float().to(device).detach()
208
            corrected_probs = probs + correction.detach()
209
210
            kl_loss = kl_scale * ((corrected_probs * (corrected_probs / unpert_probs).log()).sum())
            print(" kl_loss", kl_loss.data.cpu().numpy())
211
            loss += kl_loss
Julien Chaumond's avatar
Julien Chaumond committed
212
213

        loss_per_iter.append(loss.data.cpu().numpy())
214
        print(" pplm_loss", (loss - kl_loss).data.cpu().numpy())
Julien Chaumond's avatar
Julien Chaumond committed
215

216
        # compute gradients
Rosanne Liu's avatar
Rosanne Liu committed
217
        loss.backward()
218
219
220

        # calculate gradient norms
        if grad_norms is not None and loss_type == PPLM_BOW:
Julien Chaumond's avatar
Julien Chaumond committed
221
222
            grad_norms = [
                torch.max(grad_norms[index], torch.norm(p_.grad * window_mask))
223
224
                for index, p_ in enumerate(curr_perturbation)
            ]
Julien Chaumond's avatar
Julien Chaumond committed
225
        else:
226
            grad_norms = [
227
                (torch.norm(p_.grad * window_mask) + SMALL_CONST) for index, p_ in enumerate(curr_perturbation)
228
            ]
Julien Chaumond's avatar
Julien Chaumond committed
229

230
        # normalize gradients
Julien Chaumond's avatar
Julien Chaumond committed
231
        grad = [
232
            -stepsize * (p_.grad * window_mask / grad_norms[index] ** gamma).data.cpu().numpy()
233
234
            for index, p_ in enumerate(curr_perturbation)
        ]
Julien Chaumond's avatar
Julien Chaumond committed
235

236
237
238
239
240
        # accumulate gradient
        grad_accumulator = list(map(add, grad, grad_accumulator))

        # reset gradients, just to make sure
        for p_ in curr_perturbation:
Julien Chaumond's avatar
Julien Chaumond committed
241
242
            p_.grad.data.zero_()

243
        # removing past from the graph
Julien Chaumond's avatar
Julien Chaumond committed
244
        new_past = []
245
246
        for p_ in past:
            new_past.append(p_.detach())
Julien Chaumond's avatar
Julien Chaumond committed
247
248
        past = new_past

249
    # apply the accumulated perturbations to the past
250
    grad_accumulator = [to_var(torch.from_numpy(p_), requires_grad=True, device=device) for p_ in grad_accumulator]
251
    pert_past = list(map(add, past, grad_accumulator))
Julien Chaumond's avatar
Julien Chaumond committed
252

253
    return pert_past, new_accumulated_hidden, grad_norms, loss_per_iter
Julien Chaumond's avatar
Julien Chaumond committed
254
255
256


def get_classifier(
257
    name: Optional[str], class_label: Union[str, int], device: str
Julien Chaumond's avatar
Julien Chaumond committed
258
259
260
261
262
) -> Tuple[Optional[ClassificationHead], Optional[int]]:
    if name is None:
        return None, None

    params = DISCRIMINATOR_MODELS_PARAMS[name]
263
    classifier = ClassificationHead(class_size=params["class_size"], embed_size=params["embed_size"]).to(device)
264
265
    if "url" in params:
        resolved_archive_file = cached_path(params["url"])
266
    elif "path" in params:
267
        resolved_archive_file = params["path"]
268
    else:
269
270
        raise ValueError("Either url or path have to be specified " "in the discriminator model parameters")
    classifier.load_state_dict(torch.load(resolved_archive_file, map_location=device))
Julien Chaumond's avatar
Julien Chaumond committed
271
272
    classifier.eval()

273
274
275
    if isinstance(class_label, str):
        if class_label in params["class_vocab"]:
            label_id = params["class_vocab"][class_label]
Julien Chaumond's avatar
Julien Chaumond committed
276
277
        else:
            label_id = params["default_class"]
278
            print("class_label {} not in class_vocab".format(class_label))
Julien Chaumond's avatar
Julien Chaumond committed
279
280
281
            print("available values are: {}".format(params["class_vocab"]))
            print("using default class {}".format(label_id))

282
283
284
    elif isinstance(class_label, int):
        if class_label in set(params["class_vocab"].values()):
            label_id = class_label
Julien Chaumond's avatar
Julien Chaumond committed
285
286
        else:
            label_id = params["default_class"]
287
            print("class_label {} not in class_vocab".format(class_label))
Julien Chaumond's avatar
Julien Chaumond committed
288
289
290
291
292
293
294
295
296
            print("available values are: {}".format(params["class_vocab"]))
            print("using default class {}".format(label_id))

    else:
        label_id = params["default_class"]

    return classifier, label_id


297
def get_bag_of_words_indices(bag_of_words_ids_or_paths: List[str], tokenizer) -> List[List[List[int]]]:
Julien Chaumond's avatar
Julien Chaumond committed
298
299
300
301
302
303
304
    bow_indices = []
    for id_or_path in bag_of_words_ids_or_paths:
        if id_or_path in BAG_OF_WORDS_ARCHIVE_MAP:
            filepath = cached_path(BAG_OF_WORDS_ARCHIVE_MAP[id_or_path])
        else:
            filepath = id_or_path
        with open(filepath, "r") as f:
Piero Molino's avatar
Piero Molino committed
305
            words = f.read().strip().split("\n")
306
        bow_indices.append([tokenizer.encode(word.strip(), add_prefix_space=True) for word in words])
Julien Chaumond's avatar
Julien Chaumond committed
307
308
309
    return bow_indices


310
def build_bows_one_hot_vectors(bow_indices, tokenizer, device="cuda"):
Julien Chaumond's avatar
Julien Chaumond committed
311
312
313
314
315
316
    if bow_indices is None:
        return None

    one_hot_bows_vectors = []
    for single_bow in bow_indices:
        single_bow = list(filter(lambda x: len(x) <= 1, single_bow))
317
        single_bow = torch.tensor(single_bow).to(device)
Julien Chaumond's avatar
Julien Chaumond committed
318
        num_words = single_bow.shape[0]
319
        one_hot_bow = torch.zeros(num_words, tokenizer.vocab_size).to(device)
Julien Chaumond's avatar
Julien Chaumond committed
320
321
322
323
324
        one_hot_bow.scatter_(1, single_bow, 1)
        one_hot_bows_vectors.append(one_hot_bow)
    return one_hot_bows_vectors


325
def full_text_generation(
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
    model,
    tokenizer,
    context=None,
    num_samples=1,
    device="cuda",
    bag_of_words=None,
    discrim=None,
    class_label=None,
    length=100,
    stepsize=0.02,
    temperature=1.0,
    top_k=10,
    sample=False,
    num_iterations=3,
    grad_length=10000,
    horizon_length=1,
    window_length=0,
    decay=False,
    gamma=1.5,
    gm_scale=0.9,
    kl_scale=0.01,
    **kwargs
348
):
349
    classifier, class_id = get_classifier(discrim, class_label, device)
Julien Chaumond's avatar
Julien Chaumond committed
350

351
352
    bow_indices = []
    if bag_of_words:
353
        bow_indices = get_bag_of_words_indices(bag_of_words.split(";"), tokenizer)
Piero Molino's avatar
Piero Molino committed
354

355
    if bag_of_words and classifier:
Julien Chaumond's avatar
Julien Chaumond committed
356
        print("Both PPLM-BoW and PPLM-Discrim are on. This is not optimized.")
357
        loss_type = PPLM_BOW_DISCRIM
Julien Chaumond's avatar
Julien Chaumond committed
358

359
360
    elif bag_of_words:
        loss_type = PPLM_BOW
Julien Chaumond's avatar
Julien Chaumond committed
361
362
363
        print("Using PPLM-BoW")

    elif classifier is not None:
364
        loss_type = PPLM_DISCRIM
Julien Chaumond's avatar
Julien Chaumond committed
365
366
367
        print("Using PPLM-Discrim")

    else:
368
        raise Exception("Specify either a bag of words or a discriminator")
Julien Chaumond's avatar
Julien Chaumond committed
369

370
    unpert_gen_tok_text, _, _ = generate_text_pplm(
371
        model=model, tokenizer=tokenizer, context=context, device=device, length=length, sample=sample, perturb=False
372
    )
373
    if device == "cuda":
374
        torch.cuda.empty_cache()
Julien Chaumond's avatar
Julien Chaumond committed
375

376
377
378
    pert_gen_tok_texts = []
    discrim_losses = []
    losses_in_time = []
Piero Molino's avatar
Piero Molino committed
379

380
    for i in range(num_samples):
381
        pert_gen_tok_text, discrim_loss, loss_in_time = generate_text_pplm(
382
            model=model,
383
            tokenizer=tokenizer,
384
385
386
387
388
            context=context,
            device=device,
            perturb=True,
            bow_indices=bow_indices,
            classifier=classifier,
389
            class_label=class_id,
390
391
392
393
394
            loss_type=loss_type,
            length=length,
            stepsize=stepsize,
            temperature=temperature,
            top_k=top_k,
395
396
397
            sample=sample,
            num_iterations=num_iterations,
            grad_length=grad_length,
398
            horizon_length=horizon_length,
399
            window_length=window_length,
400
401
            decay=decay,
            gamma=gamma,
402
403
            gm_scale=gm_scale,
            kl_scale=kl_scale,
404
        )
405
        pert_gen_tok_texts.append(pert_gen_tok_text)
Julien Chaumond's avatar
Julien Chaumond committed
406
        if classifier is not None:
407
408
            discrim_losses.append(discrim_loss.data.cpu().numpy())
        losses_in_time.append(loss_in_time)
Julien Chaumond's avatar
Julien Chaumond committed
409

410
    if device == "cuda":
411
        torch.cuda.empty_cache()
Julien Chaumond's avatar
Julien Chaumond committed
412

413
    return unpert_gen_tok_text, pert_gen_tok_texts, discrim_losses, losses_in_time
Julien Chaumond's avatar
Julien Chaumond committed
414

415
416

def generate_text_pplm(
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
    model,
    tokenizer,
    context=None,
    past=None,
    device="cuda",
    perturb=True,
    bow_indices=None,
    classifier=None,
    class_label=None,
    loss_type=0,
    length=100,
    stepsize=0.02,
    temperature=1.0,
    top_k=10,
    sample=False,
    num_iterations=3,
    grad_length=10000,
    horizon_length=1,
    window_length=0,
    decay=False,
    gamma=1.5,
    gm_scale=0.9,
    kl_scale=0.01,
440
):
441
442
443
444
445
446
    output_so_far = None
    if context:
        context_t = torch.tensor(context, device=device, dtype=torch.long)
        while len(context_t.shape) < 2:
            context_t = context_t.unsqueeze(0)
        output_so_far = context_t
Julien Chaumond's avatar
Julien Chaumond committed
447

448
    # collect one hot vectors for bags of words
449
    one_hot_bows_vectors = build_bows_one_hot_vectors(bow_indices, tokenizer, device)
450

Julien Chaumond's avatar
Julien Chaumond committed
451
    grad_norms = None
452
    last = None
453
    unpert_discrim_loss = 0
Julien Chaumond's avatar
Julien Chaumond committed
454
    loss_in_time = []
455
    for i in trange(length, ascii=True):
Julien Chaumond's avatar
Julien Chaumond committed
456
457

        # Get past/probs for current output, except for last word
458
        # Note that GPT takes 2 inputs: past + current_token
Julien Chaumond's avatar
Julien Chaumond committed
459

460
461
462
        # run model forward to obtain unperturbed
        if past is None and output_so_far is not None:
            last = output_so_far[:, -1:]
463
464
            if output_so_far.shape[1] > 1:
                _, past, _ = model(output_so_far[:, :-1])
Piero Molino's avatar
Piero Molino committed
465

466
467
        unpert_logits, unpert_past, unpert_all_hidden = model(output_so_far)
        unpert_last_hidden = unpert_all_hidden[-1]
Piero Molino's avatar
Piero Molino committed
468

469
        # check if we are abowe grad max length
470
471
        if i >= grad_length:
            current_stepsize = stepsize * 0
Julien Chaumond's avatar
Julien Chaumond committed
472
        else:
473
            current_stepsize = stepsize
Julien Chaumond's avatar
Julien Chaumond committed
474

475
        # modify the past if necessary
476
        if not perturb or num_iterations == 0:
477
            pert_past = past
Julien Chaumond's avatar
Julien Chaumond committed
478
479

        else:
480
            accumulated_hidden = unpert_last_hidden[:, :-1, :]
Julien Chaumond's avatar
Julien Chaumond committed
481
482
            accumulated_hidden = torch.sum(accumulated_hidden, dim=1)

483
484
485
486
487
488
489
490
491
492
            if past is not None:
                pert_past, _, grad_norms, loss_this_iter = perturb_past(
                    past,
                    model,
                    last,
                    unpert_past=unpert_past,
                    unpert_logits=unpert_logits,
                    accumulated_hidden=accumulated_hidden,
                    grad_norms=grad_norms,
                    stepsize=current_stepsize,
493
                    one_hot_bows_vectors=one_hot_bows_vectors,
494
                    classifier=classifier,
495
                    class_label=class_label,
496
497
498
                    loss_type=loss_type,
                    num_iterations=num_iterations,
                    horizon_length=horizon_length,
499
                    window_length=window_length,
500
501
                    decay=decay,
                    gamma=gamma,
502
503
                    kl_scale=kl_scale,
                    device=device,
504
505
506
507
                )
                loss_in_time.append(loss_this_iter)
            else:
                pert_past = past
Piero Molino's avatar
Piero Molino committed
508

509
510
511
        pert_logits, past, pert_all_hidden = model(last, past=pert_past)
        pert_logits = pert_logits[:, -1, :] / temperature  # + SMALL_CONST
        pert_probs = F.softmax(pert_logits, dim=-1)
Julien Chaumond's avatar
Julien Chaumond committed
512
513

        if classifier is not None:
Piero Molino's avatar
Piero Molino committed
514
            ce_loss = torch.nn.CrossEntropyLoss()
515
            prediction = classifier(torch.mean(unpert_last_hidden, dim=1))
516
            label = torch.tensor([class_label], device=device, dtype=torch.long)
517
            unpert_discrim_loss = ce_loss(prediction, label)
518
            print("unperturbed discrim loss", unpert_discrim_loss.data.cpu().numpy())
Julien Chaumond's avatar
Julien Chaumond committed
519
        else:
520
            unpert_discrim_loss = 0
Piero Molino's avatar
Piero Molino committed
521
522

        # Fuse the modified model and original model
Julien Chaumond's avatar
Julien Chaumond committed
523
524
        if perturb:

525
            unpert_probs = F.softmax(unpert_logits[:, -1, :], dim=-1)
Piero Molino's avatar
Piero Molino committed
526

527
528
            pert_probs = (pert_probs ** gm_scale) * (unpert_probs ** (1 - gm_scale))  # + SMALL_CONST
            pert_probs = top_k_filter(pert_probs, k=top_k, probs=True)  # + SMALL_CONST
Julien Chaumond's avatar
Julien Chaumond committed
529

530
531
532
            # rescale
            if torch.sum(pert_probs) <= 1:
                pert_probs = pert_probs / torch.sum(pert_probs)
Julien Chaumond's avatar
Julien Chaumond committed
533
534

        else:
535
536
            pert_logits = top_k_filter(pert_logits, k=top_k)  # + SMALL_CONST
            pert_probs = F.softmax(pert_logits, dim=-1)
Julien Chaumond's avatar
Julien Chaumond committed
537

538
        # sample or greedy
Julien Chaumond's avatar
Julien Chaumond committed
539
        if sample:
540
541
            last = torch.multinomial(pert_probs, num_samples=1)

Julien Chaumond's avatar
Julien Chaumond committed
542
        else:
543
544
545
            _, last = torch.topk(pert_probs, k=1, dim=-1)

        # update context/output_so_far appending the new token
546
        output_so_far = last if output_so_far is None else torch.cat((output_so_far, last), dim=1)
547

548
        print(tokenizer.decode(output_so_far.tolist()[0]))
549
550

    return output_so_far, unpert_discrim_loss, loss_in_time
Julien Chaumond's avatar
Julien Chaumond committed
551
552


553
554
def set_generic_model_params(discrim_weights, discrim_meta):
    if discrim_weights is None:
555
        raise ValueError("When using a generic discriminator, " "discrim_weights need to be specified")
556
    if discrim_meta is None:
557
        raise ValueError("When using a generic discriminator, " "discrim_meta need to be specified")
558

559
    with open(discrim_meta, "r") as discrim_meta_file:
560
        meta = json.load(discrim_meta_file)
561
562
    meta["path"] = discrim_weights
    DISCRIMINATOR_MODELS_PARAMS["generic"] = meta
563
564


565
def run_pplm_example(
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
    pretrained_model="gpt2-medium",
    cond_text="",
    uncond=False,
    num_samples=1,
    bag_of_words=None,
    discrim=None,
    discrim_weights=None,
    discrim_meta=None,
    class_label=-1,
    length=100,
    stepsize=0.02,
    temperature=1.0,
    top_k=10,
    sample=False,
    num_iterations=3,
    grad_length=10000,
    horizon_length=1,
    window_length=0,
    decay=False,
    gamma=1.5,
    gm_scale=0.9,
    kl_scale=0.01,
    seed=0,
    no_cuda=False,
    colorama=False,
591
):
592
    # set Random seed
593
594
    torch.manual_seed(seed)
    np.random.seed(seed)
Julien Chaumond's avatar
Julien Chaumond committed
595

596
    # set the device
597
598
    device = "cuda" if torch.cuda.is_available() and not no_cuda else "cpu"

599
    if discrim == "generic":
600
        set_generic_model_params(discrim_weights, discrim_meta)
Julien Chaumond's avatar
Julien Chaumond committed
601

602
    if discrim is not None:
603
604
        pretrained_model = DISCRIMINATOR_MODELS_PARAMS[discrim]["pretrained_model"]
        print("discrim = {}, pretrained_model set " "to discriminator's = {}".format(discrim, pretrained_model))
605

606
    # load pretrained model
607
    model = GPT2LMHeadModel.from_pretrained(pretrained_model, output_hidden_states=True)
Julien Chaumond's avatar
Julien Chaumond committed
608
609
610
    model.to(device)
    model.eval()

611
612
613
    # load tokenizer
    tokenizer = GPT2Tokenizer.from_pretrained(pretrained_model)

Piero Molino's avatar
Piero Molino committed
614
    # Freeze GPT-2 weights
Julien Chaumond's avatar
Julien Chaumond committed
615
616
617
    for param in model.parameters():
        param.requires_grad = False

618
    # figure out conditioning text
619
    if uncond:
620
        tokenized_cond_text = tokenizer.encode([tokenizer.bos_token])
Julien Chaumond's avatar
Julien Chaumond committed
621
    else:
622
        raw_text = cond_text
Julien Chaumond's avatar
Julien Chaumond committed
623
        while not raw_text:
624
            print("Did you forget to add `--cond_text`? ")
Julien Chaumond's avatar
Julien Chaumond committed
625
            raw_text = input("Model prompt >>> ")
626
        tokenized_cond_text = tokenizer.encode(tokenizer.bos_token + raw_text)
Piero Molino's avatar
Piero Molino committed
627

628
    print("= Prefix of sentence =")
629
    print(tokenizer.decode(tokenized_cond_text))
630
    print()
Piero Molino's avatar
Piero Molino committed
631

632
    # generate unperturbed and perturbed texts
Piero Molino's avatar
Piero Molino committed
633

634
635
636
    # full_text_generation returns:
    # unpert_gen_tok_text, pert_gen_tok_texts, discrim_losses, losses_in_time
    unpert_gen_tok_text, pert_gen_tok_texts, _, _ = full_text_generation(
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
        model=model,
        tokenizer=tokenizer,
        context=tokenized_cond_text,
        device=device,
        num_samples=num_samples,
        bag_of_words=bag_of_words,
        discrim=discrim,
        class_label=class_label,
        length=length,
        stepsize=stepsize,
        temperature=temperature,
        top_k=top_k,
        sample=sample,
        num_iterations=num_iterations,
        grad_length=grad_length,
        horizon_length=horizon_length,
        window_length=window_length,
        decay=decay,
        gamma=gamma,
        gm_scale=gm_scale,
        kl_scale=kl_scale,
658
659
660
    )

    # untokenize unperturbed text
661
    unpert_gen_text = tokenizer.decode(unpert_gen_tok_text.tolist()[0])
Piero Molino's avatar
Piero Molino committed
662

663
664
665
666
    print("=" * 80)
    print("= Unperturbed generated text =")
    print(unpert_gen_text)
    print()
Piero Molino's avatar
Piero Molino committed
667

668
669
    generated_texts = []

670
    bow_word_ids = set()
671
    if bag_of_words and colorama:
672
        bow_indices = get_bag_of_words_indices(bag_of_words.split(";"), tokenizer)
673
674
675
676
677
        for single_bow_list in bow_indices:
            # filtering all words in the list composed of more than 1 token
            filtered = list(filter(lambda x: len(x) <= 1, single_bow_list))
            # w[0] because we are sure w has only 1 item because previous fitler
            bow_word_ids.update(w[0] for w in filtered)
678
679
680
681
682

    # iterate through the perturbed texts
    for i, pert_gen_tok_text in enumerate(pert_gen_tok_texts):
        try:
            # untokenize unperturbed text
683
            if colorama:
Piero Molino's avatar
Piero Molino committed
684
685
                import colorama

686
                pert_gen_text = ""
687
                for word_id in pert_gen_tok_text.tolist()[0]:
688
                    if word_id in bow_word_ids:
689
690
                        pert_gen_text += "{}{}{}".format(
                            colorama.Fore.RED, tokenizer.decode([word_id]), colorama.Style.RESET_ALL
691
                        )
Piero Molino's avatar
Piero Molino committed
692
                    else:
693
                        pert_gen_text += tokenizer.decode([word_id])
Piero Molino's avatar
Piero Molino committed
694
            else:
695
                pert_gen_text = tokenizer.decode(pert_gen_tok_text.tolist()[0])
Julien Chaumond's avatar
Julien Chaumond committed
696

697
698
699
            print("= Perturbed generated text {} =".format(i + 1))
            print(pert_gen_text)
            print()
700
701
        except Exception as exc:
            print("Ignoring error while generating perturbed text:", exc)
Julien Chaumond's avatar
Julien Chaumond committed
702

703
        # keep the prefix, perturbed seq, original seq for each index
704
        generated_texts.append((tokenized_cond_text, pert_gen_tok_text, unpert_gen_tok_text))
Julien Chaumond's avatar
Julien Chaumond committed
705

Piero Molino's avatar
Piero Molino committed
706
    return
Julien Chaumond's avatar
Julien Chaumond committed
707
708


709
if __name__ == "__main__":
710
711
712
713
714
715
716
717
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--pretrained_model",
        "-M",
        type=str,
        default="gpt2-medium",
        help="pretrained model name or path to local checkpoint",
    )
718
719
    parser.add_argument("--cond_text", type=str, default="The lake", help="Prefix texts to condition on")
    parser.add_argument("--uncond", action="store_true", help="Generate from end-of-text as prefix")
720
    parser.add_argument(
721
        "--num_samples", type=int, default=1, help="Number of samples to generate from the modified latents",
722
    )
723
724
725
726
727
728
    parser.add_argument(
        "--bag_of_words",
        "-B",
        type=str,
        default=None,
        help="Bags of words used for PPLM-BoW. "
729
730
        "Either a BOW id (see list in code) or a filepath. "
        "Multiple BoWs separated by ;",
731
732
733
734
735
736
737
738
739
    )
    parser.add_argument(
        "--discrim",
        "-D",
        type=str,
        default=None,
        choices=("clickbait", "sentiment", "toxicity", "generic"),
        help="Discriminator to use",
    )
740
    parser.add_argument("--discrim_weights", type=str, default=None, help="Weights for the generic discriminator")
741
    parser.add_argument(
742
743
744
745
        "--discrim_meta", type=str, default=None, help="Meta information for the generic discriminator"
    )
    parser.add_argument(
        "--class_label", type=int, default=-1, help="Class label used for the discriminator",
746
747
    )
    parser.add_argument("--length", type=int, default=100)
748
    parser.add_argument("--stepsize", type=float, default=0.02)
749
750
    parser.add_argument("--temperature", type=float, default=1.0)
    parser.add_argument("--top_k", type=int, default=10)
751
    parser.add_argument("--sample", action="store_true", help="Generate from end-of-text as prefix")
752
753
754
    parser.add_argument("--num_iterations", type=int, default=3)
    parser.add_argument("--grad_length", type=int, default=10000)
    parser.add_argument(
755
        "--window_length",
756
        type=int,
757
        default=0,
758
        help="Length of past which is being optimized; " "0 corresponds to infinite window length",
759
760
    )
    parser.add_argument(
761
        "--horizon_length", type=int, default=1, help="Length of future to optimize over",
762
    )
763
    parser.add_argument("--decay", action="store_true", help="whether to decay or not")
764
    parser.add_argument("--gamma", type=float, default=1.5)
765
766
767
768
    parser.add_argument("--gm_scale", type=float, default=0.9)
    parser.add_argument("--kl_scale", type=float, default=0.01)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--no_cuda", action="store_true", help="no cuda")
769
    parser.add_argument("--colorama", action="store_true", help="colors keywords")
770
771
772

    args = parser.parse_args()
    run_pplm_example(**vars(args))