



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). ''' ,		_lowerCAmelCase ,		)
class 				_a				(   _lowerCAmelCase					):
			A      		=	RobertaConfig
			A      		=	'''roberta'''

			def __init__(self,   SCREAMING_SNAKE_CASE_					)       ->   List[str]:
						super().__init__(SCREAMING_SNAKE_CASE_					)

						UpperCAmelCase_:		str							      =     RobertaEmbeddings(SCREAMING_SNAKE_CASE_					)
						self.init_weights()







@add_start_docstrings(
    '''RoBERTa Model (with early exiting - DeeRoBERTa) with a classifier on top,
    also takes care of multi-layer training. ''' ,		_lowerCAmelCase ,		)
class 				_a				(   _lowerCAmelCase					):
			A      		=	RobertaConfig
			A      		=	'''roberta'''

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

						UpperCAmelCase_:		Tuple							      =     DeeRobertaModel(SCREAMING_SNAKE_CASE_					)
						UpperCAmelCase_:		Optional[int]							      =     nn.Dropout(config.hidden_dropout_prob					)
						UpperCAmelCase_:		Union[str, Any]							      =     nn.Linear(config.hidden_size,   self.config.num_labels					)

			@add_start_docstrings_to_model_forward(SCREAMING_SNAKE_CASE_					)
			def      __snake_case       (self,   SCREAMING_SNAKE_CASE_=None,   SCREAMING_SNAKE_CASE_=None,   SCREAMING_SNAKE_CASE_=None,   SCREAMING_SNAKE_CASE_=None,   SCREAMING_SNAKE_CASE_=None,   SCREAMING_SNAKE_CASE_=None,   SCREAMING_SNAKE_CASE_=None,   SCREAMING_SNAKE_CASE_=-1,   SCREAMING_SNAKE_CASE_=False,   )       ->   Dict:
						UpperCAmelCase_:		List[Any]							      =     self.num_layers
						try:
									UpperCAmelCase_:		Union[str, Any]							      =     self.roberta(
									    SCREAMING_SNAKE_CASE_,   attention_mask=SCREAMING_SNAKE_CASE_,   token_type_ids=SCREAMING_SNAKE_CASE_,   position_ids=SCREAMING_SNAKE_CASE_,   head_mask=SCREAMING_SNAKE_CASE_,   inputs_embeds=SCREAMING_SNAKE_CASE_,   )

									UpperCAmelCase_:		List[str]							      =     outputs[1]

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

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

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

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

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











from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


a      :  Optional[Any]			     =			{
    'configuration_mvp': ['MVP_PRETRAINED_CONFIG_ARCHIVE_MAP', 'MvpConfig', 'MvpOnnxConfig'],
    'tokenization_mvp': ['MvpTokenizer'],
}

try:
	if not is_tokenizers_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	a      :  str			     =			['MvpTokenizerFast']

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	a      :  Optional[Any]			     =			[
	    'MVP_PRETRAINED_MODEL_ARCHIVE_LIST',
	    'MvpForCausalLM',
	    'MvpForConditionalGeneration',
	    'MvpForQuestionAnswering',
	    'MvpForSequenceClassification',
	    'MvpModel',
	    'MvpPreTrainedModel',
	]

if TYPE_CHECKING:
	from .configuration_mvp import MVP_PRETRAINED_CONFIG_ARCHIVE_MAP, MvpConfig, MvpOnnxConfig
	from .tokenization_mvp import MvpTokenizer

	try:
		if not is_tokenizers_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .tokenization_mvp_fast import MvpTokenizerFast

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_mvp import (
		    MVP_PRETRAINED_MODEL_ARCHIVE_LIST,
		    MvpForCausalLM,
		    MvpForConditionalGeneration,
		    MvpForQuestionAnswering,
		    MvpForSequenceClassification,
		    MvpModel,
		    MvpPreTrainedModel,
		)

else:
	import sys

	a      :  Optional[Any]			     =			_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)






