run_pplm.py 26.7 KB
Newer Older
Piero Molino's avatar
Piero Molino committed
1
#! /usr/bin/env python3
Julien Chaumond's avatar
Julien Chaumond committed
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
# coding=utf-8
# Copyright 2018 The Uber AI Team Authors.
#
# 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
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# 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.

# TODO: add code for training a custom discriminator

"""
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:
python examples/run_pplm.py -D sentiment --label_class 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
"""

import argparse
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

37
from examples.run_pplm_discrim_train import ClassificationHead
Julien Chaumond's avatar
Julien Chaumond committed
38
39
40
41
42
43
44
45
from transformers import GPT2Tokenizer
from transformers.file_utils import cached_path
from transformers.modeling_gpt2 import GPT2LMHeadModel

PPLM_BOW = 1
PPLM_DISCRIM = 2
PPLM_BOW_DISCRIM = 3
SMALL_CONST = 1e-15
Piero Molino's avatar
Piero Molino committed
46
SmallConst = 1e-15
Julien Chaumond's avatar
Julien Chaumond committed
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
TOKENIZER = GPT2Tokenizer.from_pretrained("gpt2-medium")

BAG_OF_WORDS_ARCHIVE_MAP = {
    'kitchen': "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/kitchen.txt",
    '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",
    'monsters': "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/monsters.txt",
    'politics': "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/politics.txt",
    'positive_words': "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/bow/positive_words.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",
}

DISCRIMINATOR_MODELS_PARAMS = {
    "clickbait": {
        "url": "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/discriminators/clickbait_classifierhead.pt",
        "class_size": 2,
        "embed_size": 1024,
        "class_vocab": {"non_clickbait": 0, "clickbait": 1},
        "default_class": 1,
    },
    "sentiment": {
Piero Molino's avatar
Piero Molino committed
71
        "url": "http://s.yosinski.com/SST_classifier_head.pt",
Julien Chaumond's avatar
Julien Chaumond committed
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
        "class_size": 5,
        "embed_size": 1024,
        "class_vocab": {"very_positive": 2, "very_negative": 3},
        "default_class": 3,
    },
    "toxicity": {
        "url": "https://s3.amazonaws.com/models.huggingface.co/bert/pplm/discriminators/toxicity_classifierhead.pt",
        "class_size": 2,
        "embed_size": 1024,
        "class_vocab": {"non_toxic": 0, "toxic": 1},
        "default_class": 0,
    },
}


Piero Molino's avatar
Piero Molino committed
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
def to_var(x, requires_grad=False, volatile=False):
    if torch.cuda.is_available():
        x = x.cuda()
    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:
            return torch.where(logits < batch_mins,
                               torch.ones_like(logits) * 0.0, logits)
        return torch.where(logits < batch_mins, torch.ones_like(logits) * -1e10,
                           logits)


111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
def perturb_past(
        past,
        model,
        prev,
        unpert_past=None,
        unpert_logits=None,
        accumulated_hidden=None,
        grad_norms=None,
        stepsize=0.01,
        classifier=None,
        label_class=None,
        one_hot_bows_vectors=None,
        loss_type=0,
        num_iterations=3,
        kl_scale=0.01,
        window_length=0,
        horizon_length=1,
        decay=False,
        gamma=1.5,
):
Piero Molino's avatar
Piero Molino committed
131
132
133
134
    # Generate inital perturbed past
    past_perturb_orig = [
        (np.random.uniform(0.0, 0.0, p.shape).astype('float32'))
        for p in past]
Julien Chaumond's avatar
Julien Chaumond committed
135
136
137
138

    if accumulated_hidden is None:
        accumulated_hidden = 0

139
    if decay:
Piero Molino's avatar
Piero Molino committed
140
141
        decay_mask = torch.arange(0., 1.0 + SmallConst, 1.0 / (window_length))[
                     1:]
Julien Chaumond's avatar
Julien Chaumond committed
142
143
144
    else:
        decay_mask = 1.0

