run_pplm.py 30.6 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
from transformers import GPT2Tokenizer
from transformers.file_utils import cached_path
from transformers.modeling_gpt2 import GPT2LMHeadModel

Piero Molino's avatar
Piero Molino committed
42

Julien Chaumond's avatar
Julien Chaumond committed
43
44
45
46
PPLM_BOW = 1
PPLM_DISCRIM = 2
PPLM_BOW_DISCRIM = 3
SMALL_CONST = 1e-15
Piero Molino's avatar
Piero Molino committed
47
SmallConst = 1e-15
Julien Chaumond's avatar
Julien Chaumond committed
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
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
72
        "url": "http://s.yosinski.com/SST_classifier_head.pt",
Julien Chaumond's avatar
Julien Chaumond committed
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
        "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
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
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)


112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
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,
):

    #def perturb_past(past, model, prev, classifier, good_index=None,
    #             stepsize=0.01, vocab_size=50257,
    #             original_probs=None, accumulated_hidden=None, true_past=None,
    #             grad_norms=None):

    # one_hot_bows_vectors = []
    # for good_list in good_index:
    #     good_list = list(filter(lambda x: len(x) <= 1, good_list))
    #     good_list = torch.tensor(good_list).cuda()
    #     num_good = good_list.shape[0]
    #     one_hot_good = torch.zeros(num_good, vocab_size).cuda()
    #     one_hot_good.scatter_(1, good_list, 1)
    #     one_hot_bows_vectors.append(one_hot_good)
Piero Molino's avatar
Piero Molino committed
146
147
148
149
150

    # 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
151
152
153
154

    if accumulated_hidden is None:
        accumulated_hidden = 0

155
    if decay:
Piero Molino's avatar
Piero Molino committed
156
157
        decay_mask = torch.arange(0., 1.0 + SmallConst, 1.0 / (window_length))[
                     1:]
Julien Chaumond's avatar
Julien Chaumond committed
158
159
160
    else:
        decay_mask = 1.0

Piero Molino's avatar
Piero Molino committed
161
162
163
164
165
166
167
168
169
170
171
    # 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
172
173
174
175
176

        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
177
178
        window_mask = torch.cat((ones_mask, torch.zeros(zeros_key_val_shape)),
                                dim=-2).cuda()
Julien Chaumond's avatar
Julien Chaumond committed
179
180
181
182
    else:
        window_mask = torch.ones_like(past[0]).cuda()

    loss_per_iter = []
183
    for i in range(num_iterations):
Julien Chaumond's avatar
Julien Chaumond committed
184
        print("Iteration ", i + 1)
Piero Molino's avatar
Piero Molino committed
185
186
        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
187

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

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

Piero Molino's avatar
Piero Molino committed
192
193
        # _, future_past = model(prev, past=perturbed_past)
        # hidden = model.hidden_states
Julien Chaumond's avatar
Julien Chaumond committed
194

Piero Molino's avatar
Piero Molino committed
195
196
197
198
199
200
201
202
203
204
205
        # 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 = []
206
207
        if loss_type == 1 or loss_type == 3:
            for one_hot_good in one_hot_bows_vectors:
Piero Molino's avatar
Piero Molino committed
208
209
210
211
212
213
214
215
216
                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())

217
        if loss_type == 2 or loss_type == 3:
Julien Chaumond's avatar
Julien Chaumond committed
218
            ce_loss = torch.nn.CrossEntropyLoss()
219
220
            new_true_past = unpert_past
            for i in range(horizon_length):
Piero Molino's avatar
Piero Molino committed
221
222
223
224
225
226
227
228
229
230
231
                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
232
233
                    inputs_embeds=inputs_embeds
                )
Piero Molino's avatar
Piero Molino committed
234
235
236
237
                future_hidden = future_hidden[-1]

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

Piero Molino's avatar
Piero Molino committed
239
            predicted_sentiment = classifier(new_accumulated_hidden / (
240
                        current_length + 1 + horizon_length))
