'''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). '''      ,	__SCREAMING_SNAKE_CASE      ,	)
class  __magic_name__  (		__SCREAMING_SNAKE_CASE    ):





   """simple docstring"""




   lowerCAmelCase      :    int  			=   RobertaConfig
   lowerCAmelCase      :    Any  			=   '''roberta'''



   def __init__(	self					:  Dict		,  _lowercase					:  Optional[int]							):



    """simple docstring"""



    super().__init__(__snake_case							)

    _UpperCamelCase:     Optional[int]							= RobertaEmbeddings(__snake_case							)
    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. '''      ,	__SCREAMING_SNAKE_CASE      ,	)
class  __magic_name__  (		__SCREAMING_SNAKE_CASE    ):





   """simple docstring"""




   lowerCAmelCase      :    Dict  			=   RobertaConfig
   lowerCAmelCase      :    Tuple  			=   '''roberta'''



   def __init__(	self					:  List[str]		,  _lowercase					:  Optional[Any]							):



    """simple docstring"""



    super().__init__(__snake_case							)
    _UpperCamelCase:     Any							= config.num_labels
    _UpperCamelCase:     Union[str, Any]							= config.num_hidden_layers

    _UpperCamelCase:     Optional[int]							= DeeRobertaModel(__snake_case							)
    _UpperCamelCase:     Dict							= nn.Dropout(config.hidden_dropout_prob							)
    _UpperCamelCase:     Optional[Any]							= nn.Linear(config.hidden_size		,  self.config.num_labels							)



   @add_start_docstrings_to_model_forward(__snake_case							)
   def   lowerCAmelCase						(	self					:  int		,  _lowercase					:  Optional[int]=None		,  _lowercase					:  int=None		,  _lowercase					:  Optional[Any]=None		,  _lowercase					:  Union[str, Any]=None		,  _lowercase					:  Any=None		,  _lowercase					:  int=None		,  _lowercase					:  Union[str, Any]=None		,  _lowercase					:  Dict=-1		,  _lowercase					:  Optional[Any]=False		,  ):



    """simple docstring"""



    _UpperCamelCase:     List[str]							= self.num_layers
    try:
     _UpperCamelCase:     Tuple							= self.roberta(
         __snake_case		,  attention_mask=__snake_case		,  token_type_ids=__snake_case		,  position_ids=__snake_case		,  head_mask=__snake_case		,  inputs_embeds=__snake_case		,  )

     _UpperCamelCase:     int							= outputs[1]

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

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

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

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


UpperCAmelCase_			    =							{'''configuration_xlnet''': ['''XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''XLNetConfig''']}

try:
  if not is_sentencepiece_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase_			    =							['''XLNetTokenizer''']

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase_			    =							['''XLNetTokenizerFast''']

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase_			    =							[
      '''XLNET_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''XLNetForMultipleChoice''',
      '''XLNetForQuestionAnswering''',
      '''XLNetForQuestionAnsweringSimple''',
      '''XLNetForSequenceClassification''',
      '''XLNetForTokenClassification''',
      '''XLNetLMHeadModel''',
      '''XLNetModel''',
      '''XLNetPreTrainedModel''',
      '''load_tf_weights_in_xlnet''',
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase_			    =							[
      '''TF_XLNET_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''TFXLNetForMultipleChoice''',
      '''TFXLNetForQuestionAnsweringSimple''',
      '''TFXLNetForSequenceClassification''',
      '''TFXLNetForTokenClassification''',
      '''TFXLNetLMHeadModel''',
      '''TFXLNetMainLayer''',
      '''TFXLNetModel''',
      '''TFXLNetPreTrainedModel''',
  ]


if TYPE_CHECKING:
  from .configuration_xlnet import XLNET_PRETRAINED_CONFIG_ARCHIVE_MAP, XLNetConfig

  try:
    if not is_sentencepiece_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_xlnet import XLNetTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_xlnet_fast import XLNetTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_xlnet import (
        XLNET_PRETRAINED_MODEL_ARCHIVE_LIST,
        XLNetForMultipleChoice,
        XLNetForQuestionAnswering,
        XLNetForQuestionAnsweringSimple,
        XLNetForSequenceClassification,
        XLNetForTokenClassification,
        XLNetLMHeadModel,
        XLNetModel,
        XLNetPreTrainedModel,
        load_tf_weights_in_xlnet,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_xlnet import (
        TF_XLNET_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFXLNetForMultipleChoice,
        TFXLNetForQuestionAnsweringSimple,
        TFXLNetForSequenceClassification,
        TFXLNetForTokenClassification,
        TFXLNetLMHeadModel,
        TFXLNetMainLayer,
        TFXLNetModel,
        TFXLNetPreTrainedModel,
    )

else:
  import sys

  UpperCAmelCase_			    =							_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)