
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 		_UpperCAmelCase   (	UpperCamelCase:     List[str]  ):






	"""simple docstring"""


	__lowerCAmelCase         =							torch.exp(UpperCamelCase  )
	__lowerCAmelCase         =							torch.sum(UpperCamelCase			,							dim=1  )  # sum of exp(x_i)
	__lowerCAmelCase         =							torch.sum(x * exp_x			,							dim=1  )  # sum of x_i * exp(x_i)
	return torch.log(UpperCamelCase  ) - B / A

class       a		(   nn.Module					):







							def __init__(						self  :  Tuple      ,				snake_case__  :  str ):
								"""simple docstring"""







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

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







							def       UpperCAmelCase__					(						self  :  List[Any]      ,				snake_case__  :  List[str] ):
								"""simple docstring"""







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







							def       UpperCAmelCase__					(						self  :  List[str]      ,				snake_case__  :  Tuple ):
								"""simple docstring"""







								__lowerCAmelCase         =							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  :  Union[str, Any]      ,				snake_case__  :  Union[str, Any]      ,				snake_case__  :  Dict=None      ,				snake_case__  :  Union[str, Any]=None      ,				snake_case__  :  str=None      ,				snake_case__  :  Optional[Any]=None      ,				):
								"""simple docstring"""







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

									__lowerCAmelCase         =							layer_module(
									    snake_case__      ,				snake_case__      ,				head_mask[i]      ,				snake_case__      ,				snake_case__ )
									__lowerCAmelCase         =							layer_outputs[0]

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

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

									__lowerCAmelCase         =							self.highway[i](snake_case__ )
									# logits, pooled_output

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

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

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

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

								__lowerCAmelCase         =							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). '      ,	__UpperCAmelCase      ,	)
