
"""simple docstring"""

from dataclasses import dataclass
from typing import Optional, Tuple

import torch
from torch import nn
from transformers import RobertaPreTrainedModel, XLMRobertaConfig, XLMRobertaModel
from transformers.utils import ModelOutput

@dataclass
class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):
     lowercase		     =		None
     lowercase		     =		None
     lowercase		     =		None
     lowercase		     =		None


class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):





     def __init__(			self			,  __UpperCAmelCase=1			,  __UpperCAmelCase=0			,  __UpperCAmelCase=2			,  __UpperCAmelCase=512			,  __UpperCAmelCase="cls"			,  __UpperCAmelCase=False			,  __UpperCAmelCase=True			,  **__UpperCAmelCase			,  ):

      '''simple docstring'''




      super().__init__(pad_token_id=__UpperCAmelCase			,  bos_token_id=__UpperCAmelCase			,  eos_token_id=__UpperCAmelCase			,  **__UpperCAmelCase   )
      __UpperCamelCase       						= project_dim
      __UpperCamelCase       						= pooler_fn
      __UpperCamelCase       						= learn_encoder
      __UpperCamelCase       						= use_attention_mask


class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):
     lowercase		     =		[r"pooler", r"logit_scale"]
     lowercase		     =		[r"position_ids", r"predictions.decoder.bias"]
     lowercase		     =		"roberta"
     lowercase		     =		RobertaSeriesConfig





     def __init__(			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      super().__init__(__UpperCAmelCase   )
      __UpperCamelCase       						= XLMRobertaModel(__UpperCAmelCase   )
      __UpperCamelCase       						= nn.Linear(config.hidden_size			,  config.project_dim   )
      __UpperCamelCase       						= getattr(__UpperCAmelCase			,  'has_pre_transformation'			,  __UpperCAmelCase   )
      if self.has_pre_transformation:
       __UpperCamelCase       						= nn.Linear(config.hidden_size			,  config.project_dim   )
       __UpperCamelCase       						= nn.LayerNorm(config.hidden_size			,  eps=config.layer_norm_eps   )
      self.post_init()







     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase = None			,  __UpperCAmelCase = None			,  __UpperCAmelCase = None			,  __UpperCAmelCase = None			,  __UpperCAmelCase = None			,  __UpperCAmelCase = None			,  __UpperCAmelCase = None			,  __UpperCAmelCase = None			,  __UpperCAmelCase = None			,  __UpperCAmelCase = None			,  __UpperCAmelCase = None			,  ):

      '''simple docstring'''




      __UpperCamelCase       						= return_dict if return_dict is not None else self.config.use_return_dict

      __UpperCamelCase       						= self.base_model(
          input_ids=__UpperCAmelCase			,  attention_mask=__UpperCAmelCase			,  token_type_ids=__UpperCAmelCase			,  position_ids=__UpperCAmelCase			,  head_mask=__UpperCAmelCase			,  inputs_embeds=__UpperCAmelCase			,  encoder_hidden_states=__UpperCAmelCase			,  encoder_attention_mask=__UpperCAmelCase			,  output_attentions=__UpperCAmelCase			,  output_hidden_states=True if self.has_pre_transformation else output_hidden_states			,  return_dict=__UpperCAmelCase			,  )

      if self.has_pre_transformation:
       __UpperCamelCase       						= outputs['hidden_states'][-2]
       __UpperCamelCase       						= self.pre_LN(__UpperCAmelCase   )
       __UpperCamelCase       						= self.transformation_pre(__UpperCAmelCase   )

       return TransformationModelOutput(
           projection_state=__UpperCAmelCase			,  last_hidden_state=outputs.last_hidden_state			,  hidden_states=outputs.hidden_states			,  attentions=outputs.attentions			,  )
      else:
       __UpperCamelCase       						= self.transformation(outputs.last_hidden_state   )
       return TransformationModelOutput(
           projection_state=__UpperCAmelCase			,  last_hidden_state=outputs.last_hidden_state			,  hidden_states=outputs.hidden_states			,  attentions=outputs.attentions			,  )





"""simple docstring"""

import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.bert.modeling_bert import (
    BERT_INPUTS_DOCSTRING,
    BERT_START_DOCSTRING,
    BertEmbeddings,
    BertLayer,
    BertPooler,
    BertPreTrainedModel,
)






