text: # Copyright (c) 2020, Salesforce.com, Inc.
  # Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
  # Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
  #
  # 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.
  
  from __future__ import absolute_import, division, print_function, unicode_literals
  
  import json
  import logging
  import math
  import os
  import sys
  from io import open
  
  import torch
  from torch import nn
  from torch.nn import CrossEntropyLoss, MSELoss
  
  
  from pytorch_transformers.modeling_utils import (WEIGHTS_NAME, CONFIG_NAME, PretrainedConfig, 
                                                   PreTrainedModel, prune_linear_layer, add_start_docstrings)
  from pytorch_transformers.modeling_bert import BertPreTrainedModel, BertModel, BertLayer, BertPooler
  
  
  class BertPointer(BertPreTrainedModel):
      def __init__(self, config):
          super(BertPointer, self).__init__(config)
          self.num_labels = config.num_labels
  
          self.bert = BertModel(config)
          self.dropout = nn.Dropout(config.hidden_dropout_prob)
  
          # classifiers
          self.ext_start_classifier = nn.Linear(config.hidden_size, 1, bias=False)
          self.ext_end_classifier = nn.Linear(config.hidden_size, 1, bias=False)
          self.aug_start_classifier = nn.Linear(config.hidden_size, 1, bias=False)
          self.aug_end_classifier = nn.Linear(config.hidden_size, 1, bias=False)
  
          self.label_classifier = nn.Linear(config.hidden_size, self.config.num_labels)
  
          self.apply(self.init_weights)
  
      def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None,
                  position_ids=None, head_mask=None, 
                  ext_mask=None, ext_start_labels=None, ext_end_labels=None,
                  aug_mask=None, aug_start_labels=None, aug_end_labels=None,
                  loss_lambda=1.0):
          # run through bert
          bert_outputs = self.bert(input_ids, position_ids=position_ids, token_type_ids=token_type_ids,
                                   attention_mask=attention_mask, head_mask=head_mask)
  
          # label classifier
          pooled_output = bert_outputs[1]
          pooled_output = self.dropout(pooled_output)
          label_logits = self.label_classifier(pooled_output)
  
          # extraction classifier
          output = bert_outputs[0]
          ext_mask = ext_mask.unsqueeze(-1)
          ext_start_logits = self.ext_start_classifier(output) * ext_mask
          ext_end_logits = self.ext_end_classifier(output) * ext_mask
  
          # augmentation classifier
          output = bert_outputs[0]
          aug_mask = aug_mask.unsqueeze(-1)
          aug_start_logits = self.aug_start_classifier(output) * aug_mask
          aug_end_logits = self.aug_end_classifier(output) * aug_mask
  
          span_logits = (ext_start_logits, ext_end_logits, aug_start_logits, aug_end_logits,)
          outputs = (label_logits,) + span_logits + bert_outputs[2:]
  
          if labels is not None and \
                  ext_start_labels is not None and ext_end_labels is not None and \
                  aug_start_labels is not None and aug_end_labels is not None:
              if self.num_labels == 1:
                  #  We are doing regression
                  loss_fct = MSELoss()
                  loss = loss_fct(label_logits.view(-1), labels.view(-1))
              else:
                  loss_fct = CrossEntropyLoss()
  
                  # label loss
                  labels_loss = loss_fct(label_logits.view(-1, self.num_labels), labels.view(-1))
  
                  # extraction loss
                  ext_start_loss = loss_fct(ext_start_logits.squeeze(), ext_start_labels)
                  ext_end_loss = loss_fct(ext_end_logits.squeeze(), ext_end_labels)
  
                  # augmentation loss
                  aug_start_loss = loss_fct(aug_start_logits.squeeze(), aug_start_labels)
                  aug_end_loss = loss_fct(aug_end_logits.squeeze(), aug_end_labels)
  
                  span_loss = (ext_start_loss + ext_end_loss + aug_start_loss + aug_end_loss) / 4
  
                  # combined loss
                  loss = labels_loss + loss_lambda * span_loss
  
              outputs = (loss, labels_loss, span_loss, ext_start_loss, ext_end_loss, aug_start_loss, aug_end_loss) + outputs
  
          return outputs  # (loss), (logits), (hidden_states), (attentions)
  