

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 (	__a      :			Optional[Any]	)			->     str:

					"""simple docstring"""

					_a	:    List[Any]         =		torch.exp(__a	)
					_a	:    Optional[Any]         =		torch.sum(__a		,dim=1	)  # sum of exp(x_i)
					_a	:    str         =		torch.sum(x * exp_x		,dim=1	)  # sum of x_i * exp(x_i)
					return torch.log(__a	) - B / A

class 				UpperCAmelCase_		(			nn.Module   ):
							"""simple docstring"""







							def __init__(			self       ,							_a       )    ->    Union[str, Any]:
												super().__init__()
												_a	:    List[str]         =		config.output_attentions
												_a	:    List[str]         =		config.output_hidden_states
												_a	:    Tuple         =		nn.ModuleList([BertLayer(_a       ) for _ in range(config.num_hidden_layers       )]       )
												_a	:    Tuple         =		nn.ModuleList([BertHighway(_a       ) for _ in range(config.num_hidden_layers       )]       )

												_a	:    Optional[Any]         =		[-1 for _ in range(config.num_hidden_layers       )]


							def 				__lowercase				(			self       ,							_a       )    ->    Any:
												if (type(_a       ) is float) or (type(_a       ) is int):
																	for i in range(len(self.early_exit_entropy       )       ):
																						_a	:    List[str]         =		x
												else:
																	_a	:    Tuple         =		x


							def 				__lowercase				(			self       ,							_a       )    ->    Any:
												_a	:    List[Any]         =		pooler.state_dict()
												for highway in self.highway:
																	for name, param in highway.pooler.state_dict().items():
																						param.copy_(loaded_model[name]       )




							def 				__lowercase				(			self       ,							_a       ,							_a=None       ,							_a=None       ,							_a=None       ,							_a=None       ,							)    ->    int:
												_a	:    Optional[int]         =		()
												_a	:    Any         =		()
												_a	:    int         =		()
												for i, layer_module in enumerate(self.layer       ):
																	if self.output_hidden_states:
																						_a	:    List[str]         =		all_hidden_states + (hidden_states,)

																	_a	:    List[Any]         =		layer_module(
																	    _a       ,							_a       ,							head_mask[i]       ,							_a       ,							_a       )
																	_a	:    Union[str, Any]         =		layer_outputs[0]

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

																	_a	:    Optional[int]         =		(hidden_states,)
																	if self.output_hidden_states:
																						_a	:    int         =		current_outputs + (all_hidden_states,)
																	if self.output_attentions:
																						_a	:    Tuple         =		current_outputs + (all_attentions,)

																	_a	:    Dict         =		self.highway[i](_a       )
																	# logits, pooled_output

																	if not self.training:
																						_a	:    Optional[int]         =		highway_exit[0]
																						_a	:    List[str]         =		entropy(_a       )
																						_a	:    Optional[Any]         =		highway_exit + (highway_entropy,)  # logits, hidden_states(?), entropy
																						_a	:    Optional[int]         =		all_highway_exits + (highway_exit,)

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

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

												_a	:    Union[str, Any]         =		(hidden_states,)
												if self.output_hidden_states:
																	_a	:    List[Any]         =		outputs + (all_hidden_states,)
												if self.output_attentions:
																	_a	:    Optional[int]         =		outputs + (all_attentions,)

												_a	:    Optional[Any]         =		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). "				,       __lowercase				,       )