def 							A		(       snake_case   :Optional[int]	)   ->			str:
 __UpperCamelCase       						= torch.exp(snake_case	)
 __UpperCamelCase       						= torch.sum(snake_case       ,      dim=1	)  # sum of exp(x_i)
 __UpperCamelCase       						= torch.sum(x * exp_x       ,      dim=1	)  # sum of x_i * exp(x_i)
 return torch.log(snake_case	) - B / A

class __lowerCAmelCase	(   nn.Module    ):





     def __init__(			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      super().__init__()
      __UpperCamelCase       						= config.output_attentions
      __UpperCamelCase       						= config.output_hidden_states
      __UpperCamelCase       						= nn.ModuleList([BertLayer(__UpperCAmelCase   ) for _ in range(config.num_hidden_layers   )]   )
      __UpperCamelCase       						= nn.ModuleList([BertHighway(__UpperCAmelCase   ) for _ in range(config.num_hidden_layers   )]   )

      __UpperCamelCase       						= [-1 for _ in range(config.num_hidden_layers   )]





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      if (type(__UpperCAmelCase   ) is float) or (type(__UpperCAmelCase   ) is int):
       for i in range(len(self.early_exit_entropy   )   ):
        __UpperCamelCase       						= x
      else:
       __UpperCamelCase       						= x





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      __UpperCamelCase       						= pooler.state_dict()
      for highway in self.highway:
       for name, param in highway.pooler.state_dict().items():
        param.copy_(loaded_model[name]   )







     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  ):

      '''simple docstring'''




      __UpperCamelCase       						= ()
      __UpperCamelCase       						= ()
      __UpperCamelCase       						= ()
      for i, layer_module in enumerate(self.layer   ):
       if self.output_hidden_states:
        __UpperCamelCase       						= all_hidden_states + (hidden_states,)

       __UpperCamelCase       						= layer_module(
           __UpperCAmelCase			,  __UpperCAmelCase			,  head_mask[i]			,  __UpperCAmelCase			,  __UpperCAmelCase   )
       __UpperCamelCase       						= layer_outputs[0]

       if self.output_attentions:
        __UpperCamelCase       						= all_attentions + (layer_outputs[1],)

       __UpperCamelCase       						= (hidden_states,)
       if self.output_hidden_states:
        __UpperCamelCase       						= current_outputs + (all_hidden_states,)
       if self.output_attentions:
        __UpperCamelCase       						= current_outputs + (all_attentions,)

       __UpperCamelCase       						= self.highway[i](__UpperCAmelCase   )
       # logits, pooled_output

       if not self.training:
        __UpperCamelCase       						= highway_exit[0]
        __UpperCamelCase       						= entropy(__UpperCAmelCase   )
        __UpperCamelCase       						= highway_exit + (highway_entropy,)  # logits, hidden_states(?), entropy
        __UpperCamelCase       						= all_highway_exits + (highway_exit,)

        if highway_entropy < self.early_exit_entropy[i]:
         __UpperCamelCase       						= (highway_logits,) + current_outputs[1:] + (all_highway_exits,)
         raise HighwayException(__UpperCAmelCase			,  i + 1   )
       else:
        __UpperCamelCase       						= all_highway_exits + (highway_exit,)

        # Add last layer
      if self.output_hidden_states:
       __UpperCamelCase       						= all_hidden_states + (hidden_states,)

      __UpperCamelCase       						= (hidden_states,)
      if self.output_hidden_states:
       __UpperCamelCase       						= outputs + (all_hidden_states,)
      if self.output_attentions:
       __UpperCamelCase       						= outputs + (all_attentions,)

      __UpperCamelCase       						= outputs + (all_highway_exits,)
      return outputs  # last-layer hidden state, (all hidden states), (all attentions), all highway exits


@add_start_docstrings(
    "The Bert Model transformer with early exiting (DeeBERT). "    ,						__SCREAMING_SNAKE_CASE    ,						)
