run_pplm.py 30 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
37
38
39
40
# 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

from transformers import GPT2Tokenizer
from transformers.file_utils import cached_path
from transformers.modeling_gpt2 import GPT2LMHeadModel

Piero Molino's avatar
Piero Molino committed
41

Julien Chaumond's avatar
Julien Chaumond committed
42
43
44
45
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)


Julien Chaumond's avatar
Julien Chaumond committed
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
class ClassificationHead(torch.nn.Module):
    """ Classification Head for the transformer """

    def __init__(self, class_size=5, embed_size=2048):
        super(ClassificationHead, self).__init__()
        self.class_size = class_size
        self.embed_size = embed_size
        # self.mlp1 = torch.nn.Linear(embed_size, embed_size)
        # self.mlp2 = (torch.nn.Linear(embed_size, class_size))
        self.mlp = torch.nn.Linear(embed_size, class_size)

    def forward(self, hidden_state):
        # hidden_state = F.relu(self.mlp1(hidden_state))
        # hidden_state = self.mlp2(hidden_state)
        logits = self.mlp(hidden_state)
        return logits


Piero Molino's avatar
Piero Molino committed
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
def perturb_past(past, model, prev, args, classifier, good_index=None,
                 stepsize=0.01, vocab_size=50257,
                 original_probs=None, accumulated_hidden=None, true_past=None,
                 grad_norms=None):
    window_length = args.window_length
    gm_scale, kl_scale = args.gm_scale, args.kl_scale
    one_hot_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_vectors.append(one_hot_good)

    # 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
148
149
150
151

    if accumulated_hidden is None:
        accumulated_hidden = 0

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

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

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

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

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

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

Piero Molino's avatar
Piero Molino committed
189
190
        # _, future_past = model(prev, past=perturbed_past)
        # hidden = model.hidden_states
Julien Chaumond's avatar
Julien Chaumond committed
191

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

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

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

Piero Molino's avatar
Piero Molino committed
236
237
            predicted_sentiment = classifier(new_accumulated_hidden / (
                        current_length + 1 + args.horizon_length))
Julien Chaumond's avatar
Julien Chaumond committed
238

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

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

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

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

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

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

Piero Molino's avatar
Piero Molino committed
279
        for p_ in past_perturb:
Julien Chaumond's avatar
Julien Chaumond committed
280
281
282
            p_.grad.data.zero_()

        new_past = []
Piero Molino's avatar
Piero Molino committed
283
284
285
        for p in past:
            new_past.append(p.detach())

Julien Chaumond's avatar
Julien Chaumond committed
286
287
        past = new_past

Piero Molino's avatar
Piero Molino committed
288
289
290
    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
291

Piero Molino's avatar
Piero Molino committed
292
    return perturbed_past, new_accumulated_hidden, grad_norms, loss_per_iter
Julien Chaumond's avatar
Julien Chaumond committed
293
294
295


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


Piero Molino's avatar
Piero Molino committed
367
def latent_perturb(model, args, context=None, sample=True, device='cuda'):
Julien Chaumond's avatar
Julien Chaumond committed
368
    classifier, class_id = get_classifier(
Piero Molino's avatar
Piero Molino committed
369
370
        args.discrim,
        args.label_class,
Julien Chaumond's avatar
Julien Chaumond committed
371
372
373
        device
    )

Piero Molino's avatar
Piero Molino committed
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
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
    # 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]

    good_index = []
    actual_words = None
    if args.bag_of_words:
        good_index = get_bag_of_words_indices(args.bag_of_words.split(";"))

        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]

    if args.bag_of_words and classifier:
Julien Chaumond's avatar
Julien Chaumond committed
453
        print("Both PPLM-BoW and PPLM-Discrim are on. This is not optimized.")
Piero Molino's avatar
Piero Molino committed
454
        args.loss_type = PPLM_BOW_DISCRIM
Julien Chaumond's avatar
Julien Chaumond committed
455

Piero Molino's avatar
Piero Molino committed
456
457
    elif args.bag_of_words:
        args.loss_type = PPLM_BOW
Julien Chaumond's avatar
Julien Chaumond committed
458
459
460
        print("Using PPLM-BoW")

    elif classifier is not None:
Piero Molino's avatar
Piero Molino committed
461
        args.loss_type = PPLM_DISCRIM
Julien Chaumond's avatar
Julien Chaumond committed
462
463
464
465
466
        print("Using PPLM-Discrim")

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

Piero Molino's avatar
Piero Molino committed
467
468
469
470
    original, _, _ = sample_from_hidden(model=model, args=args, context=context,
                                        device=device,
                                        perturb=False, good_index=good_index,
                                        classifier=classifier)
Julien Chaumond's avatar
Julien Chaumond committed
471
472
    torch.cuda.empty_cache()