Piero Molino's avatar
Piero Molino committed
145
146
147
148
149
150
151
152
153
154
155
    # Generate a mask is gradient perturbated is based on a past window
    _, _, _, current_length, _ = past[0].shape

    if current_length > window_length and window_length > 0:
        ones_key_val_shape = tuple(past[0].shape[:-2]) + tuple(
            [window_length]) + tuple(
            past[0].shape[-1:])

        zeros_key_val_shape = tuple(past[0].shape[:-2]) + tuple(
            [current_length - window_length]) + tuple(
            past[0].shape[-1:])
Julien Chaumond's avatar
Julien Chaumond committed
156
157
158
159
160

        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)

Piero Molino's avatar
Piero Molino committed
161
162
        window_mask = torch.cat((ones_mask, torch.zeros(zeros_key_val_shape)),
                                dim=-2).cuda()
Julien Chaumond's avatar
Julien Chaumond committed
163
164
165
166
    else:
        window_mask = torch.ones_like(past[0]).cuda()

    loss_per_iter = []
167
    for i in range(num_iterations):
Julien Chaumond's avatar
Julien Chaumond committed
168
        print("Iteration ", i + 1)
Piero Molino's avatar
Piero Molino committed
169
170
        past_perturb = [torch.from_numpy(p_) for p_ in past_perturb_orig]
        past_perturb = [to_var(p_, requires_grad=True) for p_ in past_perturb]
Julien Chaumond's avatar
Julien Chaumond committed
171

Piero Molino's avatar
Piero Molino committed
172
        perturbed_past = list(map(add, past, past_perturb))
Julien Chaumond's avatar
Julien Chaumond committed
173

Piero Molino's avatar
Piero Molino committed
174
        _, _, _, current_length, _ = past_perturb[0].shape
Julien Chaumond's avatar
Julien Chaumond committed
175

Piero Molino's avatar
Piero Molino committed
176
177
        # _, future_past = model(prev, past=perturbed_past)
        # hidden = model.hidden_states
Julien Chaumond's avatar
Julien Chaumond committed
178

Piero Molino's avatar
Piero Molino committed
179
180
181
182
183
184
185
186
187
188
189
        # Piero modified model call
        logits, _, all_hidden = model(prev, past=perturbed_past)
        hidden = all_hidden[-1]
        new_accumulated_hidden = accumulated_hidden + torch.sum(hidden,
                                                                dim=1).detach()

        # TODO: Check the layer-norm consistency of this with trained discriminator
        logits = logits[:, -1, :]
        probabs = F.softmax(logits, dim=-1)
        loss = 0.0
        loss_list = []
190
191
        if loss_type == 1 or loss_type == 3:
            for one_hot_good in one_hot_bows_vectors:
Piero Molino's avatar
Piero Molino committed
192
193
194
195
196
197
198
199
200
                good_logits = torch.mm(probabs, torch.t(one_hot_good))
                loss_word = good_logits
                loss_word = torch.sum(loss_word)
                loss_word = -torch.log(loss_word)
                # loss_word = torch.sum(loss_word) /torch.sum(one_hot_good)
                loss += loss_word
                loss_list.append(loss_word)
            print(" pplm_bow_loss:", loss.data.cpu().numpy())

201
        if loss_type == 2 or loss_type == 3:
Julien Chaumond's avatar
Julien Chaumond committed
202
            ce_loss = torch.nn.CrossEntropyLoss()
203
204
            new_true_past = unpert_past
            for i in range(horizon_length):
Piero Molino's avatar
Piero Molino committed
205
206
207
208
209
210
211
212
213
214
215
                future_probabs = F.softmax(logits, dim=-1)  # Get softmax
                future_probabs = torch.unsqueeze(future_probabs, dim=1)

                # _, new_true_past = model(future_probabs, past=new_true_past)
                # future_hidden = model.hidden_states  # Get expected hidden states

                # Piero modified model call
                wte = model.resize_token_embeddings()
                inputs_embeds = torch.matmul(future_probabs, wte.weight.data)
                _, new_true_past, future_hidden = model(
                    past=new_true_past,
Julien Chaumond's avatar
Julien Chaumond committed
216
217
                    inputs_embeds=inputs_embeds
                )
Piero Molino's avatar
Piero Molino committed
218
219
220
221
                future_hidden = future_hidden[-1]

                new_accumulated_hidden = new_accumulated_hidden + torch.sum(
                    future_hidden, dim=1)