class 				UpperCAmelCase_		(			__lowercase   ):
							"""simple docstring"""







							def __init__(			self       ,							_a       )    ->    List[str]:
												super().__init__(_a       )
												_a	:    Tuple         =		config

												_a	:    int         =		BertEmbeddings(_a       )
												_a	:    Dict         =		DeeBertEncoder(_a       )
												_a	:    str         =		BertPooler(_a       )

												self.init_weights()


							def 				__lowercase				(			self       )    ->    int:
												self.encoder.init_highway_pooler(self.pooler       )


							def 				__lowercase				(			self       )    ->    Union[str, Any]:
												return self.embeddings.word_embeddings


							def 				__lowercase				(			self       ,							_a       )    ->    Tuple:
												_a	:    Optional[int]         =		value


							def 				__lowercase				(			self       ,							_a       )    ->    Optional[Any]:
												for layer, heads in heads_to_prune.items():
																	self.encoder.layer[layer].attention.prune_heads(_a       )




							@add_start_docstrings_to_model_forward(_a       )
							def 				__lowercase				(			self       ,							_a=None       ,							_a=None       ,							_a=None       ,							_a=None       ,							_a=None       ,							_a=None       ,							_a=None       ,							_a=None       ,							)    ->    Tuple:
												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:
																	_a	:    Optional[int]         =		input_ids.size()
												elif inputs_embeds is not None:
																	_a	:    Optional[int]         =		inputs_embeds.size()[:-1]
												else:
																	raise ValueError('''You have to specify either input_ids or inputs_embeds'''       )

												_a	:    Optional[Any]         =		input_ids.device if input_ids is not None else inputs_embeds.device

												if attention_mask is None:
																	_a	:    Dict         =		torch.ones(_a       ,							device=_a       )
												if encoder_attention_mask is None:
																	_a	:    List[str]         =		torch.ones(_a       ,							device=_a       )
												if token_type_ids is None:
																	_a	:    Optional[int]         =		torch.zeros(_a       ,							dtype=torch.long       ,							device=_a       )

												# 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.
												_a	:    torch.Tensor         =		self.get_extended_attention_mask(_a       ,							_a       ,							_a       )

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

												_a	:    int         =		encoder_extended_attention_mask.to(
												    dtype=next(self.parameters()       ).dtype       )  # fp16 compatibility
												_a	:    Tuple         =		(1.0 - encoder_extended_attention_mask) * -1_0000.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]
												_a	:    str         =		self.get_head_mask(_a       ,							self.config.num_hidden_layers       )

												_a	:    List[str]         =		self.embeddings(
												    input_ids=_a       ,							position_ids=_a       ,							token_type_ids=_a       ,							inputs_embeds=_a       )
												_a	:    Optional[Any]         =		self.encoder(
												    _a       ,							attention_mask=_a       ,							head_mask=_a       ,							encoder_hidden_states=_a       ,							encoder_attention_mask=_a       ,							)
												_a	:    Union[str, Any]         =		encoder_outputs[0]
												_a	:    str         =		self.pooler(_a       )

												_a	:    Dict         =		(
												    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 				UpperCAmelCase_		(			__lowercase   ):
							"""simple docstring"""







							def __init__(			self       ,							_a       ,							_a       )    ->    Dict:
												_a	:    str         =		message
												_a	:    Optional[int]         =		exit_layer  # start from 1!



class 				UpperCAmelCase_		(			nn.Module   ):
							"""simple docstring"""







							def __init__(			self       ,							_a       )    ->    Any:
												super().__init__()
												_a	:    Optional[Any]         =		BertPooler(_a       )
												_a	:    int         =		nn.Dropout(config.hidden_dropout_prob       )
												_a	:    Union[str, Any]         =		nn.Linear(config.hidden_size       ,							config.num_labels       )




							def 				__lowercase				(			self       ,							_a       )    ->    Union[str, Any]:
												# Pooler
												_a	:    Any         =		encoder_outputs[0]
												_a	:    Any         =		self.pooler(_a       )
												# "return" pooler_output

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

												# Dropout and classification
												_a	:    Dict         =		bmodel_output[1]

												_a	:    Optional[Any]         =		self.dropout(_a       )
												_a	:    List[Any]         =		self.classifier(_a       )

												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. "				,       __lowercase				,       )
class 				UpperCAmelCase_		(			__lowercase   ):
							"""simple docstring"""







							def __init__(			self       ,							_a       )    ->    Any:
												super().__init__(_a       )
												_a	:    List[Any]         =		config.num_labels
												_a	:    str         =		config.num_hidden_layers

												_a	:    Tuple         =		DeeBertModel(_a       )
												_a	:    Tuple         =		nn.Dropout(config.hidden_dropout_prob       )
												_a	:    Dict         =		nn.Linear(config.hidden_size       ,							self.config.num_labels       )

												self.init_weights()




							@add_start_docstrings_to_model_forward(_a       )
							def 				__lowercase				(			self       ,							_a=None       ,							_a=None       ,							_a=None       ,							_a=None       ,							_a=None       ,							_a=None       ,							_a=None       ,							_a=-1       ,							_a=False       ,							)    ->    Optional[Any]:
												_a	:    List[Any]         =		self.num_layers
												try:
																	_a	:    str         =		self.bert(
																	    _a       ,							attention_mask=_a       ,							token_type_ids=_a       ,							position_ids=_a       ,							head_mask=_a       ,							inputs_embeds=_a       ,							)
																	# sequence_output, pooled_output, (hidden_states), (attentions), highway exits

																	_a	:    Any         =		outputs[1]

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

												if not self.training:
																	_a	:    Union[str, Any]         =		entropy(_a       )
																	_a	:    Optional[int]         =		[]
																	_a	:    Optional[int]         =		[]
												if labels is not None:
																	if self.num_labels == 1:
																						#  We are doing regression
																						_a	:    Dict         =		MSELoss()
																						_a	:    Optional[int]         =		loss_fct(logits.view(-1       )       ,							labels.view(-1       )       )
																	else:
																						_a	:    Union[str, Any]         =		CrossEntropyLoss()
																						_a	:    Optional[int]         =		loss_fct(logits.view(-1       ,							self.num_labels       )       ,							labels.view(-1       )       )

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

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


a__				=		logging.get_logger(__name__)

a__				=		{
    '''junnyu/roformer_chinese_small''': '''https://huggingface.co/junnyu/roformer_chinese_small/resolve/main/config.json''',
    '''junnyu/roformer_chinese_base''': '''https://huggingface.co/junnyu/roformer_chinese_base/resolve/main/config.json''',
    '''junnyu/roformer_chinese_char_small''': (
        '''https://huggingface.co/junnyu/roformer_chinese_char_small/resolve/main/config.json'''
    ),
    '''junnyu/roformer_chinese_char_base''': (
        '''https://huggingface.co/junnyu/roformer_chinese_char_base/resolve/main/config.json'''
    ),
    '''junnyu/roformer_small_discriminator''': (
        '''https://huggingface.co/junnyu/roformer_small_discriminator/resolve/main/config.json'''
    ),
    '''junnyu/roformer_small_generator''': (
        '''https://huggingface.co/junnyu/roformer_small_generator/resolve/main/config.json'''
    ),
    # See all RoFormer models at https://huggingface.co/models?filter=roformer
}

class 				UpperCAmelCase_		(			__lowercase   ):
							"""simple docstring"""







							UpperCAmelCase__       :						Optional[Any]	    =    "roformer"


							def __init__(			self       ,							_a=5_0_0_0_0       ,							_a=None       ,							_a=7_6_8       ,							_a=1_2       ,							_a=1_2       ,							_a=3_0_7_2       ,							_a="gelu"       ,							_a=0.1       ,							_a=0.1       ,							_a=1_5_3_6       ,							_a=2       ,							_a=0.02       ,							_a=1e-1_2       ,							_a=0       ,							_a=False       ,							_a=True       ,							**_a       ,							)    ->    List[str]:
												super().__init__(pad_token_id=_a       ,							**_a       )

												_a	:    Tuple         =		vocab_size
												_a	:    List[Any]         =		hidden_size if embedding_size is None else embedding_size
												_a	:    Any         =		hidden_size
												_a	:    Any         =		num_hidden_layers
												_a	:    List[Any]         =		num_attention_heads
												_a	:    str         =		hidden_act
												_a	:    Any         =		intermediate_size
												_a	:    Dict         =		hidden_dropout_prob
												_a	:    Optional[Any]         =		attention_probs_dropout_prob
												_a	:    str         =		max_position_embeddings
												_a	:    Dict         =		type_vocab_size
												_a	:    List[Any]         =		initializer_range
												_a	:    Dict         =		layer_norm_eps
												_a	:    Dict         =		rotary_value
												_a	:    Dict         =		use_cache



class 				UpperCAmelCase_		(			__lowercase   ):
							"""simple docstring"""







							@property
							def 				__lowercase				(			self       )    ->    Mapping[str, Mapping[int, str]]:
												if self.task == "multiple-choice":
																	_a	:    Dict         =		{0: '''batch''', 1: '''choice''', 2: '''sequence'''}
												else:
																	_a	:    List[Any]         =		{0: '''batch''', 1: '''sequence'''}
												_a	:    Union[str, Any]         =		{0: '''batch''', 1: '''sequence'''}
												return OrderedDict(
												    [
												        ('''input_ids''', dynamic_axis),
												        ('''attention_mask''', dynamic_axis),
												        ('''token_type_ids''', dynamic_axis),
												    ]       )



