import math import torch import torch.nn as nn from .utils import log_sum_exp import pdb import sys sys.path.append('../../') from pytorch_transformers.modeling_bert import BertEmbeddings import torch.nn.functional as F class CARA(nn.Module): def __init__(self, encoder, decoder, tokenizer_encoder, tokenizer_decoder, args): # super(CARA, self).__init__() self.encoder = encoder self.decoder = decoder self.tokenizer_encoder = tokenizer_encoder self.tokenizer_decoder = tokenizer_decoder self.args = args self.nz = args.latent_size self.bos_token_id_list = self.tokenizer_decoder.encode(self.tokenizer_decoder.bos_token) self.pad_token_id = self.tokenizer_decoder.encode(self.tokenizer_decoder.pad_token)[0] # connector: from Bert hidden units to the latent space self.linear = nn.Linear(encoder.config.hidden_size, self.nz, bias=False) # # Standard Normal prior # loc = torch.zeros(self.nz, device=args.device) # scale = torch.ones(self.nz, device=args.device) # self.prior = torch.distributions.normal.Normal(loc, scale) self.label_embedding = nn.Embedding(args.label_size, self.nz, padding_idx=0) # use the same size as latent_z so as to use the same decoder.linear() self.latent_generator = nn.Linear(self.nz, self.nz) self.latent_classifier = nn.Linear(self.nz, args.label_size if args.label_size > 2 else 1) self.latent_discriminator = nn.Linear(self.nz, 1) self.gpt_embeddings = nn.Embedding(self.decoder.config.vocab_size, self.decoder.config.n_embd) self.gpt_embeddings.weight.data = decoder.transformer.wte.weight.data self.conv1 = nn.Conv1d(self.encoder.config.hidden_size, self.encoder.config.hidden_size, 3) self.classifier = nn.Linear(self.encoder.config.hidden_size, 1 if args.label_size <= 2 else args.label_size) self.CrossEntropyLoss = torch.nn.CrossEntropyLoss() self.BCEWithLogitsLoss = torch.nn.BCEWithLogitsLoss() def forward(self, input_seq_ids, tgt_seq_ids, cond_labels, attention_mask): # inputs: (B, seq_len) # labels: (B, seq_len) # cond_labels: (B), conditional labels. ones_label = torch.ones_like(cond_labels).to(dtype=torch.float32) zeros_label = torch.zeros_like(cond_labels).to(dtype=torch.float32) random_noise = torch.nn.init.normal_(torch.empty(input_seq_ids.size(0), self.nz)).to(device=input_seq_ids.device, dtype=torch.float32) # Encode inputs outputs = self.encoder(input_seq_ids, attention_mask=attention_mask) pooled_hidden_fea = outputs[1] # (B, dim_h) # Encode z latent_z = self.linear(pooled_hidden_fea) # (B, nz) # Generate z gen_z = self.latent_generator(random_noise) # (B, nz) #################### Latent discriminator for sampling from a simple distribution #################### prob_encode_z_dis = self.latent_discriminator(latent_z).squeeze(1).float() # (B) prob_gen_z_dis = self.latent_discriminator(gen_z).squeeze(1).float() # (B) # Train latent discriminator loss_lsd = self.BCEWithLogitsLoss(prob_gen_z_dis, zeros_label) + self.BCEWithLogitsLoss(prob_encode_z_dis, ones_label) acc_encode_z_dis = ((prob_encode_z_dis >= 0).float() == ones_label).float() acc_gen_z_dis = ((prob_gen_z_dis >= 0).float() == zeros_label).float() # Train sampler adversarially loss_lsg = self.BCEWithLogitsLoss(prob_gen_z_dis, ones_label) #################### Latent classifier for disentanglement #################### prob_encode_z_cls = self.latent_classifier(latent_z) # (B, n_labels) if self.args.label_size <= 2: prob_encode_z_cls = prob_encode_z_cls.squeeze(1) # (B) # Train latent classifier loss_lsc = self.BCEWithLogitsLoss(prob_encode_z_cls, cond_labels.float()) acc_encode_z_cls = ((prob_encode_z_cls >= 0).float() == cond_labels.float()).float() # Train encoder adversarially loss_encoder = 1 - self.BCEWithLogitsLoss(prob_encode_z_cls, cond_labels.float()) else: # Train latent classifier loss_lsc = self.CrossEntropyLoss(prob_encode_z_cls, cond_labels) acc_encode_z_cls = (torch.argmax(prob_encode_z_cls, dim=-1) == cond_labels).float() # Train encoder adversarially loss_encoder = 1 - self.CrossEntropyLoss(prob_encode_z_cls, cond_labels) #################### Recontruction loss with latent z and label emb #################### # Embed labels label_emb = self.label_embedding(cond_labels) # (B, hidden_size) # past_label = self.decoder.linear(label_emb) # (B, n_blocks * hidden_size) # todo: use the same linear layer for latent_z for now. if self.args.label_size <= 2: sampled_cond_labels = 1 - cond_labels else: raise NotImplementedError # todo: currently only implemented for binary labels. need to change for multi-class labels. sampled_label_emb = self.label_embedding(sampled_cond_labels) # (B, hidden_size) # past_sampled_label = self.decoder.linear(sampled_label_emb) # (B, n_blocks * hidden_size) # todo: use the same linear layer for latent_z for now. past_sampled_label = sampled_label_emb # Generate based on encoded z and gt labels. (reconstruction) # past_z = self.decoder.linear(latent_z) # (B, n_blocks * hidden_size) past_z = latent_z # gen_past_z = self.decoder.linear(gen_z) # (B, n_blocks * hidden_size) gen_past_z = gen_z # (B, n_blocks * hidden_size) # past = torch.cat([past_z.unsqueeze(1), past_label.unsqueeze(1)], dim=1) # (B, 2, n_blocks * hidden_size) past = latent_z + label_emb # (B, n_blocks * hidden_size) outputs = self.decoder(input_ids=tgt_seq_ids, past=past, labels=tgt_seq_ids, label_ignore=self.pad_token_id) loss_rec = outputs[0] #################### Train a classifier in the observation space #################### tgt_emb = self.gpt_embeddings(tgt_seq_ids) tgt_encode = self.conv1(tgt_emb.transpose(1, 2)) # (B, dim_h, seq_len) tgt_encode = torch.mean(tgt_encode, dim=-1) # (B, dim_h) prob_cls = self.classifier(tgt_encode) # (B, n_labels) if self.args.label_size <= 2: prob_cls = prob_cls.squeeze(1) loss_cls = self.BCEWithLogitsLoss(prob_cls, cond_labels.float()) pred_cls = (prob_cls >= 0).to(dtype=torch.long) else: loss_cls = self.CrossEntropyLoss(prob_cls, cond_labels) pred_cls = torch.argmax(prob_cls, dim=-1) acc_cls = (pred_cls == cond_labels).float() # Generate based on encoded z and sampled labels (attribute transfer) # at_past = torch.cat([past_z.unsqueeze(1), past_sampled_label.unsqueeze(1)], dim=1) # (B, 2, n_blocks * hidden_size) # at_generated_soft = self.sample_sequence_conditional_batch_soft(past=at_past, context=self.bos_token_id_list) # (B, seq_len, vocab_size) # # Classifier on attribute transfer generated sentences. Train Generator on attribute transfer. # at_soft_emb = torch.matmul(at_generated_soft, self.gpt_embeddings.weight) # at_soft_encode = self.conv1(at_soft_emb.transpose(1, 2)) # (B, dim_h, seq_len) # at_soft_encode = torch.mean(at_soft_encode, dim=-1) # (B, dim_h) # prob_at_soft_cls = self.classifier(at_soft_encode) # (B, 1) # if self.args.label_size <= 2: # prob_at_soft_cls = prob_at_soft_cls.squeeze(1) # loss_at_soft_cls = self.BCEWithLogitsLoss(prob_at_soft_cls, sampled_cond_labels.float()) # pred_at_soft_cls = (prob_at_soft_cls >= 0).to(torch.long) # else: # loss_at_soft_cls = self.CrossEntropyLoss(prob_at_soft_cls, sampled_cond_labels) # pred_at_soft_cls = torch.argmax(prob_at_soft_cls, dim=-1) # acc_at_soft_cls = (pred_at_soft_cls == sampled_cond_labels).float() # Loss loss_latent_space = (loss_encoder + loss_lsc) + (loss_lsd + loss_lsg) + self.args.beta_cls * loss_cls # + loss_at_soft_cls loss = loss_rec + 0.0 * loss_latent_space if not self.training: # Generate based on encoded z and gt labels generated = self.sample_sequence_conditional_batch(past=past, context=self.bos_token_id_list) # Generate based on encoded z and sampled labels (attribute transfer) # at_past = torch.cat([past_z.unsqueeze(1), past_sampled_label.unsqueeze(1)], dim=1) # (B, 2, n_blocks * hidden_size) at_past = past_z + past_sampled_label # (B, n_blocks * hidden_size) at_generated = self.sample_sequence_conditional_batch(past=at_past, context=self.bos_token_id_list) # (B, seq_len) # Generate based on sampled z and sampled labels. (conditional generation) # cg_past = torch.cat([gen_past_z.unsqueeze(1), past_sampled_label.unsqueeze(1)], dim=1) # (B, 2, n_blocks * hidden_size) cg_past = gen_past_z + past_sampled_label # (B, n_blocks * hidden_size) cg_generated = self.sample_sequence_conditional_batch(past=cg_past, context=self.bos_token_id_list) # (B, seq_len) # classifier on gt generated sentences. ge_emb = self.gpt_embeddings(generated) ge_encode = self.conv1(ge_emb.transpose(1, 2)) # (B, dim_h, seq_len) ge_encode = torch.mean(ge_encode, dim=-1) # (B, dim_h) prob_ge_cls = self.classifier(ge_encode) # (B, 1) if self.args.label_size <= 2: pred_ge_cls = (prob_ge_cls.squeeze(1) >= 0).to(torch.long) else: pred_ge_cls = torch.argmax(prob_ge_cls, dim=-1) acc_ge_cls = (pred_ge_cls == cond_labels).float() # classifier on attribute transfer generated sentences. at_emb = self.gpt_embeddings(at_generated) at_encode = self.conv1(at_emb.transpose(1, 2)) # (B, dim_h, seq_len) at_encode = torch.mean(at_encode, dim=-1) # (B, dim_h) prob_at_cls = self.classifier(at_encode) # (B, 1) if self.args.label_size <= 2: pred_at_cls = (prob_at_cls.squeeze(1) >= 0).to(torch.long) else: pred_at_cls = torch.argmax(prob_at_cls, dim=-1) acc_at_cls = (pred_at_cls == sampled_cond_labels).float() # classifier on conditional generated sentences. cg_emb = self.gpt_embeddings(cg_generated) cg_encode = self.conv1(cg_emb.transpose(1, 2)) # (B, dim_h, seq_len) cg_encode = torch.mean(cg_encode, dim=-1) # (B, dim_h) prob_cg_cls = self.classifier(cg_encode) # (B, 1) if self.args.label_size <= 2: pred_cg_cls = (prob_cg_cls.squeeze(1) >= 0).to(torch.long) else: pred_cg_cls = torch.argmax(prob_cg_cls, dim=-1) acc_cg_cls = (pred_cg_cls == sampled_cond_labels).float() result = { 'sampled_cond_labels': sampled_cond_labels, 'cond_labels': cond_labels, 'tgt_seq_ids': tgt_seq_ids, 'generated': generated, 'at_generated': at_generated, 'cg_generated': cg_generated, 'acc_encode_z_dis': acc_encode_z_dis, 'acc_gen_z_dis': acc_gen_z_dis, 'acc_encode_z_cls': acc_encode_z_cls, 'acc_cls': acc_cls, 'acc_ge_cls': acc_ge_cls, 'acc_at_cls': acc_at_cls, 'acc_cg_cls': acc_cg_cls, 'pred_cls': pred_cls, 'pred_ge_cls': pred_ge_cls, 'pred_at_cls': pred_at_cls, 'pred_cg_cls': pred_cg_cls, } return result loss_dict = { 'loss': loss, 'loss_rec': loss_rec, 'loss_encoder': loss_encoder, 'loss_lsc': loss_lsc, 'loss_lsd': loss_lsd, 'loss_lsg': loss_lsg, 'loss_cls': loss_cls, # 'loss_at_soft_cls': loss_at_soft_cls, } acc_dict = { 'acc_encode_z_dis': acc_encode_z_dis, 'acc_gen_z_dis': acc_gen_z_dis, 'acc_encode_z_cls': acc_encode_z_cls, 'acc_cls': acc_cls, # 'acc_at_soft_cls': acc_at_soft_cls, } return loss_dict, acc_dict def sample_sequence_conditional_batch(self, past, context): # context: a single id of # past: (B, past_seq_len dim_h) num_samples = past.size(0) context = torch.tensor(context, dtype=torch.long, device=past.device) context = context.unsqueeze(0).repeat(num_samples, 1) generated = context # (B, 1) # with torch.no_grad(): while generated.size(-1) < self.args.block_size: inputs = {'input_ids': generated, 'past': past} outputs = self.decoder(**inputs) # Note: we could also use 'past' with GPT-2/Transfo-XL/XLNet (cached hidden-states) lm_logits = outputs[0] # softmax sample next_tokens_logits = lm_logits[:, -1, :] / self.args.temperature # (B, 1, vocab_size) filtered_logits = self.top_k_top_p_filtering_batch(next_tokens_logits, top_k=self.args.top_k, top_p=self.args.top_p) # (B, 1, vocab_size) filtered_logits = F.softmax(filtered_logits, dim=-1) next_tokens = torch.multinomial(filtered_logits, num_samples=1) # (B, 1) generated = torch.cat((generated, next_tokens), dim=1) # (B, seq_len+1) not_finished = next_tokens != self.tokenizer_decoder.encode('')[0] if torch.sum(not_finished) == 0: break return generated # (B, seq_len) def top_k_top_p_filtering_batch(self, logits, top_k=0, top_p=0.0, filter_value=-float('Inf')): """ Filter a distribution of logits using top-k and/or nucleus (top-p) filtering Args: logits: logits distribution shape (vocabulary size) top_k > 0: keep only top k tokens with highest probability (top-k filtering). top_p > 0.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering). Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751) From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317 """ # assert logits.dim() == 1 # batch size 1 for now - could be updated for more but the code would be less clear top_k = min(top_k, logits.size(-1)) # Safety check if top_k > 0: # Remove all tokens with a probability less than the last token of the top-k threshold = torch.topk(logits, top_k, dim=-1)[0][:, -1, None] logits.masked_fill_(logits < threshold, filter_value) # (B, vocab_size) if top_p > 0.0: sorted_logits, sorted_indices = torch.sort(logits, descending=True) # (B, vocab_size) cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1) # (B, vocab_size) # Remove tokens with cumulative probability above the threshold sorted_indices_to_remove = cumulative_probs > top_p # Shift the indices to the right to keep also the first token above the threshold sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone() sorted_indices_to_remove[..., 0] = 0 indices_to_remove = sorted_indices[sorted_indices_to_remove] logits.masked_fill_(indices_to_remove, filter_value) return logits def sample_sequence_conditional_batch_soft(self, past, context): # context: a single id of # past: (B, past_seq_len dim_h) num_samples = past.size(0) context = torch.tensor(context, dtype=torch.long, device=past.device).unsqueeze(0).repeat(num_samples, 1) # (B, 1) context_soft = torch.FloatTensor(num_samples, self.decoder.config.vocab_size).zero_().to(device=past.device) # (B, vocab_size) context_soft.scatter_(1, context, 1) # (B, vocab_size) generated_soft = context_soft.unsqueeze(1) # (B, 1, vocab_size) # with torch.no_grad(): while generated_soft.size(1) < self.args.block_size: # generated_soft: (B, seq_len, vocab_size) inputs = {'soft_ids': generated_soft, 'past': past} outputs = self.decoder(**inputs) # Note: we could also use 'past' with GPT-2/Transfo-XL/XLNet (cached hidden-states) lm_logits = outputs[0] # (B, seq_len, vocab_size) # Gumbel softmax sample next_tokens_soft = gumbel_softmax(logits=lm_logits[:, -1:, :], temperature=self.args.soft_temperature, hard=False) # (B, 1, vocab_size) generated_soft = torch.cat((generated_soft, next_tokens_soft), dim=1) # (B, seq_len+1, vocab_size) # # softmax sample # next_tokens_logits = lm_logits[:, -1, :] / self.args.temperature # (B, 1, vocab_size) # filtered_logits = self.top_k_top_p_filtering_batch(next_tokens_logits, top_k=self.args.top_k, top_p=self.args.top_p) # (B, 1, vocab_size) # filtered_logits = F.softmax(filtered_logits, dim=-1) # next_tokens = torch.multinomial(filtered_logits, num_samples=1) # (B, 1) # generated = torch.cat((generated, next_tokens), dim=1) # (B, seq_len+1) next_tokens = torch.argmax(next_tokens_soft, dim=-1) # (B, 1) not_finished = next_tokens != self.tokenizer_decoder.encode('')[0] if torch.sum(not_finished) == 0: break return generated_soft # (B, seq_len, vocab_size) ### Gumbel Softmax def gumbel_softmax(logits, temperature, hard=False): """Sample from the Gumbel-Softmax distribution and optionally discretize. Args: logits: [..., n_class] unnormalized log-probs temperature: non-negative scalar hard: if True, take argmax, but differentiate w.r.t. soft sample y Returns: [..., n_class] sample from the Gumbel-Softmax distribution. If hard=True, then the returned sample will be one-hot, otherwise it will be a probabilitiy distribution that sums to 1 across classes """ y = gumbel_softmax_sample(logits, temperature) # (..., n_class) if hard: # return onehot shape = y.size() _, ind = y.max(dim=-1) y_hard = torch.zeros_like(y).view(-1, shape[-1]) y_hard.scatter_(1, ind.view(-1, 1), 1) # one hot y_hard = y_hard.view(*shape) # Set gradients w.r.t. y_hard gradients w.r.t. y y = (y_hard - y).detach() + y return y # (..., n_class) from torch.nn import functional as F def gumbel_softmax_sample(logits, temperature): y = logits + sample_gumbel(logits.size(), logits.device) return F.softmax(y / temperature, dim=-1) def sample_gumbel(shape, device, eps=1e-20): U = torch.rand(shape).to(device=device) return -torch.log(-torch.log(U + eps) + eps)