


'''simple docstring'''




from __future__ import absolute_import, division, print_function, unicode_literals

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

from transformers import RobertaConfig
from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.roberta.modeling_roberta import (
    ROBERTA_INPUTS_DOCSTRING,
    ROBERTA_START_DOCSTRING,
    RobertaEmbeddings,
)

from .modeling_highway_bert import BertPreTrainedModel, DeeBertModel, HighwayException, entropy



@add_start_docstrings(
    """The RoBERTa Model transformer with early exiting (DeeRoBERTa). """      ,				__lowercase      ,				)
class  lowerCAmelCase				(				__lowercase							):
	_lowerCamelCase   :					str			      = RobertaConfig
	_lowerCamelCase   :					Dict			      = 'roberta'



	def __init__(					self ,					snake_case__    ):
		super().__init__(__A    )

		lowerCAmelCase					:    List[Any]      							= RobertaEmbeddings(__A    )
		self.init_weights()



@add_start_docstrings(
    """RoBERTa Model (with early exiting - DeeRoBERTa) with a classifier on top,\n    also takes care of multi-layer training. """      ,				__lowercase      ,				)
class  lowerCAmelCase				(				__lowercase							):
	_lowerCamelCase   :					int			      = RobertaConfig
	_lowerCamelCase   :					Tuple			      = 'roberta'



	def __init__(					self ,					snake_case__    ):
		super().__init__(__A    )
		lowerCAmelCase					:    Union[str, Any]      							= config.num_labels
		lowerCAmelCase					:    Any      							= config.num_hidden_layers

		lowerCAmelCase					:    List[Any]      							= DeeRobertaModel(__A    )
		lowerCAmelCase					:    Union[str, Any]      							= nn.Dropout(config.hidden_dropout_prob    )
		lowerCAmelCase					:    List[Any]      							= nn.Linear(config.hidden_size ,					self.config.num_labels    )



	@add_start_docstrings_to_model_forward(__A    )
	def 				lowercase     (					self ,					snake_case__=None ,					snake_case__=None ,					snake_case__=None ,					snake_case__=None ,					snake_case__=None ,					snake_case__=None ,					snake_case__=None ,					snake_case__=-1 ,					snake_case__=False ,					):
		lowerCAmelCase					:    Optional[Any]      							= self.num_layers
		try:
			lowerCAmelCase					:    Any      							= self.roberta(
			    __A ,					attention_mask=__A ,					token_type_ids=__A ,					position_ids=__A ,					head_mask=__A ,					inputs_embeds=__A ,					)

			lowerCAmelCase					:    Dict      							= outputs[1]

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

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

			# work with highway exits
			lowerCAmelCase					:    str      							= []
			for highway_exit in outputs[-1]:
				lowerCAmelCase					:    int      							= 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
					lowerCAmelCase					:    Any      							= MSELoss()
					lowerCAmelCase					:    List[Any]      							= loss_fct(highway_logits.view(-1    ) ,					labels.view(-1    )    )
				else:
					lowerCAmelCase					:    List[str]      							= CrossEntropyLoss()
					lowerCAmelCase					:    Optional[int]      							= loss_fct(highway_logits.view(-1 ,					self.num_labels    ) ,					labels.view(-1    )    )
				highway_losses.append(__A    )

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

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









'''simple docstring'''




from collections import OrderedDict
from typing import Mapping

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


_lowerCAmelCase  :		Tuple		     =	logging.get_logger(__name__)