Julien Chaumond's avatar
Julien Chaumond committed
222

Piero Molino's avatar
Piero Molino committed
223
            predicted_sentiment = classifier(new_accumulated_hidden / (
224
                    current_length + 1 + horizon_length))
Julien Chaumond's avatar
Julien Chaumond committed
225

226
            label = torch.tensor([label_class], device='cuda',
Piero Molino's avatar
Piero Molino committed
227
228
                                 dtype=torch.long)
            discrim_loss = ce_loss(predicted_sentiment, label)
Julien Chaumond's avatar
Julien Chaumond committed
229
            print(" pplm_discrim_loss:", discrim_loss.data.cpu().numpy())
Piero Molino's avatar
Piero Molino committed
230
231
            loss += discrim_loss
            loss_list.append(discrim_loss)
Julien Chaumond's avatar
Julien Chaumond committed
232

Piero Molino's avatar
Piero Molino committed
233
234
        kl_loss = 0.0
        if kl_scale > 0.0:
235
            p = (F.softmax(unpert_logits[:, -1, :], dim=-1))
Piero Molino's avatar
Piero Molino committed
236
237
238
239
240
            p = p + SmallConst * (p <= SmallConst).type(
                torch.FloatTensor).cuda().detach()
            correction = SmallConst * (probabs <= SmallConst).type(
                torch.FloatTensor).cuda().detach()
            corrected_probabs = probabs + correction.detach()
Rosanne Liu's avatar
Rosanne Liu committed
241
            kl_loss = kl_scale * (
Piero Molino's avatar
Piero Molino committed
242
                (corrected_probabs * (corrected_probabs / p).log()).sum())
Julien Chaumond's avatar
Julien Chaumond committed
243
            print(' kl_loss', (kl_loss).data.cpu().numpy())
Piero Molino's avatar
Piero Molino committed
244
            loss += kl_loss  # + discrim_loss
Julien Chaumond's avatar
Julien Chaumond committed
245
246

        loss_per_iter.append(loss.data.cpu().numpy())
Piero Molino's avatar
Piero Molino committed
247

Julien Chaumond's avatar
Julien Chaumond committed
248
249
        print(' pplm_loss', (loss - kl_loss).data.cpu().numpy())

Rosanne Liu's avatar
Rosanne Liu committed
250
        loss.backward()
251
        if grad_norms is not None and loss_type == 1:
Julien Chaumond's avatar
Julien Chaumond committed
252
253
            grad_norms = [
                torch.max(grad_norms[index], torch.norm(p_.grad * window_mask))
Piero Molino's avatar
Piero Molino committed
254
255
                for index, p_ in
                enumerate(past_perturb)]
Julien Chaumond's avatar
Julien Chaumond committed
256
        else:
Piero Molino's avatar
Piero Molino committed
257
258
            grad_norms = [(torch.norm(p_.grad * window_mask) + SmallConst) for
                          index, p_ in enumerate(past_perturb)]
Julien Chaumond's avatar
Julien Chaumond committed
259
260

        grad = [
Piero Molino's avatar
Piero Molino committed
261
            -stepsize * (p_.grad * window_mask / grad_norms[
262
                index] ** gamma).data.cpu().numpy()
Piero Molino's avatar
Piero Molino committed
263
264
            for index, p_ in enumerate(past_perturb)]
        past_perturb_orig = list(map(add, grad, past_perturb_orig))
Julien Chaumond's avatar
Julien Chaumond committed
265

Piero Molino's avatar
Piero Molino committed
266
        for p_ in past_perturb:
Julien Chaumond's avatar
Julien Chaumond committed
267
268
269
            p_.grad.data.zero_()

        new_past = []
Piero Molino's avatar
Piero Molino committed
270
271
272
        for p in past:
            new_past.append(p.detach())

Julien Chaumond's avatar
Julien Chaumond committed
273
274
        past = new_past

Piero Molino's avatar
Piero Molino committed
275
276
277
    past_perturb = [torch.from_numpy(p_) for p_ in past_perturb_orig]
    past_perturb = [to_var(p_, requires_grad=True) for p_ in past_perturb]
    perturbed_past = list(map(add, past, past_perturb))