Piero Molino's avatar
Piero Molino committed
473
474
475
476
477
478
479
480
481
482
483
484
485
    perturbed_list = []
    discrim_loss_list = []
    loss_in_time_list = []

    for i in range(args.num_samples):
        perturbed, discrim_loss, loss_in_time = sample_from_hidden(model=model,
                                                                   args=args,
                                                                   context=context,
                                                                   device=device,
                                                                   perturb=True,
                                                                   good_index=good_index,
                                                                   classifier=classifier)
        perturbed_list.append(perturbed)
Julien Chaumond's avatar
Julien Chaumond committed
486
        if classifier is not None:
Piero Molino's avatar
Piero Molino committed
487
488
            discrim_loss_list.append(discrim_loss.data.cpu().numpy())
        loss_in_time_list.append(loss_in_time)
Julien Chaumond's avatar
Julien Chaumond committed
489
490
491

    torch.cuda.empty_cache()

Piero Molino's avatar
Piero Molino committed
492
493
    return original, perturbed_list, discrim_loss_list, loss_in_time_list, actual_words

Julien Chaumond's avatar
Julien Chaumond committed
494

Piero Molino's avatar
Piero Molino committed
495
496
497
498
499
def sample_from_hidden(model, args, classifier, context=None, past=None,
                       device='cuda',
                       sample=True, perturb=True, good_index=None):
    output = torch.tensor(context, device=device, dtype=torch.long).unsqueeze(
        0) if context else None
Julien Chaumond's avatar
Julien Chaumond committed
500
501
502

    grad_norms = None
    loss_in_time = []
Piero Molino's avatar
Piero Molino committed
503
    for i in trange(args.length, ascii=True):
Julien Chaumond's avatar
Julien Chaumond committed
504
505

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

Piero Molino's avatar
Piero Molino committed
509
510
511
512
513
        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
514

Piero Molino's avatar
Piero Molino committed
515
516
517
518
            # Piero modified model call
            _, past, _ = model(output[:, :-1])
            original_probs, true_past, unpert_all_hidden = model(output)
            true_hidden = unpert_all_hidden[-1]
Julien Chaumond's avatar
Julien Chaumond committed
519
520

        else:
Piero Molino's avatar
Piero Molino committed
521
522
            # original_probs, true_past = model(output)
            # true_hidden = model.hidden_states
Julien Chaumond's avatar
Julien Chaumond committed
523

Piero Molino's avatar
Piero Molino committed
524
525
526
527
528
529
530
531
            # Piero modified model call
            original_probs, true_past, unpert_all_hidden = model(output)
            true_hidden = unpert_all_hidden[-1]

        # Modify the past if necessary

        if i >= args.grad_length:
            current_stepsize = args.stepsize * 0
Julien Chaumond's avatar
Julien Chaumond committed
532
        else:
Piero Molino's avatar
Piero Molino committed
533
            current_stepsize = args.stepsize
Julien Chaumond's avatar
Julien Chaumond committed
534

Piero Molino's avatar
Piero Molino committed
535
536
        if not perturb or args.num_iterations == 0:
            perturbed_past = past
Julien Chaumond's avatar
Julien Chaumond committed
537
538

        else:
Piero Molino's avatar
Piero Molino committed
539
540
541
            # Piero modified model call
            # accumulated_hidden = model.hidden_states[:, :-1, :]
            accumulated_hidden = true_hidden[:, :-1, :]
Julien Chaumond's avatar
Julien Chaumond committed
542
543
            accumulated_hidden = torch.sum(accumulated_hidden, dim=1)

Piero Molino's avatar
Piero Molino committed
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
            perturbed_past, _, grad_norms, loss_per_iter = perturb_past(past,
                                                                        model,
                                                                        prev,
                                                                        args,
                                                                        good_index=good_index,
                                                                        stepsize=current_stepsize,
                                                                        original_probs=original_probs,
                                                                        true_past=true_past,
                                                                        accumulated_hidden=accumulated_hidden,
                                                                        classifier=classifier,
                                                                        grad_norms=grad_norms)
            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
562
563

        if classifier is not None:
Piero Molino's avatar
Piero Molino committed
564
565
566
567
568
569
            ce_loss = torch.nn.CrossEntropyLoss()
            predicted_sentiment = classifier(torch.mean(true_hidden, dim=1))
            label = torch.tensor([args.label_class], device='cuda',
                                 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
570
        else:
Piero Molino's avatar
Piero Molino committed
571
572
573
574
575
576
577
578
            true_discrim_loss = 0

        # Piero modified model call
        # hidden = model.hidden_states  # update hidden
        # logits = model.forward_hidden(hidden)
        logits = logits[:, -1, :] / args.temperature  # + SmallConst

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

Piero Molino's avatar
Piero Molino committed
580
581
582
        log_probs = F.softmax(logits, dim=-1)

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

Piero Molino's avatar
Piero Molino committed
585
586
587
588
589
590
591
592
            # original_probs = top_k_filter(original_probs[:, -1, :]) #+ SmallConst
            original_probs = F.softmax(original_probs[:, -1, :], dim=-1)
            # likelywords = torch.topk(original_probs, k=10, dim=-1)
            # print(TOKENIZER.decode(likelywords[1].tolist()[0]))

            gm_scale = args.gm_scale
            log_probs = ((log_probs ** gm_scale) * (
                        original_probs ** (1 - gm_scale)))  # + SmallConst
Julien Chaumond's avatar
Julien Chaumond committed
593

Piero Molino's avatar
Piero Molino committed
594
595
            log_probs = top_k_filter(log_probs, k=args.top_k,
                                     probs=True)  # + SmallConst
Julien Chaumond's avatar
Julien Chaumond committed
596

Piero Molino's avatar
Piero Molino committed
597
598
            if torch.sum(log_probs) <= 1:
                log_probs = log_probs / torch.sum(log_probs)
Julien Chaumond's avatar
Julien Chaumond committed
599
600

        else:
Piero Molino's avatar
Piero Molino committed
601
602
            logits = top_k_filter(logits, k=args.top_k)  # + SmallConst
            log_probs = F.softmax(logits, dim=-1)
Julien Chaumond's avatar
Julien Chaumond committed
603
604

        if sample:
Piero Molino's avatar
Piero Molino committed
605
606
607
608
            # 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
609
        else:
Piero Molino's avatar
Piero Molino committed
610
611
612
613
614
615
            _, 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
616

Piero Molino's avatar
Piero Molino committed
617
    return output, true_discrim_loss, loss_in_time
Julien Chaumond's avatar
Julien Chaumond committed
618
619
620
621


def run_model():
    parser = argparse.ArgumentParser()
Piero Molino's avatar
Piero Molino committed
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
    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
637
638
639
640
641
642
    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
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
    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
661
662
663
664
665
666

    args = parser.parse_args()

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

Piero Molino's avatar
Piero Molino committed
667
    device = 'cpu' if args.nocuda else 'cuda'
Julien Chaumond's avatar
Julien Chaumond committed
668
669
670
671
672
673
674
675

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

Piero Molino's avatar
Piero Molino committed
676
    # Freeze GPT-2 weights
Julien Chaumond's avatar
Julien Chaumond committed
677
678
    for param in model.parameters():
        param.requires_grad = False
Piero Molino's avatar
Piero Molino committed
679
    pass
Julien Chaumond's avatar
Julien Chaumond committed
680
681

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

Julien Chaumond's avatar
Julien Chaumond committed
684
685
686
    else:
        raw_text = args.cond_text
        while not raw_text:
Piero Molino's avatar
Piero Molino committed
687
            print('Did you forget to add `--cond-text`? ')
Julien Chaumond's avatar
Julien Chaumond committed
688
            raw_text = input("Model prompt >>> ")
Piero Molino's avatar
Piero Molino committed
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
        seq = [[50256] + TOKENIZER.encode(raw_text)]

    collect_gen = dict()
    current_index = 0
    for out in seq:

        text = TOKENIZER.decode(out)
        print("=" * 40 + " Prefix of sentence " + "=" * 40)
        print(text)
        print("=" * 80)

        out1, out_perturb, discrim_loss_list, loss_in_time_list, actual_words = latent_perturb(
            model=model, args=args, context=out,
            device=device)

        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 = ''
                for out in output_tokens:
                    if out in keyword_tokens:
                        text_whole += '%s%s%s' % (
                        colorama.Fore.GREEN, TOKENIZER.decode([out]),
                        colorama.Style.RESET_ALL)
                    else:
                        text_whole += TOKENIZER.decode([out])
            else:
                text_whole = TOKENIZER.decode(out_perturb.tolist()[0])
Julien Chaumond's avatar
Julien Chaumond committed
741

Piero Molino's avatar
Piero Molino committed
742
743
            print(text_whole)
            print("=" * 80)
Julien Chaumond's avatar
Julien Chaumond committed
744

Piero Molino's avatar
Piero Molino committed
745
            collect_gen[current_index] = [out, out_perturb, out1]
Julien Chaumond's avatar
Julien Chaumond committed
746

Piero Molino's avatar
Piero Molino committed
747
            current_index = current_index + 1
Julien Chaumond's avatar
Julien Chaumond committed
748
749


Piero Molino's avatar
Piero Molino committed
750
    return
Julien Chaumond's avatar
Julien Chaumond committed
751
752


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