_lowerCAmelCase  :		List[str]		     =	{
    'facebook/xmod-base': 'https://huggingface.co/facebook/xmod-base/resolve/main/config.json',
    'facebook/xmod-large-prenorm': 'https://huggingface.co/facebook/xmod-large-prenorm/resolve/main/config.json',
    'facebook/xmod-base-13-125k': 'https://huggingface.co/facebook/xmod-base-13-125k/resolve/main/config.json',
    'facebook/xmod-base-30-125k': 'https://huggingface.co/facebook/xmod-base-30-125k/resolve/main/config.json',
    'facebook/xmod-base-30-195k': 'https://huggingface.co/facebook/xmod-base-30-195k/resolve/main/config.json',
    'facebook/xmod-base-60-125k': 'https://huggingface.co/facebook/xmod-base-60-125k/resolve/main/config.json',
    'facebook/xmod-base-60-265k': 'https://huggingface.co/facebook/xmod-base-60-265k/resolve/main/config.json',
    'facebook/xmod-base-75-125k': 'https://huggingface.co/facebook/xmod-base-75-125k/resolve/main/config.json',
    'facebook/xmod-base-75-269k': 'https://huggingface.co/facebook/xmod-base-75-269k/resolve/main/config.json',
}



class  lowerCAmelCase				(				a							):
	_lowerCamelCase   :					int			      = """xmod"""



	def __init__(					self ,					snake_case__=3_0522 ,					snake_case__=768 ,					snake_case__=12 ,					snake_case__=12 ,					snake_case__=3072 ,					snake_case__="gelu" ,					snake_case__=0.1 ,					snake_case__=0.1 ,					snake_case__=512 ,					snake_case__=2 ,					snake_case__=0.0_2 ,					snake_case__=1e-1_2 ,					snake_case__=1 ,					snake_case__=0 ,					snake_case__=2 ,					snake_case__="absolute" ,					snake_case__=True ,					snake_case__=None ,					snake_case__=False ,					snake_case__=2 ,					snake_case__=False ,					snake_case__=True ,					snake_case__=True ,					snake_case__=("en_XX",) ,					snake_case__=None ,					**snake_case__ ,					):
		super().__init__(pad_token_id=snake_case__ ,					bos_token_id=snake_case__ ,					eos_token_id=snake_case__ ,					**snake_case__    )

		lowerCAmelCase					:    Dict      							= vocab_size
		lowerCAmelCase					:    Optional[Any]      							= hidden_size
		lowerCAmelCase					:    int      							= num_hidden_layers
		lowerCAmelCase					:    List[Any]      							= num_attention_heads
		lowerCAmelCase					:    List[Any]      							= hidden_act
		lowerCAmelCase					:    Optional[int]      							= intermediate_size
		lowerCAmelCase					:    Optional[int]      							= hidden_dropout_prob
		lowerCAmelCase					:    Optional[Any]      							= attention_probs_dropout_prob
		lowerCAmelCase					:    str      							= max_position_embeddings
		lowerCAmelCase					:    int      							= type_vocab_size
		lowerCAmelCase					:    List[Any]      							= initializer_range
		lowerCAmelCase					:    Any      							= layer_norm_eps
		lowerCAmelCase					:    Dict      							= position_embedding_type
		lowerCAmelCase					:    Optional[Any]      							= use_cache
		lowerCAmelCase					:    Union[str, Any]      							= classifier_dropout
		lowerCAmelCase					:    int      							= pre_norm
		lowerCAmelCase					:    Optional[Any]      							= adapter_reduction_factor
		lowerCAmelCase					:    Any      							= adapter_layer_norm
		lowerCAmelCase					:    Dict      							= adapter_reuse_layer_norm
		lowerCAmelCase					:    Any      							= ln_before_adapter
		lowerCAmelCase					:    Optional[Any]      							= list(snake_case__    )
		lowerCAmelCase					:    List[Any]      							= default_language



class  lowerCAmelCase				(				a							):



	@property
	def 				lowercase     (					self    ):
		if self.task == "multiple-choice":
			lowerCAmelCase					:    List[Any]      							= {0: 'batch', 1: 'choice', 2: 'sequence'}
		else:
			lowerCAmelCase					:    Optional[int]      							= {0: 'batch', 1: 'sequence'}
		return OrderedDict(
		    [
		        ('input_ids', dynamic_axis),
		        ('attention_mask', dynamic_axis),
		    ]    )