Julien Chaumond's avatar
Julien Chaumond committed
278

Piero Molino's avatar
Piero Molino committed
279
    return perturbed_past, new_accumulated_hidden, grad_norms, loss_per_iter
Julien Chaumond's avatar
Julien Chaumond committed
280
281
282


def get_classifier(
Piero Molino's avatar
Piero Molino committed
283
284
        name: Optional[str], label_class: Union[str, int],
        device: Union[str, torch.device]
Julien Chaumond's avatar
Julien Chaumond committed
285
286
287
288
289
290
291
292
293
294
) -> Tuple[Optional[ClassificationHead], Optional[int]]:
    if name is None:
        return None, None

    params = DISCRIMINATOR_MODELS_PARAMS[name]
    classifier = ClassificationHead(
        class_size=params['class_size'],
        embed_size=params['embed_size']
    ).to(device)
    resolved_archive_file = cached_path(params["url"])
Piero Molino's avatar
Piero Molino committed
295
296
    classifier.load_state_dict(
        torch.load(resolved_archive_file, map_location=device))
Julien Chaumond's avatar
Julien Chaumond committed
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
    classifier.eval()

    if isinstance(label_class, str):
        if label_class in params["class_vocab"]:
            label_id = params["class_vocab"][label_class]
        else:
            label_id = params["default_class"]
            print("label_class {} not in class_vocab".format(label_class))
            print("available values are: {}".format(params["class_vocab"]))
            print("using default class {}".format(label_id))

    elif isinstance(label_class, int):
        if label_class in set(params["class_vocab"].values()):
            label_id = label_class
        else:
            label_id = params["default_class"]
            print("label_class {} not in class_vocab".format(label_class))
            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


Piero Molino's avatar
Piero Molino committed
323
324
def get_bag_of_words_indices(bag_of_words_ids_or_paths: List[str]) -> List[
    List[List[int]]]:
Julien Chaumond's avatar
Julien Chaumond committed
325
326
327
328
329
330
331
    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
332
333
334
335
            words = f.read().strip().split("\n")
        bow_indices.append(
            [TOKENIZER.encode(word.strip(), add_prefix_space=True) for word in
             words])
Julien Chaumond's avatar
Julien Chaumond committed
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
    return bow_indices


def build_bows_one_hot_vectors(bow_indices):
    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))
        single_bow = torch.tensor(single_bow).cuda()
        num_words = single_bow.shape[0]
        one_hot_bow = torch.zeros(num_words, TOKENIZER.vocab_size).cuda()
        one_hot_bow.scatter_(1, single_bow, 1)
        one_hot_bows_vectors.append(one_hot_bow)
    return one_hot_bows_vectors


354
def full_text_generation(
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
        model,
        context=None,
        num_samples=1,
        device="cuda",
        sample=True,
        discrim=None,
        label_class=None,
        bag_of_words=None,
        length=100,
        grad_length=10000,
        stepsize=0.02,
        num_iterations=3,
        temperature=1.0,
        gm_scale=0.9,
        kl_scale=0.01,
        top_k=10,
        window_length=0,
        horizon_length=1,
        decay=False,
        gamma=1.5,
        **kwargs
):
Julien Chaumond's avatar
Julien Chaumond committed
377
    classifier, class_id = get_classifier(
378
379
        discrim,
        label_class,
Julien Chaumond's avatar
Julien Chaumond committed
380
381
382
        device
    )

383
384
385
    bow_indices = []
    if bag_of_words:
        bow_indices = get_bag_of_words_indices(bag_of_words.split(";"))
Piero Molino's avatar
Piero Molino committed
386

387
    if bag_of_words and classifier:
Julien Chaumond's avatar
Julien Chaumond committed
388
        print("Both PPLM-BoW and PPLM-Discrim are on. This is not optimized.")
389
        loss_type = PPLM_BOW_DISCRIM
Julien Chaumond's avatar
Julien Chaumond committed
390

391
392
    elif bag_of_words:
        loss_type = PPLM_BOW
Julien Chaumond's avatar
Julien Chaumond committed
393
394
395
        print("Using PPLM-BoW")

    elif classifier is not None:
396
        loss_type = PPLM_DISCRIM