class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):





     def __init__(			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      super().__init__(__UpperCAmelCase   )
      __UpperCamelCase       						= config

      __UpperCamelCase       						= BertEmbeddings(__UpperCAmelCase   )
      __UpperCamelCase       						= DeeBertEncoder(__UpperCAmelCase   )
      __UpperCamelCase       						= BertPooler(__UpperCAmelCase   )

      self.init_weights()





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      self.encoder.init_highway_pooler(self.pooler   )





     def 		UpperCAmelCase      (			self   ):

      '''simple docstring'''




      return self.embeddings.word_embeddings





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      __UpperCamelCase       						= value





     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      for layer, heads in heads_to_prune.items():
       self.encoder.layer[layer].attention.prune_heads(__UpperCAmelCase   )







     @add_start_docstrings_to_model_forward(__UpperCAmelCase   )
     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  ):

      '''simple docstring'''




      if input_ids is not None and inputs_embeds is not None:
       raise ValueError('You cannot specify both input_ids and inputs_embeds at the same time'   )
      elif input_ids is not None:
       __UpperCamelCase       						= input_ids.size()
      elif inputs_embeds is not None:
       __UpperCamelCase       						= inputs_embeds.size()[:-1]
      else:
       raise ValueError('You have to specify either input_ids or inputs_embeds'   )

      __UpperCamelCase       						= input_ids.device if input_ids is not None else inputs_embeds.device

      if attention_mask is None:
       __UpperCamelCase       						= torch.ones(__UpperCAmelCase			,  device=__UpperCAmelCase   )
      if encoder_attention_mask is None:
       __UpperCamelCase       						= torch.ones(__UpperCAmelCase			,  device=__UpperCAmelCase   )
      if token_type_ids is None:
       __UpperCamelCase       						= torch.zeros(__UpperCAmelCase			,  dtype=torch.long			,  device=__UpperCAmelCase   )

      # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
      # ourselves in which case we just need to make it broadcastable to all heads.
      __UpperCamelCase       						= self.get_extended_attention_mask(__UpperCAmelCase			,  __UpperCAmelCase			,  __UpperCAmelCase   )

      # If a 2D ou 3D attention mask is provided for the cross-attention
      # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
      if encoder_attention_mask.dim() == 3:
       __UpperCamelCase       						= encoder_attention_mask[:, None, :, :]
      if encoder_attention_mask.dim() == 2:
       __UpperCamelCase       						= encoder_attention_mask[:, None, None, :]

      __UpperCamelCase       						= encoder_extended_attention_mask.to(
          dtype=next(self.parameters()   ).dtype   )  # fp16 compatibility
      __UpperCamelCase       						= (1.0 - encoder_extended_attention_mask) * -1_0_0_0_0.0

      # Prepare head mask if needed
      # 1.0 in head_mask indicate we keep the head
      # attention_probs has shape bsz x n_heads x N x N
      # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
      # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
      __UpperCamelCase       						= self.get_head_mask(__UpperCAmelCase			,  self.config.num_hidden_layers   )

      __UpperCamelCase       						= self.embeddings(
          input_ids=__UpperCAmelCase			,  position_ids=__UpperCAmelCase			,  token_type_ids=__UpperCAmelCase			,  inputs_embeds=__UpperCAmelCase   )
      __UpperCamelCase       						= self.encoder(
          __UpperCAmelCase			,  attention_mask=__UpperCAmelCase			,  head_mask=__UpperCAmelCase			,  encoder_hidden_states=__UpperCAmelCase			,  encoder_attention_mask=__UpperCAmelCase			,  )
      __UpperCamelCase       						= encoder_outputs[0]
      __UpperCamelCase       						= self.pooler(__UpperCAmelCase   )

      __UpperCamelCase       						= (
          sequence_output,
          pooled_output,
      ) + encoder_outputs[
          1:
      ]  # add hidden_states and attentions if they are here
      return outputs  # sequence_output, pooled_output, (hidden_states), (attentions), highway exits


class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):





     def __init__(			self			,  __UpperCAmelCase			,  __UpperCAmelCase   ):

      '''simple docstring'''




      __UpperCamelCase       						= message
      __UpperCamelCase       						= exit_layer  # start from 1!