Julien Chaumond's avatar
Julien Chaumond committed
241

242
            label = torch.tensor([label_class], device='cuda',
Piero Molino's avatar
Piero Molino committed
243
244
                                 dtype=torch.long)
            discrim_loss = ce_loss(predicted_sentiment, label)
Julien Chaumond's avatar
Julien Chaumond committed
245
            print(" pplm_discrim_loss:", discrim_loss.data.cpu().numpy())
Piero Molino's avatar
Piero Molino committed
246
247
            loss += discrim_loss
            loss_list.append(discrim_loss)
Julien Chaumond's avatar
Julien Chaumond committed
248

Piero Molino's avatar
Piero Molino committed
249
250
        kl_loss = 0.0
        if kl_scale > 0.0:
251
            p = (F.softmax(unpert_logits[:, -1, :], dim=-1))
Piero Molino's avatar
Piero Molino committed
252
253
254
255
256
            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
257
            kl_loss = kl_scale * (
Piero Molino's avatar
Piero Molino committed
258
                (corrected_probabs * (corrected_probabs / p).log()).sum())
Julien Chaumond's avatar
Julien Chaumond committed
259
            print(' kl_loss', (kl_loss).data.cpu().numpy())
Piero Molino's avatar
Piero Molino committed
260
            loss += kl_loss  # + discrim_loss
Julien Chaumond's avatar
Julien Chaumond committed
261
262

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

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

Rosanne Liu's avatar
Rosanne Liu committed
266
        loss.backward()
267
        if grad_norms is not None and loss_type == 1:
Julien Chaumond's avatar
Julien Chaumond committed
268
269
            grad_norms = [
                torch.max(grad_norms[index], torch.norm(p_.grad * window_mask))
Piero Molino's avatar
Piero Molino committed
270
271
                for index, p_ in
                enumerate(past_perturb)]
Julien Chaumond's avatar
Julien Chaumond committed
272
        else:
Piero Molino's avatar
Piero Molino committed
273
274
            grad_norms = [(torch.norm(p_.grad * window_mask) + SmallConst) for
                          index, p_ in enumerate(past_perturb)]
Julien Chaumond's avatar
Julien Chaumond committed
275
276

        grad = [
Piero Molino's avatar
Piero Molino committed
277
            -stepsize * (p_.grad * window_mask / grad_norms[
278
                index] ** gamma).data.cpu().numpy()
Piero Molino's avatar
Piero Molino committed
279
280
            for index, p_ in enumerate(past_perturb)]
        past_perturb_orig = list(map(add, grad, past_perturb_orig))
Julien Chaumond's avatar
Julien Chaumond committed
281

Piero Molino's avatar
Piero Molino committed
282
        for p_ in past_perturb:
Julien Chaumond's avatar
Julien Chaumond committed
283
284
285
            p_.grad.data.zero_()

        new_past = []
Piero Molino's avatar
Piero Molino committed
286
287
288
        for p in past:
            new_past.append(p.detach())

Julien Chaumond's avatar
Julien Chaumond committed
289
290
        past = new_past

Piero Molino's avatar
Piero Molino committed
291
292
293
    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
294

Piero Molino's avatar
Piero Molino committed
295
    return perturbed_past, new_accumulated_hidden, grad_norms, loss_per_iter
Julien Chaumond's avatar
Julien Chaumond committed
296
297
298


def get_classifier(
Piero Molino's avatar
Piero Molino committed
299
300
        name: Optional[str], label_class: Union[str, int],
        device: Union[str, torch.device]
Julien Chaumond's avatar
Julien Chaumond committed
301
302
303
304
305
306
307
308
309
310
) -> 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
311
312
    classifier.load_state_dict(
        torch.load(resolved_archive_file, map_location=device))
Julien Chaumond's avatar
Julien Chaumond committed
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
    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
339
340
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
341
342
343
344
345
346
347
    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
348
349
350
351
            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
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
    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