Julien Chaumond's avatar
Julien Chaumond committed
397
398
399
        print("Using PPLM-Discrim")

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

402
    unpert_gen_tok_text, _, _ = generate_text_pplm(
403
404
405
406
407
408
        model=model,
        context=context,
        device=device,
        length=length,
        perturb=False
    )
Julien Chaumond's avatar
Julien Chaumond committed
409
410
    torch.cuda.empty_cache()

411
412
413
    pert_gen_tok_texts = []
    discrim_losses = []
    losses_in_time = []
Piero Molino's avatar
Piero Molino committed
414

415
    for i in range(num_samples):
416
        pert_gen_tok_text, discrim_loss, loss_in_time = 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
            model=model,
            context=context,
            device=device,
            sample=sample,
            perturb=True,
            bow_indices=bow_indices,
            classifier=classifier,
            label_class=class_id,
            loss_type=loss_type,
            length=length,
            grad_length=grad_length,
            stepsize=stepsize,
            num_iterations=num_iterations,
            temperature=temperature,
            gm_scale=gm_scale,
            kl_scale=kl_scale,
            top_k=top_k,
            window_length=window_length,
            horizon_length=horizon_length,
            decay=decay,
            gamma=gamma,
        )
439
        pert_gen_tok_texts.append(pert_gen_tok_text)
Julien Chaumond's avatar
Julien Chaumond committed
440
        if classifier is not None:
441
442
            discrim_losses.append(discrim_loss.data.cpu().numpy())
        losses_in_time.append(loss_in_time)
Julien Chaumond's avatar
Julien Chaumond committed
443
444
445

    torch.cuda.empty_cache()

446
    return unpert_gen_tok_text, pert_gen_tok_texts, discrim_losses, losses_in_time
Julien Chaumond's avatar
Julien Chaumond committed
447

448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472

def generate_text_pplm(
        model,
        context=None,
        past=None,
        device="cuda",
        sample=True,
        perturb=True,
        classifier=None,
        label_class=None,
        bow_indices=None,
        loss_type=0,
        length=100,
        grad_length=10000,
        stepsize=0.02,
        num_iterations=3,
        temperature=1.0,
        gm_scale=0.9,
        kl_scale=0.01,
        top_k=10,
        window_length=0,
        horizon_length=1,
        decay=False,
        gamma=1.5,
):
Piero Molino's avatar
Piero Molino committed
473
474
    output = torch.tensor(context, device=device, dtype=torch.long).unsqueeze(
        0) if context else None
Julien Chaumond's avatar
Julien Chaumond committed
475

476
477
478
    # collect one hot vectors for bags of words
    one_hot_bows_vectors = build_bows_one_hot_vectors(bow_indices)

Julien Chaumond's avatar
Julien Chaumond committed
479
480
    grad_norms = None
    loss_in_time = []
481
    for i in trange(length, ascii=True):
Julien Chaumond's avatar
Julien Chaumond committed
482
483

        # Get past/probs for current output, except for last word
Piero Molino's avatar
Piero Molino committed
484
485
        # Note that GPT takes 2 inputs: past + current-token
        # Therefore, use everything from before current i/p token to generate relevant past
Julien Chaumond's avatar
Julien Chaumond committed
486

Piero Molino's avatar
Piero Molino committed
487
488
489
490
491
        if past is None and output is not None:
            prev = output[:, -1:]
            # _, past = model(output[:, :-1])
            # original_probs, true_past = model(output)
            # true_hidden = model.hidden_states
Julien Chaumond's avatar
Julien Chaumond committed
492

Piero Molino's avatar
Piero Molino committed
493
494
            # Piero modified model call
            _, past, _ = model(output[:, :-1])
495
            unpert_logits, unpert_past, unpert_all_hidden = model(output)
Piero Molino's avatar
Piero Molino committed
496
            true_hidden = unpert_all_hidden[-1]
Julien Chaumond's avatar
Julien Chaumond committed
497
498

        else:
Piero Molino's avatar
Piero Molino committed
499
500
            # original_probs, true_past = model(output)
            # true_hidden = model.hidden_states
Julien Chaumond's avatar
Julien Chaumond committed
501

