# Copyright (c) Microsoft, Inc. 2020 # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. # # Zhou Bo # Date: 01/15/2020 # import copy import torch import os import random import json from .ops import * from .bert import * from .bert import BertLayer from .config import ModelConfig from .cache_utils import load_model_state from .nnmodule import NNModule # from ..utils.bad_grad_viz import register_hooks __all__ = ['WywLM'] def flatten_states(q_states, mask_index): q_states = q_states.reshape((-1, q_states.size(-1))) q_states = q_states.index_select(0, mask_index) return q_states class UGDecoder(torch.nn.Module): def __init__(self, config, vocab_size): super().__init__() self.config = config self.position_biased_input = getattr(config, 'position_biased_input', True) # self.layer = torch.nn.ModuleList([BertLayer(config) for _ in range(2)]) # self.causal_mask = torch.tril(torch.ones((input_ids.dim(0), input_ids.dim(1), input_ids.dim(1))), diagonal=0) def forward(self, ctx_layers, word_embedding, input_ids, z_states, attention_mask, \ encoder, target_ids=None, relative_pos=None, decode=False, s2s_idx=None): causal_outputs, lm_outputs = self.emd_context_layer(ctx_layers, z_states, attention_mask, encoder, target_ids, input_ids, relative_pos=relative_pos, decode=decode, word_embedding=word_embedding, s2s_idx=s2s_idx) # loss_fct = torch.nn.CrossEntropyLoss(reduction='none') # ctx_layer = mlm_ctx_layers[-1] # lm_logits = lm_logits.view(-1, lm_logits.size(-1)) return causal_outputs[-1], lm_outputs[-1] def emd_context_layer(self, encoder_layers, z_states, attention_mask, encoder, target_ids, input_ids,\ relative_pos=None, decode=False, word_embedding=None, s2s_idx=None): # if decode: # attention_mask = torch.tril(torch.ones((input_ids.shape[0], 1, input_ids.shape[1], input_ids.shape[1])), diagonal=0).to(input_ids.device) # else: if attention_mask.dim()<=2: extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2) att_mask = extended_attention_mask.byte() attention_mask = att_mask*att_mask.squeeze(-2).unsqueeze(-1) elif attention_mask.dim()==3: attention_mask = attention_mask.unsqueeze(1) if not self.position_biased_input: lm_outputs = [] # else: hidden_states = encoder_layers[-2] layers = [encoder.layer[-1] for _ in range(2)] z_states += hidden_states query_states = z_states query_mask = attention_mask rel_embeddings = encoder.get_rel_embedding() for layer in layers: # TODO: pass relative pos ids output = layer(hidden_states, query_mask, return_att=False, query_states=query_states, relative_pos=relative_pos, rel_embeddings=rel_embeddings) query_states = output lm_outputs.append(query_states) # if decode: attention_mask = torch.tril(torch.ones((input_ids.shape[0], 1, input_ids.shape[1], input_ids.shape[1])), diagonal=0).to(input_ids.device) causal_outputs = [] # with torch.no_grad(): target_embd = word_embedding(target_ids) target_embd += z_states.detach() # self attention of target output = layers[-2](target_embd, attention_mask, return_att=False, query_states=target_embd, relative_pos=relative_pos, rel_embeddings=encoder.get_rel_embedding()) causal_outputs.append(output) # cross attention output = layers[-1](output, attention_mask, return_att=False, query_states=query_states, relative_pos=relative_pos, rel_embeddings=encoder.get_rel_embedding()) causal_outputs.append(output) else: causal_outputs = [encoder_layers[-1]] lm_outputs = [encoder_layers[-1]] return causal_outputs, lm_outputs def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int): """ Shift input ids one token to the right. """ shifted_input_ids = input_ids.new_zeros(input_ids.shape) shifted_input_ids[:, 1:] = input_ids[:, :-1].clone() shifted_input_ids[:, 0] = decoder_start_token_id if pad_token_id is None: raise ValueError("self.model.config.pad_token_id has to be defined.") # replace possible -100 values in labels by `pad_token_id` shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id) return shifted_input_ids class WywLMLoss(torch.nn.Module): def __init__(self, config) -> None: super().__init__() self.loss_fn = torch.nn.CrossEntropyLoss(reduction='mean') hidden_size = getattr(config, 'embedding_size', config.hidden_size) self.compare = torch.nn.Linear(hidden_size * 3, 2) # self.mlm_head = BertLMPredictionHead(config, config.vocab_size) self.lm_head = BertLMPredictionHead(config, config.vocab_size) def forward(self, logits, lm_logits, target_ids, dict_pos, input_ids, target_ids_s2s, decode=False, ebd_weight=None, task=0): loss_compare = torch.tensor(0).to(logits).float() mlm_loss = torch.tensor(0).to(logits).float() lm_loss = torch.tensor(0).to(logits).float() # else: if task == 1: compare_logits = [] compare_labels = [] for bi, sampel_pos in enumerate(dict_pos): num_pos = int((sampel_pos > 0).sum().detach().cpu().numpy() / 4) - 1 if num_pos <= 1: continue for pi in range(num_pos): pos = sampel_pos[pi] entry_logits = logits[bi][pos[0]: pos[1]] desc_logits = logits[bi][pos[2]: pos[3]] neg_num = random.randint(0, num_pos) # torch.randint(low=0, high=num_pos, size=(1,)) ids_neg = input_ids[bi][sampel_pos[neg_num][0]: sampel_pos[neg_num][1]] ids_pos = input_ids[bi][pos[0]: pos[1]] if neg_num == pi or (ids_neg.shape == ids_pos.shape and torch.all(ids_neg == ids_pos)): neg_num = -1 for ni in range(num_pos): neg_num = random.randint(0, num_pos)# torch.randint(low=0, high=num_pos, size=(1,)) ids_neg = input_ids[bi][sampel_pos[neg_num][0]: sampel_pos[neg_num][1]] if neg_num != pi and (ids_neg.shape != ids_pos.shape or not torch.all(ids_neg == ids_pos)): break else: neg_num = -1 if neg_num == -1: continue neg_desc_logits = logits[bi][sampel_pos[neg_num][2]: sampel_pos[neg_num][3]] if torch.any(torch.isnan(neg_desc_logits)): print('error') entry_logits = entry_logits.mean(dim=0, keepdim=True).float() desc_logits = desc_logits.mean(dim=0, keepdim=True).float() neg_desc_logits = neg_desc_logits.mean(dim=0, keepdim=True).float() compare_logits.append(torch.concat([entry_logits, desc_logits, entry_logits - desc_logits], dim=1)) compare_logits.append(torch.concat([entry_logits, neg_desc_logits, entry_logits - neg_desc_logits], dim=1)) compare_labels += [1, 0] if len(compare_logits) > 0: compare_logits = torch.concat(compare_logits, dim=0).to(logits.dtype) compare_pred = self.compare(compare_logits) loss_compare = self.loss_fn(compare_pred, torch.tensor(compare_labels, dtype=torch.long, device=compare_logits.device)).mean() if torch.all(loss_compare == 0): entry_logits = logits[0][0].unsqueeze(0) compare_logits = torch.concat([entry_logits, entry_logits, entry_logits - entry_logits], dim=1) compare_pred = self.compare(compare_logits) compare_labels = [1] loss_compare = self.loss_fn(compare_pred, torch.tensor(compare_labels, dtype=torch.long, device=compare_logits.device)).mean() # if decode: # lm_labels = target_ids_s2s.index_select(0, (target_ids_s2s.sum(-1) > 0).nonzero().view(-1)[0]) # lm_labels = lm_labels.repeat(logits.shape[0], 1).clone().view(-1) # lm_labels = target_ids_s2s.clone() # target_ids_s2s = shift_tokens_right(target_ids_s2s, 0, 1) # target_ids_s2s.masked_fill_(target_ids_s2s==0, 3) if task == 0: _mask_index = (target_ids_s2s > 0).view(-1).nonzero().view(-1) lm_logits_ = flatten_states(lm_logits, _mask_index) lm_pred = self.lm_head(lm_logits_, ebd_weight).float() lm_labels = target_ids_s2s.clone().reshape(-1) lm_labels = lm_labels.index_select(0, _mask_index) # lm_pred = torch.nn.functional.log_softmax(lm_pred) # lm_loss = torch.nn.functional.nll_loss(lm_pred, lm_labels.long()) lm_loss = self.loss_fn(lm_pred, lm_labels.long()) # dot = register_hooks(lm_loss) # lm_loss.backward() # dot().save('tmp.dot') _mask_index = (target_ids > 0).view(-1).nonzero().view(-1) mlm_logits = flatten_states(logits, _mask_index) mlm_pred = self.lm_head(mlm_logits, ebd_weight).float() mlm_labels = target_ids.view(-1) mlm_labels = mlm_labels.index_select(0, _mask_index) mlm_loss = self.loss_fn(mlm_pred, mlm_labels.long()) return loss_compare, mlm_loss, lm_loss class WywLM(torch.nn.Module): """ DeBERTa encoder This module is composed of the input embedding layer with stacked transformer layers with disentangled attention. Parameters: config: A model config class instance with the configuration to build a new model. The schema is similar to `BertConfig`, \ for more details, please refer :class:`~DeBERTa.deberta.ModelConfig` pre_trained: The pre-trained DeBERTa model, it can be a physical path of a pre-trained DeBERTa model or a released configurations, \ i.e. [**base, large, base_mnli, large_mnli**] """ def __init__(self, config=None, pre_trained=None): super().__init__() state = None if pre_trained is not None: state, model_config = load_model_state(pre_trained) if config is not None and model_config is not None: for k in config.__dict__: if k not in ['hidden_size', 'intermediate_size', 'num_attention_heads', 'num_hidden_layers', 'vocab_size', 'max_position_embeddings']: model_config.__dict__[k] = config.__dict__[k] config = copy.copy(model_config) self.embeddings = BertEmbeddings(config) self.encoder = BertEncoder(config) self.config = config self.pre_trained = pre_trained self.apply_state(state) def forward(self, input_ids, attention_mask=None, token_type_ids=None, output_all_encoded_layers=True, position_ids = None, return_att = False): """ Args: input_ids: a torch.LongTensor of shape [batch_size, sequence_length] \ with the word token indices in the vocabulary attention_mask: an optional parameter for input mask or attention mask. - If it's an input mask, then it will be torch.LongTensor of shape [batch_size, sequence_length] with indices \ selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max \ input sequence length in the current batch. It's the mask that we typically use for attention when \ a batch has varying length sentences. - If it's an attention mask then it will be torch.LongTensor of shape [batch_size, sequence_length, sequence_length]. \ In this case, it's a mask indicate which tokens in the sequence should be attended by other tokens in the sequence. token_type_ids: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token \ types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to \ a `sentence B` token (see BERT paper for more details). output_all_encoded_layers: whether to output results of all encoder layers, default, True Returns: - The output of the stacked transformer layers if `output_all_encoded_layers=True`, else \ the last layer of stacked transformer layers - Attention matrix of self-attention layers if `return_att=True` Example:: # Batch of wordPiece token ids. # Each sample was padded with zero to the maxium length of the batch input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]]) # Mask of valid input ids attention_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]]) # DeBERTa model initialized with pretrained base model bert = DeBERTa(pre_trained='base') encoder_layers = bert(input_ids, attention_mask=attention_mask) """ if attention_mask is None: attention_mask = torch.ones_like(input_ids) if token_type_ids is None: token_type_ids = torch.zeros_like(input_ids) token_mask = torch.ones_like(input_ids) else: idxs = torch.flip(torch.cumsum(torch.flip(token_type_ids, [-1]), axis=1), [-1]) token_mask = idxs > 0 token_mask = token_mask.byte() ebd_output = self.embeddings(input_ids.to(torch.long), token_type_ids.to(torch.long), position_ids, token_mask) embedding_output = ebd_output['embeddings'] encoder_output = self.encoder(embedding_output, attention_mask, output_all_encoded_layers=output_all_encoded_layers, return_att = return_att) encoder_output.update(ebd_output) return encoder_output def apply_state(self, state = None): """ Load state from previous loaded model state dictionary. Args: state (:obj:`dict`, optional): State dictionary as the state returned by torch.module.state_dict(), default: `None`. \ If it's `None`, then will use the pre-trained state loaded via the constructor to re-initialize \ the `DeBERTa` model """ if self.pre_trained is None and state is None: return if state is None: state, config = load_model_state(self.pre_trained) self.config = config prefix = '' for k in state: if 'embeddings.' in k: if not k.startswith('embeddings.'): prefix = k[:k.index('embeddings.')] break missing_keys = [] unexpected_keys = [] error_msgs = [] self._load_from_state_dict(state, prefix = prefix, local_metadata=None, strict=True, missing_keys=missing_keys, unexpected_keys=unexpected_keys, error_msgs=error_msgs) class MaskedLanguageModel(NNModule): """ Masked language model """ def __init__(self, config, *wargs, **kwargs): super().__init__(config) self.backbone = WywLM(config) self.max_relative_positions = getattr(config, 'max_relative_positions', -1) self.position_buckets = getattr(config, 'position_buckets', -1) if self.max_relative_positions <1: self.max_relative_positions = config.max_position_embeddings # self.mlm_predictions = UGDecoder(self.backbone.config, self.backbone.embeddings.word_embeddings.weight.size(0)) self.lm_predictions = UGDecoder(self.backbone.config, self.backbone.embeddings.word_embeddings.weight.size(0)) self.device = None self.loss = WywLMLoss(config) # self.loss_lm = WywLMLoss(config) self.apply(self.init_weights) def forward(self, samples, position_ids=None): task = samples['task'] if task == 0: input_ids = samples['s2s_input_ids'] type_ids = samples['s2s_token_type_ids'] attention_mask = samples['s2s_attention_mask'] labels = samples['s2s_masked_lm_labels'] dict_pos = samples['dict_pos'] s2s_label = samples['s2s_label'] else: input_ids = samples['input_ids'] type_ids = samples['token_type_ids'] attention_mask = samples['attention_mask'] labels = samples['masked_lm_labels'] dict_pos = samples['dict_pos'] s2s_label = samples['s2s_label'] if self.device is None: self.device = list(self.parameters())[0].device input_ids = input_ids.to(self.device) type_ids = None lm_labels = labels.to(self.device) s2s_label = s2s_label.to(self.device) attention_mask = attention_mask.to(self.device) encoder_output = self.backbone(input_ids, attention_mask, type_ids, output_all_encoded_layers=True, position_ids = position_ids) encoder_layers = encoder_output['hidden_states'] z_states = encoder_output['position_embeddings'] ctx_layer = encoder_layers[-1] mlm_loss = torch.tensor(0).to(ctx_layer).float() lm_loss = torch.tensor(0).to(ctx_layer).float() lm_logits = None label_inputs = None loss = torch.tensor(0).to(ctx_layer).float() loss_compare = torch.tensor(0).to(ctx_layer).float() ebd_weight = self.backbone.embeddings.word_embeddings.weight lm_logits, mlm_logits = self.lm_predictions(encoder_layers, self.backbone.embeddings.word_embeddings, input_ids, z_states, attention_mask, self.backbone.encoder, target_ids=lm_labels) # if lm_labels.detach().sum() != 0: loss_compare, mlm_loss, lm_loss = self.loss(mlm_logits, lm_logits, lm_labels, dict_pos, target_ids_s2s=s2s_label, decode=False, ebd_weight=ebd_weight, input_ids=input_ids, task=task) loss = loss_compare * 10 + mlm_loss + lm_loss # if s2s_label.detach().sum() != 0: # s2s_idx = (s2s_label.sum(-1)>0).nonzero().view(-1) # s2s_label = s2s_label.index_select(0, s2s_idx) # # ebd_weight = self.backbone.embeddings.word_embeddings.weight # # lm_logits = self.lm_predictions(encoder_layers[-3], self.backbone.embeddings.word_embeddings, # # input_ids.index_select(0, s2s_idx), z_states.index_select(0, s2s_idx), # # attention_mask.index_select(0, s2s_idx), self.backbone.encoder, # # target_ids=s2s_label, # # decode=True, s2s_idx=s2s_idx) # # lm_logits = encoder_layers[-1].detach().index_select(0, s2s_idx) # _, lm_loss = self.loss_lm(lm_logits, # s2s_label, # torch.zeros_like(dict_pos), # decode=True, # ebd_weight=ebd_weight, # input_ids=input_ids.index_select(0, s2s_idx)) # lm_loss = lm_logits.max() # loss = loss + lm_loss return { 'logits' : lm_logits, 'labels' : lm_labels, 's2s_label': s2s_label, 'loss' : loss.float(), 'loss_compare': loss_compare.float(), 'lm_loss': lm_loss.float(), 'mlm_loss': mlm_loss.float() }