class       a		(   __UpperCAmelCase					):







							def __init__(						self  :  int      ,				snake_case__  :  Optional[Any] ):
								"""simple docstring"""







								super().__init__(snake_case__ )
								__lowerCAmelCase         =							config

								__lowerCAmelCase         =							BertEmbeddings(snake_case__ )
								__lowerCAmelCase         =							DeeBertEncoder(snake_case__ )
								__lowerCAmelCase         =							BertPooler(snake_case__ )

								self.init_weights()







							def       UpperCAmelCase__					(						self  :  Dict ):
								"""simple docstring"""







								self.encoder.init_highway_pooler(self.pooler )







							def       UpperCAmelCase__					(						self  :  Any ):
								"""simple docstring"""







								return self.embeddings.word_embeddings







							def       UpperCAmelCase__					(						self  :  Union[str, Any]      ,				snake_case__  :  Optional[Any] ):
								"""simple docstring"""







								__lowerCAmelCase         =							value







							def       UpperCAmelCase__					(						self  :  Any      ,				snake_case__  :  Optional[int] ):
								"""simple docstring"""







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







							@add_start_docstrings_to_model_forward(snake_case__ )
							def       UpperCAmelCase__					(						self  :  str      ,				snake_case__  :  int=None      ,				snake_case__  :  Tuple=None      ,				snake_case__  :  Union[str, Any]=None      ,				snake_case__  :  int=None      ,				snake_case__  :  Optional[Any]=None      ,				snake_case__  :  List[str]=None      ,				snake_case__  :  str=None      ,				snake_case__  :  Union[str, Any]=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:
									__lowerCAmelCase         =							input_ids.size()
								elif inputs_embeds is not None:
									__lowerCAmelCase         =							inputs_embeds.size()[:-1]
								else:
									raise ValueError("You have to specify either input_ids or inputs_embeds" )

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

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

								# 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.
								__lowerCAmelCase         =							self.get_extended_attention_mask(snake_case__      ,				snake_case__      ,				snake_case__ )

								# 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:
									__lowerCAmelCase         =							encoder_attention_mask[:, None, :, :]
								if encoder_attention_mask.dim() == 2:
									__lowerCAmelCase         =							encoder_attention_mask[:, None, None, :]

								__lowerCAmelCase         =							encoder_extended_attention_mask.to(
								    dtype=next(self.parameters() ).dtype )  # fp16 compatibility
								__lowerCAmelCase         =							(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]
								__lowerCAmelCase         =							self.get_head_mask(snake_case__      ,				self.config.num_hidden_layers )

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

								__lowerCAmelCase         =							(
								    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       a		(   __UpperCAmelCase					):







							def __init__(						self  :  Dict      ,				snake_case__  :  int      ,				snake_case__  :  List[Any] ):
								"""simple docstring"""







								__lowerCAmelCase         =							message
								__lowerCAmelCase         =							exit_layer  # start from 1!

class       a		(   nn.Module					):







							def __init__(						self  :  Any      ,				snake_case__  :  Tuple ):
								"""simple docstring"""







								super().__init__()
								__lowerCAmelCase         =							BertPooler(snake_case__ )
								__lowerCAmelCase         =							nn.Dropout(config.hidden_dropout_prob )
								__lowerCAmelCase         =							nn.Linear(config.hidden_size      ,				config.num_labels )







							def       UpperCAmelCase__					(						self  :  int      ,				snake_case__  :  List[Any] ):
								"""simple docstring"""







								__lowerCAmelCase         =							encoder_outputs[0]
								__lowerCAmelCase         =							self.pooler(snake_case__ )
								# "return" pooler_output

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

								# Dropout and classification
								__lowerCAmelCase         =							bmodel_output[1]

								__lowerCAmelCase         =							self.dropout(snake_case__ )
								__lowerCAmelCase         =							self.classifier(snake_case__ )

								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. '      ,	__UpperCAmelCase      ,	)
class       a		(   __UpperCAmelCase					):







							def __init__(						self  :  Any      ,				snake_case__  :  Union[str, Any] ):
								"""simple docstring"""







								super().__init__(snake_case__ )
								__lowerCAmelCase         =							config.num_labels
								__lowerCAmelCase         =							config.num_hidden_layers

								__lowerCAmelCase         =							DeeBertModel(snake_case__ )
								__lowerCAmelCase         =							nn.Dropout(config.hidden_dropout_prob )
								__lowerCAmelCase         =							nn.Linear(config.hidden_size      ,				self.config.num_labels )

								self.init_weights()







							@add_start_docstrings_to_model_forward(snake_case__ )
							def       UpperCAmelCase__					(						self  :  Union[str, Any]      ,				snake_case__  :  str=None      ,				snake_case__  :  int=None      ,				snake_case__  :  Union[str, Any]=None      ,				snake_case__  :  Any=None      ,				snake_case__  :  List[str]=None      ,				snake_case__  :  int=None      ,				snake_case__  :  Union[str, Any]=None      ,				snake_case__  :  Dict=-1      ,				snake_case__  :  Tuple=False      ,				):
								"""simple docstring"""







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

									__lowerCAmelCase         =							outputs[1]

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

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

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

									if train_highway:
										__lowerCAmelCase         =							(sum(highway_losses[:-1] ),) + outputs
										# exclude the final highway, of course
									else:
										__lowerCAmelCase         =							(loss,) + outputs
								if not self.training:
									__lowerCAmelCase         =							outputs + ((original_entropy, highway_entropy), exit_layer)
									if output_layer >= 0:
										__lowerCAmelCase         =							(
										    (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)








from collections import OrderedDict
from typing import Mapping

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


UpperCamelCase_   	=	logging.get_logger(__name__)

UpperCamelCase_   	=	{
    "distilbert-base-uncased": "https://huggingface.co/distilbert-base-uncased/resolve/main/config.json",
    "distilbert-base-uncased-distilled-squad": (
        "https://huggingface.co/distilbert-base-uncased-distilled-squad/resolve/main/config.json"
    ),
    "distilbert-base-cased": "https://huggingface.co/distilbert-base-cased/resolve/main/config.json",
    "distilbert-base-cased-distilled-squad": (
        "https://huggingface.co/distilbert-base-cased-distilled-squad/resolve/main/config.json"
    ),
    "distilbert-base-german-cased": "https://huggingface.co/distilbert-base-german-cased/resolve/main/config.json",
    "distilbert-base-multilingual-cased": (
        "https://huggingface.co/distilbert-base-multilingual-cased/resolve/main/config.json"
    ),
    "distilbert-base-uncased-finetuned-sst-2-english": (
        "https://huggingface.co/distilbert-base-uncased-finetuned-sst-2-english/resolve/main/config.json"
    ),
}

class       a		(   __UpperCAmelCase					):
							lowercase_		:						str   							=  'distilbert'
							lowercase_		:						Any   							=  {
							    'hidden_size': 'dim',
							    'num_attention_heads': 'n_heads',
							    'num_hidden_layers': 'n_layers',
							}







							def __init__(						self  :  Optional[int]      ,				snake_case__  :  int=30_522      ,				snake_case__  :  str=512      ,				snake_case__  :  Tuple=False      ,				snake_case__  :  Tuple=6      ,				snake_case__  :  Any=12      ,				snake_case__  :  Dict=768      ,				snake_case__  :  Any=4 * 768      ,				snake_case__  :  Union[str, Any]=0.1      ,				snake_case__  :  str=0.1      ,				snake_case__  :  Tuple="gelu"      ,				snake_case__  :  str=0.0_2      ,				snake_case__  :  Optional[Any]=0.1      ,				snake_case__  :  Optional[Any]=0.2      ,				snake_case__  :  str=0      ,				**snake_case__  :  Dict      ,				):
								"""simple docstring"""







								__lowerCAmelCase         =							vocab_size
								__lowerCAmelCase         =							max_position_embeddings
								__lowerCAmelCase         =							sinusoidal_pos_embds
								__lowerCAmelCase         =							n_layers
								__lowerCAmelCase         =							n_heads
								__lowerCAmelCase         =							dim
								__lowerCAmelCase         =							hidden_dim
								__lowerCAmelCase         =							dropout
								__lowerCAmelCase         =							attention_dropout
								__lowerCAmelCase         =							activation
								__lowerCAmelCase         =							initializer_range
								__lowerCAmelCase         =							qa_dropout
								__lowerCAmelCase         =							seq_classif_dropout
								super().__init__(**snake_case__      ,				pad_token_id=snake_case__ )

class       a		(   __UpperCAmelCase					):







							@property
							def       UpperCAmelCase__					(						self  :  Tuple ):
								"""simple docstring"""







								if self.task == "multiple-choice":
									__lowerCAmelCase         =							{0: "batch", 1: "choice", 2: "sequence"}
								else:
									__lowerCAmelCase         =							{0: "batch", 1: "sequence"}
								return OrderedDict(
								    [
								        ("input_ids", dynamic_axis),
								        ("attention_mask", dynamic_axis),
								    ] )