Piero Molino's avatar
Piero Molino committed
502
            # Piero modified model call
503
            unpert_logits, unpert_past, unpert_all_hidden = model(output)
Piero Molino's avatar
Piero Molino committed
504
505
506
507
            true_hidden = unpert_all_hidden[-1]

        # Modify the past if necessary

508
509
        if i >= grad_length:
            current_stepsize = stepsize * 0
Julien Chaumond's avatar
Julien Chaumond committed
510
        else:
511
            current_stepsize = stepsize
Julien Chaumond's avatar
Julien Chaumond committed
512

513
        if not perturb or num_iterations == 0:
Piero Molino's avatar
Piero Molino committed
514
            perturbed_past = past
Julien Chaumond's avatar
Julien Chaumond committed
515
516

        else:
Piero Molino's avatar
Piero Molino committed
517
518
519
            # Piero modified model call
            # accumulated_hidden = model.hidden_states[:, :-1, :]
            accumulated_hidden = true_hidden[:, :-1, :]
Julien Chaumond's avatar
Julien Chaumond committed
520
521
            accumulated_hidden = torch.sum(accumulated_hidden, dim=1)

522
            perturbed_past, _, grad_norms, loss_per_iter = perturb_past(
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
                past,
                model,
                prev,
                unpert_past=unpert_past,
                unpert_logits=unpert_logits,
                accumulated_hidden=accumulated_hidden,
                grad_norms=grad_norms,
                stepsize=current_stepsize,
                classifier=classifier,
                label_class=label_class,
                one_hot_bows_vectors=one_hot_bows_vectors,
                loss_type=loss_type,
                num_iterations=num_iterations,
                kl_scale=kl_scale,
                window_length=window_length,
                horizon_length=horizon_length,
                decay=decay,
                gamma=gamma,
            )
Piero Molino's avatar
Piero Molino committed
542
543
544
545
546
547
548
            loss_in_time.append(loss_per_iter)

        # Piero modified model call
        logits, past, pert_all_hidden = model(prev, past=perturbed_past)
        # test_logits = F.softmax(test_logits[:, -1, :], dim=-1)
        # likelywords = torch.topk(test_logits, k=10, dim=-1)
        # print(TOKENIZER.decode(likelywords[1].tolist()[0]))
Julien Chaumond's avatar
Julien Chaumond committed
549
550

        if classifier is not None:
Piero Molino's avatar
Piero Molino committed
551
552
            ce_loss = torch.nn.CrossEntropyLoss()
            predicted_sentiment = classifier(torch.mean(true_hidden, dim=1))
553
            label = torch.tensor([label_class], device='cuda',
Piero Molino's avatar
Piero Molino committed
554
555
556
                                 dtype=torch.long)
            true_discrim_loss = ce_loss(predicted_sentiment, label)
            print("true discrim loss", true_discrim_loss.data.cpu().numpy())
Julien Chaumond's avatar
Julien Chaumond committed
557
        else:
Piero Molino's avatar
Piero Molino committed
558
559
560
561
562
            true_discrim_loss = 0

        # Piero modified model call
        # hidden = model.hidden_states  # update hidden
        # logits = model.forward_hidden(hidden)
563
        logits = logits[:, -1, :] / temperature  # + SmallConst
Piero Molino's avatar
Piero Molino committed
564
565

        # logits = top_k_filter(logits, k=args.top_k)  # + SmallConst
Julien Chaumond's avatar
Julien Chaumond committed
566

Piero Molino's avatar
Piero Molino committed
567
568
569
        log_probs = F.softmax(logits, dim=-1)

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

Piero Molino's avatar
Piero Molino committed
572
            # original_probs = top_k_filter(original_probs[:, -1, :]) #+ SmallConst
573
            unpert_logits = F.softmax(unpert_logits[:, -1, :], dim=-1)
Piero Molino's avatar
Piero Molino committed
574
575
576
577
            # likelywords = torch.topk(original_probs, k=10, dim=-1)
            # print(TOKENIZER.decode(likelywords[1].tolist()[0]))

            log_probs = ((log_probs ** gm_scale) * (
578
                    unpert_logits ** (1 - gm_scale)))  # + SmallConst
Julien Chaumond's avatar
Julien Chaumond committed
579