class __lowerCAmelCase	(   nn.Module    ):





     def __init__(			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      super().__init__()
      __UpperCamelCase       						= BertPooler(__UpperCAmelCase   )
      __UpperCamelCase       						= nn.Dropout(config.hidden_dropout_prob   )
      __UpperCamelCase       						= nn.Linear(config.hidden_size			,  config.num_labels   )







     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      __UpperCamelCase       						= encoder_outputs[0]
      __UpperCamelCase       						= self.pooler(__UpperCAmelCase   )
      # "return" pooler_output

      # BertModel
      __UpperCamelCase       						= (pooler_input, pooler_output) + encoder_outputs[1:]
      # "return" bmodel_output

      # Dropout and classification
      __UpperCamelCase       						= bmodel_output[1]

      __UpperCamelCase       						= self.dropout(__UpperCAmelCase   )
      __UpperCamelCase       						= self.classifier(__UpperCAmelCase   )

      return logits, pooled_output


@add_start_docstrings(
    "Bert Model (with early exiting - DeeBERT) with a classifier on top,\n    also takes care of multi-layer training. "    ,						__SCREAMING_SNAKE_CASE    ,						)
class __lowerCAmelCase	(   __SCREAMING_SNAKE_CASE    ):





     def __init__(			self			,  __UpperCAmelCase   ):

      '''simple docstring'''




      super().__init__(__UpperCAmelCase   )
      __UpperCamelCase       						= config.num_labels
      __UpperCamelCase       						= config.num_hidden_layers

      __UpperCamelCase       						= DeeBertModel(__UpperCAmelCase   )
      __UpperCamelCase       						= nn.Dropout(config.hidden_dropout_prob   )
      __UpperCamelCase       						= nn.Linear(config.hidden_size			,  self.config.num_labels   )

      self.init_weights()







     @add_start_docstrings_to_model_forward(__UpperCAmelCase   )
     def 		UpperCAmelCase      (			self			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=None			,  __UpperCAmelCase=-1			,  __UpperCAmelCase=False			,  ):

      '''simple docstring'''




      __UpperCamelCase       						= self.num_layers
      try:
       __UpperCamelCase       						= self.bert(
           __UpperCAmelCase			,  attention_mask=__UpperCAmelCase			,  token_type_ids=__UpperCAmelCase			,  position_ids=__UpperCAmelCase			,  head_mask=__UpperCAmelCase			,  inputs_embeds=__UpperCAmelCase			,  )
       # sequence_output, pooled_output, (hidden_states), (attentions), highway exits

       __UpperCamelCase       						= outputs[1]

       __UpperCamelCase       						= self.dropout(__UpperCAmelCase   )
       __UpperCamelCase       						= self.classifier(__UpperCAmelCase   )
       __UpperCamelCase       						= (logits,) + outputs[2:]  # add hidden states and attention if they are here
      except HighwayException as e:
       __UpperCamelCase       						= e.message
       __UpperCamelCase       						= e.exit_layer
       __UpperCamelCase       						= outputs[0]

      if not self.training:
       __UpperCamelCase       						= entropy(__UpperCAmelCase   )
       __UpperCamelCase       						= []
       __UpperCamelCase       						= []
      if labels is not None:
       if self.num_labels == 1:
        #  We are doing regression
        __UpperCamelCase       						= MSELoss()
        __UpperCamelCase       						= loss_fct(logits.view(-1   )			,  labels.view(-1   )   )
       else:
        __UpperCamelCase       						= CrossEntropyLoss()
        __UpperCamelCase       						= loss_fct(logits.view(-1			,  self.num_labels   )			,  labels.view(-1   )   )

       # work with highway exits
       __UpperCamelCase       						= []
       for highway_exit in outputs[-1]:
        __UpperCamelCase       						= highway_exit[0]
        if not self.training:
         highway_logits_all.append(__UpperCAmelCase   )
         highway_entropy.append(highway_exit[2]   )
        if self.num_labels == 1:
         #  We are doing regression
         __UpperCamelCase       						= MSELoss()
         __UpperCamelCase       						= loss_fct(highway_logits.view(-1   )			,  labels.view(-1   )   )
        else:
         __UpperCamelCase       						= CrossEntropyLoss()
         __UpperCamelCase       						= loss_fct(highway_logits.view(-1			,  self.num_labels   )			,  labels.view(-1   )   )
        highway_losses.append(__UpperCAmelCase   )

       if train_highway:
        __UpperCamelCase       						= (sum(highway_losses[:-1]   ),) + outputs
        # exclude the final highway, of course
       else:
        __UpperCamelCase       						= (loss,) + outputs
      if not self.training:
       __UpperCamelCase       						= outputs + ((original_entropy, highway_entropy), exit_layer)
       if output_layer >= 0:
        __UpperCamelCase       						= (
            (outputs[0],) + (highway_logits_all[output_layer],) + outputs[2:]
        )  # use the highway of the last layer

      return outputs  # (loss), logits, (hidden_states), (attentions), (highway_exits)



