Spaces:
Running
Running
| # -*- coding: utf-8 -*- | |
| # @Time : 2022/4/16 12:10 下午 | |
| # @Author : JianingWang | |
| # @File : multiple_choice.py | |
| import torch | |
| from torch import nn | |
| from torch.nn import CrossEntropyLoss | |
| import torch.nn.functional as F | |
| # from transformers import MegatronBertPreTrainedModel, MegatronBertModel | |
| from transformers.models.megatron_bert import MegatronBertPreTrainedModel, MegatronBertModel | |
| from transformers.modeling_outputs import MultipleChoiceModelOutput | |
| class MegatronBertForMultipleChoice(MegatronBertPreTrainedModel): | |
| def __init__(self, config): | |
| super().__init__(config) | |
| self.bert = MegatronBertModel(config) | |
| # classifier_dropout = ( | |
| # config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob | |
| # ) | |
| classifier_dropout = 0.2 | |
| self.dropout = nn.Dropout(classifier_dropout) | |
| self.classifier = nn.Linear(config.hidden_size, 1) | |
| # Initialize weights and apply final processing | |
| self.post_init() | |
| def forward( | |
| self, | |
| input_ids=None, | |
| attention_mask=None, | |
| token_type_ids=None, | |
| position_ids=None, | |
| head_mask=None, | |
| inputs_embeds=None, | |
| labels=None, | |
| output_attentions=None, | |
| output_hidden_states=None, | |
| return_dict=None, | |
| pseudo=None | |
| ): | |
| r""" | |
| labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): | |
| Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., | |
| num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See | |
| `input_ids` above) | |
| """ | |
| return_dict = return_dict if return_dict is not None else self.config.use_return_dict | |
| num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1] | |
| input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None | |
| attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None | |
| token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None | |
| position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None | |
| inputs_embeds = ( | |
| inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1)) | |
| if inputs_embeds is not None | |
| else None | |
| ) | |
| outputs = self.bert( | |
| input_ids, | |
| attention_mask=attention_mask, | |
| token_type_ids=token_type_ids, | |
| position_ids=position_ids, | |
| head_mask=head_mask, | |
| inputs_embeds=inputs_embeds, | |
| output_attentions=output_attentions, | |
| output_hidden_states=output_hidden_states, | |
| return_dict=return_dict, | |
| ) | |
| pooled_output = outputs[1] # [batch_size, num_choices, hidden_dim] | |
| pooled_output = self.dropout(pooled_output) | |
| logits = self.classifier(pooled_output) # [batch_size, num_choices, 1] | |
| reshaped_logits = logits.view(-1, num_choices) # [batch_size, num_choices] | |
| loss = None | |
| if labels is not None: | |
| if pseudo is None: | |
| loss_fct = CrossEntropyLoss() | |
| loss = loss_fct(reshaped_logits, labels) | |
| else: | |
| loss_fct = CrossEntropyLoss(reduction="none") | |
| loss = loss_fct(reshaped_logits, labels) | |
| weight = 1 - pseudo * 0.9 | |
| loss *= weight | |
| loss = loss.mean() | |
| if not return_dict: | |
| output = (reshaped_logits,) + outputs[2:] | |
| return ((loss,) + output) if loss is not None else output | |
| return MultipleChoiceModelOutput( | |
| loss=loss, | |
| logits=reshaped_logits, | |
| hidden_states=outputs.hidden_states, | |
| attentions=outputs.attentions, | |
| ) | |
| class MegatronBertRDropForMultipleChoice(MegatronBertPreTrainedModel): | |
| def __init__(self, config): | |
| super().__init__(config) | |
| self.bert = MegatronBertModel(config) | |
| # classifier_dropout = ( | |
| # config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob | |
| # ) | |
| classifier_dropout = 0.2 | |
| self.dropout = nn.Dropout(classifier_dropout) | |
| self.classifier = nn.Linear(config.hidden_size, 1) | |
| # Initialize weights and apply final processing | |
| self.post_init() | |
| def forward( | |
| self, | |
| input_ids=None, | |
| attention_mask=None, | |
| token_type_ids=None, | |
| position_ids=None, | |
| head_mask=None, | |
| inputs_embeds=None, | |
| labels=None, | |
| output_attentions=None, | |
| output_hidden_states=None, | |
| return_dict=None, | |
| ): | |
| r""" | |
| labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): | |
| Labels for computing the multiple choice classification loss. Indices should be in `[0, ..., | |
| num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See | |
| `input_ids` above) | |
| """ | |
| return_dict = return_dict if return_dict is not None else self.config.use_return_dict | |
| num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1] | |
| input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None | |
| attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None | |
| token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None | |
| position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None | |
| inputs_embeds = ( | |
| inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1)) | |
| if inputs_embeds is not None | |
| else None | |
| ) | |
| logits_list = [] | |
| for i in range(2): | |
| outputs = self.bert( | |
| input_ids, | |
| attention_mask=attention_mask, | |
| token_type_ids=token_type_ids, | |
| position_ids=position_ids, | |
| head_mask=head_mask, | |
| inputs_embeds=inputs_embeds, | |
| output_attentions=output_attentions, | |
| output_hidden_states=output_hidden_states, | |
| return_dict=return_dict, | |
| ) | |
| pooled_output = outputs[1] | |
| pooled_output = self.dropout(pooled_output) | |
| logits = self.classifier(pooled_output) | |
| logits_list.append(logits.view(-1, num_choices)) | |
| loss = None | |
| alpha = 1.0 | |
| for logits in logits_list: | |
| if labels is not None: | |
| loss_fct = CrossEntropyLoss() | |
| l = loss_fct(logits, labels) | |
| if loss: | |
| loss += alpha * l | |
| else: | |
| loss = alpha * l | |
| if loss is not None: | |
| p = torch.log_softmax(logits_list[0], dim=-1) | |
| p_tec = torch.exp(p) | |
| q = torch.log_softmax(logits_list[-1], dim=-1) | |
| q_tec = torch.exp(q) | |
| kl_loss = F.kl_div(p, q_tec, reduction="none").sum() | |
| reverse_kl_loss = F.kl_div(q, p_tec, reduction="none").sum() | |
| loss += 0.5 * (kl_loss + reverse_kl_loss) / 2. | |
| return MultipleChoiceModelOutput( | |
| loss=loss, | |
| logits=logits_list[0], | |
| hidden_states=None, | |
| attentions=None | |
| ) | |