580
            log_probs = top_k_filter(log_probs, k=top_k,
Piero Molino's avatar
Piero Molino committed
581
                                     probs=True)  # + SmallConst
Julien Chaumond's avatar
Julien Chaumond committed
582

Piero Molino's avatar
Piero Molino committed
583
584
            if torch.sum(log_probs) <= 1:
                log_probs = log_probs / torch.sum(log_probs)
Julien Chaumond's avatar
Julien Chaumond committed
585
586

        else:
587
            logits = top_k_filter(logits, k=top_k)  # + SmallConst
Piero Molino's avatar
Piero Molino committed
588
            log_probs = F.softmax(logits, dim=-1)
Julien Chaumond's avatar
Julien Chaumond committed
589
590

        if sample:
Piero Molino's avatar
Piero Molino committed
591
592
593
594
            # likelywords = torch.topk(log_probs, k=args.top_k, dim=-1)
            # print(TOKENIZER.decode(likelywords[1].tolist()[0]))
            # print(likelywords[0].tolist())
            prev = torch.multinomial(log_probs, num_samples=1)
Julien Chaumond's avatar
Julien Chaumond committed
595
        else:
Piero Molino's avatar
Piero Molino committed
596
597
598
599
600
601
            _, prev = torch.topk(log_probs, k=1, dim=-1)
        # if perturb:
        #     prev = future
        output = prev if output is None else torch.cat((output, prev),
                                                       dim=1)  # update output
        print(TOKENIZER.decode(output.tolist()[0]))
Julien Chaumond's avatar
Julien Chaumond committed
602

Piero Molino's avatar
Piero Molino committed
603
    return output, true_discrim_loss, loss_in_time
Julien Chaumond's avatar
Julien Chaumond committed
604
605
606
607


def run_model():
    parser = argparse.ArgumentParser()
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
    parser.add_argument(
        "--model_path",
        "-M",
        type=str,
        default="gpt2-medium",
        help="pretrained model name or path to local checkpoint",
    )
    parser.add_argument(
        "--bag_of_words",
        "-B",
        type=str,
        default=None,
        help="Bags of words used for PPLM-BoW. "
             "Either a BOW id (see list in code) or a filepath. "
             "Multiple BoWs separated by ;",
    )
    parser.add_argument(
        "--discrim",
        "-D",
        type=str,
        default=None,
        choices=("clickbait", "sentiment", "toxicity"),
        help="Discriminator to use for loss-type 2",
    )
    parser.add_argument(
        "--label_class",
        type=int,
        default=-1,
        help="Class label used for the discriminator",
    )
    parser.add_argument("--stepsize", type=float, default=0.02)
Julien Chaumond's avatar
Julien Chaumond committed
639
640
641
642
643
644
    parser.add_argument("--length", type=int, default=100)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument("--temperature", type=float, default=1.0)
    parser.add_argument("--top_k", type=int, default=10)
    parser.add_argument("--gm_scale", type=float, default=0.9)
    parser.add_argument("--kl_scale", type=float, default=0.01)
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
    parser.add_argument("--no_cuda", action="store_true", help="no cuda")
    parser.add_argument(
        "--uncond", action="store_true",
        help="Generate from end-of-text as prefix"
    )
    parser.add_argument(
        "--cond_text", type=str, default="The lake",
        help="Prefix texts to condition on"
    )
    parser.add_argument("--num_iterations", type=int, default=3)
    parser.add_argument("--grad_length", type=int, default=10000)
    parser.add_argument(
        "--num_samples",
        type=int,
        default=1,
        help="Number of samples to generate from the modified latents",
    )
    parser.add_argument(
        "--horizon_length",
        type=int,
        default=1,
        help="Length of future to optimize over",
    )
    parser.add_argument(
        "--window_length",
        type=int,
        default=0,
        help="Length of past which is being optimized; "
             "0 corresponds to infinite window length",
    )
    parser.add_argument("--decay", action="store_true",
                        help="whether to decay or not")
    parser.add_argument("--gamma", type=float, default=1.5)
    parser.add_argument("--colorama", action="store_true", help="colors keywords")
Julien Chaumond's avatar
Julien Chaumond committed
679
680
681

    args = parser.parse_args()

