paper_code/gpt2tunediscrim.py

#! /usr/bin/env python3
# coding=utf-8

# This code is licensed under a non-commercial license.

import os
import sys
import argparse
from tqdm import trange
from torchtext import data as torchtext_data
from torchtext import datasets

import torch
import torch.utils.data as data

from torchtext.vocab import Vectors, GloVe, CharNGram, FastText
from nltk.tokenize.treebank import TreebankWordDetokenizer
import torch
import torch.optim
import torch.nn.functional as F
import numpy as np
from IPython import embed
from operator import add
from run_gpt2 import top_k_logits
from style_utils import to_var
import copy
import pickle
from torch.utils.data import DataLoader
from torch.utils.data.dataset import random_split
import torch.optim as optim

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

lab_root = os.path.join(os.path.abspath(os.path.dirname(__file__)), '..', '..')
sys.path.insert(1, lab_root)

from pytorch_pretrained_bert import GPT2LMHeadModel, GPT2Tokenizer
from torch.autograd import Variable

tokenizer = GPT2Tokenizer.from_pretrained('gpt-2_pt_models/345M/')

model = GPT2LMHeadModel.from_pretrained('gpt-2_pt_models/345M/')


class ClassificationHead(torch.nn.Module):
    """ Language Model 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):
        # Truncated Language modeling logits (we remove the last token)
        # h_trunc = h[:, :-1].contiguous().view(-1, self.n_embd)
        # lm_logits = F.relu(self.mlp1(hidden_state))
        # lm_logits = self.mlp2(lm_logits)
        lm_logits = self.mlp(hidden_state)
        return lm_logits


class Discriminator(torch.nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.classifierhead = ClassificationHead()
        self.model = model
        self.spltoken = Variable(torch.randn(1, 1, 1024).type(torch.FloatTensor), requires_grad=True)
        self.spltoken = self.spltoken.repeat(10, 1, 1)
        self.spltoken = self.spltoken.cuda()

    def train(self):
        for param in self.model.parameters():
            param.requires_grad = False
        pass

    def forward(self, x):
        x = model.forward_embed(x)
        x = torch.cat((x, self.spltoken), dim=1)
        _, x = model.forward_transformer_embed(x, add_one=True)
        x = self.classifierhead(x[-1][:, -1, :])
        x = F.log_softmax(x, dim=-1)
        return x


class Discriminator2(torch.nn.Module):
    def __init__(self, class_size=5, embed_size=1024):
        super(Discriminator2, self).__init__()
        self.classifierhead = ClassificationHead(class_size=class_size, embed_size=embed_size)
        self.model = model
        self.embed_size = embed_size

    def get_classifier(self):
        return self.classifierhead

    def train_custom(self):
        for param in self.model.parameters():
            param.requires_grad = False
        pass
        self.classifierhead.train()

    def forward(self, x):
        x = model.forward_embed(x)
        hidden, x = model.forward_transformer_embed(x)
        x = torch.sum(hidden, dim=1)
        x = self.classifierhead(x)
        x = F.log_softmax(x, dim=-1)
        return x

class Discriminator2mean(torch.nn.Module):
    def __init__(self, class_size=5, embed_size=1024):
        super(Discriminator2mean, self).__init__()
        self.classifierhead = ClassificationHead(class_size=class_size, embed_size=embed_size)
        self.model = model
        self.embed_size = embed_size

    def get_classifier(self):
        return self.classifierhead

    def train_custom(self):
        for param in self.model.parameters():
            param.requires_grad = False
        pass
        self.classifierhead.train()

    def forward(self, x):
        mask_src = 1 - x.eq(0).unsqueeze(1).type(torch.FloatTensor).cuda().detach()
        mask_src = mask_src.repeat(1, self.embed_size, 1)
        x = model.forward_embed(x)
        hidden, x = model.forward_transformer_embed(x)
        #  Hidden has shape batch_size x length x embed-dim

        hidden = hidden.permute(0, 2, 1)
        _, _, batch_length = hidden.shape
        hidden = hidden * mask_src  # / torch.sum(mask_src, dim=-1).unsqueeze(2).repeat(1, 1, batch_length)
        #
        hidden = hidden.permute(0, 2, 1)
        x = torch.sum(hidden, dim=1)/(torch.sum(mask_src, dim=-1).detach() + 1e-10)
        x = self.classifierhead(x)
        x = F.log_softmax(x, dim=-1)
        return x

class Dataset(data.Dataset):
    def __init__(self, X, y):
        """Reads source and target sequences from txt files."""
        self.X = X
        self.y = y

    def __len__(self):
        return len(self.X)

    def __getitem__(self, index):
        """Returns one data pair (source and target)."""
        d = {}
        d['X'] = self.X[index]
        d['y'] = self.y[index]
        return d


def collate_fn(data):
    def merge(sequences):
        lengths = [len(seq) for seq in sequences]

        padded_seqs = torch.zeros(len(sequences), max(lengths)).long().cuda()  # padding index 0
        for i, seq in enumerate(sequences):
            end = lengths[i]
            padded_seqs[i, :end] = seq[:end]
        return padded_seqs, lengths

    data.sort(key=lambda x: len(x["X"]), reverse=True)  # sort by source seq

    item_info = {}
    for key in data[0].keys():
        item_info[key] = [d[key] for d in data]

    # input
    x_batch, _ = merge(item_info['X'])
    y_batch = item_info['y']

    return x_batch, torch.tensor(y_batch, device='cuda', dtype=torch.long)


def train_epoch(data_loader, discriminator, device='cuda', args=None, epoch=1):
    optimizer = optim.Adam(discriminator.parameters(), lr=0.0001)
    discriminator.train_custom()

    for batch_idx, (data, target) in enumerate(data_loader):
        data, target = data.to(device), target.to(device)

        optimizer.zero_grad()

        output = discriminator(data)
        loss = F.nll_loss(output, target)
        loss.backward(retain_graph=True)
        optimizer.step()

        if batch_idx % args.log_interval == 0:
            print('Relu Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(data_loader.dataset),
                       100. * batch_idx / len(data_loader), loss.item()))


def test_epoch(data_loader, discriminator, device='cuda', args=None):
    discriminator.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in data_loader:
            data, target = data.to(device), target.to(device)
            output = discriminator(data)
            test_loss += F.nll_loss(output, target, reduction='sum').item()  # sum up batch loss
            pred = output.argmax(dim=1, keepdim=True)  # get the index of the max log-probability
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(data_loader.dataset)

    print('\nRelu Test set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(data_loader.dataset),
        100. * correct / len(data_loader.dataset)))


def main():
    parser = argparse.ArgumentParser(description='Train a discriminator on top of GPT-2 representations')
    parser.add_argument('--batch-size', type=int, default=64, metavar='N',
                        help='input batch size for training (default: 64)')
    parser.add_argument('--log-interval', type=int, default=10, metavar='N',
                        help='how many batches to wait before logging training status')
    parser.add_argument('--epochs', type=int, default=10, metavar='N',
                        help='Number of training epochs')
    parser.add_argument('--save-model', action='store_true', help='whether to save the model')
    parser.add_argument('--dataset-label', type=str, default='SST',choices=('SST', 'clickbait', 'toxic'))
    args = parser.parse_args()

    batch_size = args.batch_size
    device = 'cuda'
    # load sst
    if args.dataset_label == 'SST':
        text = torchtext_data.Field()
        label = torchtext_data.Field(sequential=False)
        train_data, val_data, test_data = datasets.SST.splits(text, label, fine_grained=True, train_subtrees=True,
                                                              # filter_pred=lambda ex: ex.label != 'neutral'
                                                              )
        x = []
        y = []
        d = {"positive": 0, "negative": 1, "very positive": 2, "very negative": 3, "neutral": 4}

        for i in range(len(train_data)):
            seq = TreebankWordDetokenizer().detokenize(vars(train_data[i])["text"])
            seq = tokenizer.encode(seq)
            seq = torch.tensor(seq, device=device, dtype=torch.long)
            x.append(seq)
            y.append(d[vars(train_data[i])["label"]])

        dataset = Dataset(x, y)

        test_x = []
        test_y = []
        for i in range(len(test_data)):
            seq = TreebankWordDetokenizer().detokenize(vars(test_data[i])["text"])
            seq = tokenizer.encode(seq)
            seq = torch.tensor([50256] + seq, device=device, dtype=torch.long)
            test_x.append(seq)
            test_y.append(d[vars(test_data[i])["label"]])
        test_dataset = Dataset(test_x, test_y)
        discriminator = Discriminator2mean(class_size=5).to(device)

    elif args.dataset_label == 'clickbait':
        # data = pickle.load(open("/home/gilocal/lab/exp/language/datasets/clickbait/clickbait.p", "r"))
        with open("datasets/clickbait/clickbait_train_prefix.txt") as f:
            data = []
            for d in f:
                try:
                    data.append(eval(d))
                except:
                    continue
        x = []
        y = []
        for d in data:
            try:
                # seq = tokenizer.encode("Apple's iOS 9 'App thinning' feature will give your phone's storage a boost")
                try:
                    seq = tokenizer.encode(d["text"])
                except:
                    continue
                seq = torch.tensor([50256] + seq, device=device, dtype=torch.long)
                x.append(seq)
                y.append(d['label'])
            except:
                pass

        dataset = Dataset(x, y)
        train_size = int(0.9 * len(dataset))
        test_size = len(dataset) - train_size
        dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, test_size])
        discriminator = Discriminator2mean(class_size=2).to(device)

    elif args.dataset_label == 'toxic':
        # data = pickle.load(open("/home/gilocal/lab/exp/language/datasets/clickbait/clickbait.p", "r"))
        with open("datasets/toxic/toxic_train.txt") as f:
            data = []
            for d in f:
                data.append(eval(d))

        x = []
        y = []
        for d in data:
            try:
                # seq = tokenizer.encode("Apple's iOS 9 'App thinning' feature will give your phone's storage a boost")
                seq = tokenizer.encode(d["text"])

                device = 'cuda'
                if(len(seq)<100):
                    seq = torch.tensor([50256] + seq, device=device, dtype=torch.long)
                else:
                    continue
                x.append(seq)
                y.append(int(np.sum(d['label'])>0))
            except:
                pass

        dataset = Dataset(x, y)
        print(dataset)
        print(len(dataset))
        train_size = int(0.9 * len(dataset))
        test_size = len(dataset) - train_size
        dataset, test_dataset = torch.utils.data.random_split(dataset, [train_size, test_size])
        discriminator = Discriminator2mean(class_size=2).to(device)

    data_loader = torch.utils.data.DataLoader(dataset=dataset,
                                              batch_size=batch_size,
                                              shuffle=True, collate_fn=collate_fn)
    test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                              batch_size=batch_size, collate_fn=collate_fn)

    for epoch in range(args.epochs):
        train_epoch(discriminator=discriminator, data_loader=data_loader, args=args, device=device, epoch=epoch)
        test_epoch(data_loader=test_loader, discriminator=discriminator, args=args)
        seq = tokenizer.encode("This is incredible! I love it, this is the best chicken I have ever had.")
        seq = torch.tensor([seq], device=device, dtype=torch.long)
        print(discriminator(seq))

        if (args.save_model):
            torch.save(discriminator.state_dict(),
                       "discrim_models/{}_mean_lin_discriminator_{}.pt".format(args.dataset_label, epoch))
            torch.save(discriminator.get_classifier().state_dict(),
                       "discrim_models/{}_classifierhead.pt".format(args.dataset_label))

    seq = tokenizer.encode("This is incredible! I love it, this is the best chicken I have ever had.")
    seq = torch.tensor([seq], device=device, dtype=torch.long)
    print(discriminator(seq))


if __name__ == '__main__':
    main()