370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
def full_text_generation(
            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
393
    classifier, class_id = get_classifier(
394
395
        discrim,
        label_class,
Julien Chaumond's avatar
Julien Chaumond committed
396
397
398
        device
    )

Piero Molino's avatar
Piero Molino committed
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
    # if args.discrim == 'clickbait':
    #     classifier = ClassificationHead(class_size=2, embed_size=1024).to(device)
    #     classifier.load_state_dict(torch.load("discrim_models/clickbait_classifierhead.pt"))
    #     classifier.eval()
    #     args.label_class = 1 # clickbaity
    #
    # elif args.discrim == 'sentiment':
    #     classifier = ClassificationHead(class_size=5, embed_size=1024).to(device)
    #     #classifier.load_state_dict(torch.load("discrim_models/sentiment_classifierhead.pt"))
    #     classifier.load_state_dict(torch.load("discrim_models/SST_classifier_head_epoch_16.pt"))
    #     classifier.eval()
    #     if args.label_class < 0:
    #         raise Exception('Wrong class for sentiment, use --label-class 2 for *very positive*, 3 for *very negative*')
    #     #args.label_class = 2 # very pos
    #     #args.label_class = 3 # very neg
    #
    # elif args.discrim == 'toxicity':
    #     classifier = ClassificationHead(class_size=2, embed_size=1024).to(device)
    #     classifier.load_state_dict(torch.load("discrim_models/toxicity_classifierhead.pt"))
    #     classifier.eval()
    #     args.label_class = 0 # not toxic
    #
    # elif args.discrim == 'generic':
    #     if args.discrim_weights is None:
    #         raise ValueError('When using a generic discriminator, '
    #                          'discrim_weights need to be specified')
    #     if args.discrim_meta is None:
    #         raise ValueError('When using a generic discriminator, '
    #                          'discrim_meta need to be specified')
    #
    #     with open(args.discrim_meta, 'r') as discrim_meta_file:
    #         meta = json.load(discrim_meta_file)
    #
    #     classifier = ClassificationHead(
    #         class_size=meta['class_size'],
    #         embed_size=meta['embed_size'],
    #         # todo add tokenizer from meta
    #     ).to(device)
    #     classifier.load_state_dict(torch.load(args.discrim_weights))
    #     classifier.eval()
    #     if args.label_class == -1:
    #         args.label_class = meta['default_class']
    #
    # else:
    #     classifier = None

    # Get tokens for the list of positive words
    def list_tokens(word_list):
        token_list = [TOKENIZER.encode(word, add_prefix_space=True) for word in
                      word_list]
        # token_list = []
        # for word in word_list:
        #    token_list.append(TOKENIZER.encode(" " + word))
        return token_list

    # good_index = []
    # if args.bag_of_words:
    #     bags_of_words = args.bag_of_words.split(";")
    #     for wordlist in bags_of_words:
    #         with open(wordlist, "r") as f:
    #             words = f.read().strip()
    #             words = words.split('\n')
    #         good_index.append(list_tokens(words))
    #
    #     for good_list in good_index:
    #         good_list = list(filter(lambda x: len(x) <= 1, good_list))
    #         actual_words = [(TOKENIZER.decode(ww).strip(),ww) for ww in good_list]

467
    bow_indices = []
Piero Molino's avatar
Piero Molino committed
468
    actual_words = None
469
470
    if bag_of_words:
        bow_indices = get_bag_of_words_indices(bag_of_words.split(";"))
Piero Molino's avatar
Piero Molino committed
471

472
        for good_list in bow_indices:
Piero Molino's avatar
Piero Molino committed
473
474
475
476
            good_list = list(filter(lambda x: len(x) <= 1, good_list))
            actual_words = [(TOKENIZER.decode(ww).strip(), ww) for ww in
                            good_list]

477
    if bag_of_words and classifier:
Julien Chaumond's avatar
Julien Chaumond committed
478
        print("Both PPLM-BoW and PPLM-Discrim are on. This is not optimized.")
479
        loss_type = PPLM_BOW_DISCRIM
Julien Chaumond's avatar
Julien Chaumond committed
480

481
482
    elif bag_of_words:
        loss_type = PPLM_BOW
Julien Chaumond's avatar
Julien Chaumond committed
483
484
485
        print("Using PPLM-BoW")

    elif classifier is not None:
486
        loss_type = PPLM_DISCRIM
Julien Chaumond's avatar
Julien Chaumond committed
487
488
489
490
491
        print("Using PPLM-Discrim")

    else:
        raise Exception("Specify either --bag_of_words (-B) or --discrim (-D)")

492
493
494
495
496
497
498
    original, _, _ = generate_text_pplm(
        model=model,
        context=context,
        device=device,
        length=length,
        perturb=False
    )
Julien Chaumond's avatar
Julien Chaumond committed
499
500
    torch.cuda.empty_cache()

Piero Molino's avatar
Piero Molino committed
501
502
503
504
    perturbed_list = []
    discrim_loss_list = []
    loss_in_time_list = []

505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
    for i in range(num_samples):
        perturbed, discrim_loss, loss_in_time = generate_text_pplm(
            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,
        )
Piero Molino's avatar
Piero Molino committed
529
        perturbed_list.append(perturbed)
Julien Chaumond's avatar
Julien Chaumond committed
530
        if classifier is not None:
Piero Molino's avatar
Piero Molino committed
531
532
            discrim_loss_list.append(discrim_loss.data.cpu().numpy())
        loss_in_time_list.append(loss_in_time)
Julien Chaumond's avatar
Julien Chaumond committed
533
534
535

    torch.cuda.empty_cache()

Piero Molino's avatar
Piero Molino committed
536
537
    return original, perturbed_list, discrim_loss_list, loss_in_time_list, actual_words

Julien Chaumond's avatar
Julien Chaumond committed
538

539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563

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
564
565
    output = torch.tensor(context, device=device, dtype=torch.long).unsqueeze(
        0) if context else None
Julien Chaumond's avatar
Julien Chaumond committed
566

567
568
569
    # 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
570
571
    grad_norms = None
    loss_in_time = []
572
    for i in trange(length, ascii=True):
Julien Chaumond's avatar
Julien Chaumond committed
573
574

        # Get past/probs for current output, except for last word
Piero Molino's avatar
Piero Molino committed
575
576
        # 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
577

Piero Molino's avatar
Piero Molino committed
578
579
580
581
582
        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
583

Piero Molino's avatar
Piero Molino committed
584
585
            # Piero modified model call
            _, past, _ = model(output[:, :-1])
586
            unpert_logits, unpert_past, unpert_all_hidden = model(output)
Piero Molino's avatar
Piero Molino committed
587
            true_hidden = unpert_all_hidden[-1]
Julien Chaumond's avatar
Julien Chaumond committed
588
589

        else:
Piero Molino's avatar
Piero Molino committed
590
591
            # original_probs, true_past = model(output)
            # true_hidden = model.hidden_states
Julien Chaumond's avatar
Julien Chaumond committed
592

Piero Molino's avatar
Piero Molino committed
593
            # Piero modified model call
594
            unpert_logits, unpert_past, unpert_all_hidden = model(output)
Piero Molino's avatar
Piero Molino committed
595
596
597
598
            true_hidden = unpert_all_hidden[-1]

        # Modify the past if necessary

599
600
        if i >= grad_length:
            current_stepsize = stepsize * 0
Julien Chaumond's avatar
Julien Chaumond committed
601
        else:
602
            current_stepsize = stepsize
Julien Chaumond's avatar
Julien Chaumond committed
603

604
        if not perturb or num_iterations == 0:
Piero Molino's avatar
Piero Molino committed
605
            perturbed_past = past
Julien Chaumond's avatar
Julien Chaumond committed
606
607

        else:
Piero Molino's avatar
Piero Molino committed
608
609
610
            # Piero modified model call
            # accumulated_hidden = model.hidden_states[:, :-1, :]
            accumulated_hidden = true_hidden[:, :-1, :]
Julien Chaumond's avatar
Julien Chaumond committed
611
612
            accumulated_hidden = torch.sum(accumulated_hidden, dim=1)

613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
            perturbed_past, _, grad_norms, loss_per_iter = perturb_past(
                    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
633
634
635
636
637
638
639
            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
640
641

        if classifier is not None:
Piero Molino's avatar
Piero Molino committed
642
643
            ce_loss = torch.nn.CrossEntropyLoss()
            predicted_sentiment = classifier(torch.mean(true_hidden, dim=1))
644
            label = torch.tensor([label_class], device='cuda',
Piero Molino's avatar
Piero Molino committed
645
646
647
                                 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
648
        else:
Piero Molino's avatar
Piero Molino committed
649
650
651
652
653
            true_discrim_loss = 0

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

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

Piero Molino's avatar
Piero Molino committed
658
659
660
        log_probs = F.softmax(logits, dim=-1)

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

Piero Molino's avatar
Piero Molino committed
663
            # original_probs = top_k_filter(original_probs[:, -1, :]) #+ SmallConst
664
            unpert_logits = F.softmax(unpert_logits[:, -1, :], dim=-1)
Piero Molino's avatar
Piero Molino committed
665
666
667
668
            # likelywords = torch.topk(original_probs, k=10, dim=-1)
            # print(TOKENIZER.decode(likelywords[1].tolist()[0]))

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

671
            log_probs = top_k_filter(log_probs, k=top_k,
Piero Molino's avatar
Piero Molino committed
672
                                     probs=True)  # + SmallConst
Julien Chaumond's avatar
Julien Chaumond committed
673

Piero Molino's avatar
Piero Molino committed
674
675
            if torch.sum(log_probs) <= 1:
                log_probs = log_probs / torch.sum(log_probs)
Julien Chaumond's avatar
Julien Chaumond committed
676
677

        else:
678
            logits = top_k_filter(logits, k=top_k)  # + SmallConst
Piero Molino's avatar
Piero Molino committed
679
            log_probs = F.softmax(logits, dim=-1)
Julien Chaumond's avatar
Julien Chaumond committed
680
681

        if sample:
Piero Molino's avatar
Piero Molino committed
682
683
684
685
            # 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
686
        else:
Piero Molino's avatar
Piero Molino committed
687
688
689
690
691
692
            _, 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
693

Piero Molino's avatar
Piero Molino committed
694
    return output, true_discrim_loss, loss_in_time
Julien Chaumond's avatar
Julien Chaumond committed
695
696
697
698


def run_model():
    parser = argparse.ArgumentParser()
Piero Molino's avatar
Piero Molino committed
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
    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. Multiple BoWs separated by ;')
    parser.add_argument('--discrim', '-D', type=str, default=None,
                        choices=(
                        'clickbait', 'sentiment', 'toxicity', 'generic'),
                        help='Discriminator to use for loss-type 2')
    parser.add_argument('--discrim_weights', type=str, default=None,
                        help='Weights for the generic discriminator')
    parser.add_argument('--discrim_meta', type=str, default=None,
                        help='Meta information for the generic discriminator')
    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
714
715
716
717
718
719
    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)
Piero Molino's avatar
Piero Molino committed
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
    parser.add_argument('--nocuda', 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('--force-token', action='store_true', help='no cuda')
    parser.add_argument('--window_length', type=int, default=0,
                        help='Length of past which is being optimizer; 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='no cuda')
Julien Chaumond's avatar
Julien Chaumond committed
738
739
740
741
742
743

    args = parser.parse_args()

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)

Piero Molino's avatar
Piero Molino committed
744
    device = 'cpu' if args.nocuda else 'cuda'
Julien Chaumond's avatar
Julien Chaumond committed
745
746
747
748
749
750
751
752

    model = GPT2LMHeadModel.from_pretrained(
        args.model_path,
        output_hidden_states=True
    )
    model.to(device)
    model.eval()

Piero Molino's avatar
Piero Molino committed
753
    # Freeze GPT-2 weights
Julien Chaumond's avatar
Julien Chaumond committed
754
755
    for param in model.parameters():
        param.requires_grad = False
Piero Molino's avatar
Piero Molino committed
756
    pass
Julien Chaumond's avatar
Julien Chaumond committed
757
758

    if args.uncond:
Piero Molino's avatar
Piero Molino committed
759
760
        seq = [[50256, 50256]]

Julien Chaumond's avatar
Julien Chaumond committed
761
762
763
    else:
        raw_text = args.cond_text
        while not raw_text:
Piero Molino's avatar
Piero Molino committed
764
            print('Did you forget to add `--cond-text`? ')
Julien Chaumond's avatar
Julien Chaumond committed
765
            raw_text = input("Model prompt >>> ")
Piero Molino's avatar
Piero Molino committed
766
767
768
769
        seq = [[50256] + TOKENIZER.encode(raw_text)]

    collect_gen = dict()
    current_index = 0
770
    for tokenized_cond_text in seq:
Piero Molino's avatar
Piero Molino committed
771

772
        text = TOKENIZER.decode(tokenized_cond_text)
Piero Molino's avatar
Piero Molino committed
773
774
775
776
        print("=" * 40 + " Prefix of sentence " + "=" * 40)
        print(text)
        print("=" * 80)

777
778
779
        out1, out_perturb, discrim_loss_list, loss_in_time_list, actual_words = full_text_generation(
            model=model, context=tokenized_cond_text, device=device, **vars(args)
        )
Piero Molino's avatar
Piero Molino committed
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808

        text_whole = TOKENIZER.decode(out1.tolist()[0])

        print("=" * 80)
        print("=" * 40 + " Whole sentence (Original)" + "=" * 40)
        print(text_whole)
        print("=" * 80)

        out_perturb_copy = out_perturb

        for out_perturb in out_perturb_copy:
            # try:
            #    print("=" * 40 + " Whole sentence (Perturbed)" + "=" * 40)
            #    text_whole = TOKENIZER.decode(out_perturb.tolist()[0])
            #    print(text_whole)
            #    print("=" * 80)
            # except:
            #    pass
            # collect_gen[current_index] = [out, out_perturb, out1]
            ## Save the prefix, perturbed seq, original seq for each index
            print("=" * 40 + " Whole sentence (Perturbed)" + "=" * 40)
            keyword_tokens = [aa[-1][0] for aa in
                              actual_words] if actual_words else []
            output_tokens = out_perturb.tolist()[0]

            if args.colorama:
                import colorama

                text_whole = ''
809
810
                for tokenized_cond_text in output_tokens:
                    if tokenized_cond_text in keyword_tokens:
Piero Molino's avatar
Piero Molino committed
811
                        text_whole += '%s%s%s' % (
812
                        colorama.Fore.GREEN, TOKENIZER.decode([tokenized_cond_text]),
Piero Molino's avatar
Piero Molino committed
813
814
                        colorama.Style.RESET_ALL)
                    else:
815
                        text_whole += TOKENIZER.decode([tokenized_cond_text])
Piero Molino's avatar
Piero Molino committed
816
817
            else:
                text_whole = TOKENIZER.decode(out_perturb.tolist()[0])
Julien Chaumond's avatar
Julien Chaumond committed
818

Piero Molino's avatar
Piero Molino committed
819
820
            print(text_whole)
            print("=" * 80)
Julien Chaumond's avatar
Julien Chaumond committed
821

822
            collect_gen[current_index] = [tokenized_cond_text, out_perturb, out1]
Julien Chaumond's avatar
Julien Chaumond committed
823

Piero Molino's avatar
Piero Molino committed
824
            current_index = current_index + 1
Julien Chaumond's avatar
Julien Chaumond committed
825
826


Piero Molino's avatar
Piero Molino committed
827
    return
Julien Chaumond's avatar
Julien Chaumond committed
828
829


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