682
    # set Random seed
Julien Chaumond's avatar
Julien Chaumond committed
683
684
685
    torch.manual_seed(args.seed)
    np.random.seed(args.seed)

686
687
    # set the device
    device = "cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu"
Julien Chaumond's avatar
Julien Chaumond committed
688

689
    # load pretrained model
Julien Chaumond's avatar
Julien Chaumond committed
690
691
692
693
694
695
696
    model = GPT2LMHeadModel.from_pretrained(
        args.model_path,
        output_hidden_states=True
    )
    model.to(device)
    model.eval()

Piero Molino's avatar
Piero Molino committed
697
    # Freeze GPT-2 weights
Julien Chaumond's avatar
Julien Chaumond committed
698
699
700
    for param in model.parameters():
        param.requires_grad = False

701
    # figure out conditioning text
Julien Chaumond's avatar
Julien Chaumond committed
702
    if args.uncond:
703
704
705
        tokenized_cond_text = TOKENIZER.encode(
            [TOKENIZER.bos_token]
        )
Julien Chaumond's avatar
Julien Chaumond committed
706
707
708
    else:
        raw_text = args.cond_text
        while not raw_text:
709
            print("Did you forget to add `--cond_text`? ")
Julien Chaumond's avatar
Julien Chaumond committed
710
            raw_text = input("Model prompt >>> ")
711
        tokenized_cond_text = TOKENIZER.encode(TOKENIZER.bos_token + raw_text)
Piero Molino's avatar
Piero Molino committed
712

713
714
715
    print("= Prefix of sentence =")
    print(TOKENIZER.decode(tokenized_cond_text))
    print()
Piero Molino's avatar
Piero Molino committed
716

717
    # generate unperturbed and perturbed texts
Piero Molino's avatar
Piero Molino committed
718

719
720
721
722
723
724
725
726
    # 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(
        model=model, context=tokenized_cond_text, device=device, **vars(args)
    )

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

728
729
730
731
    print("=" * 80)
    print("= Unperturbed generated text =")
    print(unpert_gen_text)
    print()
Piero Molino's avatar
Piero Molino committed
732

733
734
    generated_texts = []

735
736
737
738
739
740
741
742
    bow_word_ids = set()
    if args.bag_of_words and args.colorama:
        bow_indices = get_bag_of_words_indices(args.bag_of_words.split(";"))
        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)
743
744
745
746
747

    # iterate through the perturbed texts
    for i, pert_gen_tok_text in enumerate(pert_gen_tok_texts):
        try:
            # untokenize unperturbed text
Piero Molino's avatar
Piero Molino committed
748
749
750
            if args.colorama:
                import colorama

751
752
                pert_gen_text = ''
                for word_id in pert_gen_tok_text.tolist()[0]:
753
                    if word_id in bow_word_ids:
754
755
756
757
758
                        pert_gen_text += '{}{}{}'.format(
                            colorama.Fore.RED,
                            TOKENIZER.decode([word_id]),
                            colorama.Style.RESET_ALL
                        )
Piero Molino's avatar
Piero Molino committed
759
                    else:
760
                        pert_gen_text += TOKENIZER.decode([word_id])
Piero Molino's avatar
Piero Molino committed
761
            else:
762
                pert_gen_text = TOKENIZER.decode(pert_gen_tok_text.tolist()[0])
Julien Chaumond's avatar
Julien Chaumond committed
763

764
765
766
767
768
            print("= Perturbed generated text {} =".format(i + 1))
            print(pert_gen_text)
            print()
        except:
            pass
Julien Chaumond's avatar
Julien Chaumond committed
769

770
771
772
773
        # keep the prefix, perturbed seq, original seq for each index
        generated_texts.append(
            (tokenized_cond_text, pert_gen_tok_text, unpert_gen_tok_text)
        )
Julien Chaumond's avatar
Julien Chaumond committed
774

Piero Molino's avatar
Piero Molino committed
775
    return
Julien Chaumond's avatar
Julien Chaumond committed
776
777


Piero Molino's avatar
Piero Molino committed
778
if __name__ == '__main__':
Julien Chaumond's avatar
Julien Chaumond committed
779
    run_model()