



"""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 					a	(  __lowercase       ):
  UpperCamelCase	:				str     						=					RobertaConfig
  UpperCamelCase	:				Dict     						=					'roberta'



  def __init__(						self :					List[Any]       ,		lowerCAmelCase :					Tuple		)     ->					Union[str, Any]:


     '''simple docstring'''




     super().__init__(__A		)

     SCREAMING_SNAKE_CASE_: Optional[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 					a	(  __lowercase       ):
  UpperCamelCase	:				int     						=					RobertaConfig
  UpperCamelCase	:				Tuple     						=					'roberta'



  def __init__(						self :					List[Any]       ,		lowerCAmelCase :					Dict		)     ->					Dict:


     '''simple docstring'''




     super().__init__(__A		)
     SCREAMING_SNAKE_CASE_: List[Any]     			=config.num_labels
     SCREAMING_SNAKE_CASE_: List[Any]     			=config.num_hidden_layers

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





  @add_start_docstrings_to_model_forward(__A		)
  def 	lowerCamelCase__       (						self :					List[str]       ,		lowerCAmelCase :					Dict=None       ,		lowerCAmelCase :					Dict=None       ,		lowerCAmelCase :					Tuple=None       ,		lowerCAmelCase :					Tuple=None       ,		lowerCAmelCase :					Optional[Any]=None       ,		lowerCAmelCase :					Tuple=None       ,		lowerCAmelCase :					Any=None       ,		lowerCAmelCase :					Tuple=-1       ,		lowerCAmelCase :					Union[str, Any]=False       ,		)     ->					str:


     '''simple docstring'''




     SCREAMING_SNAKE_CASE_: Union[str, Any]     			=self.num_layers
     try:
        SCREAMING_SNAKE_CASE_: int     			=self.roberta(
            __A       ,		attention_mask=__A       ,		token_type_ids=__A       ,		position_ids=__A       ,		head_mask=__A       ,		inputs_embeds=__A       ,		)

        SCREAMING_SNAKE_CASE_: List[Any]     			=outputs[1]

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

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

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

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








"""simple docstring"""
from typing import TYPE_CHECKING

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


_UpperCAmelCase     =			{
    """configuration_xlm""": ["""XLM_PRETRAINED_CONFIG_ARCHIVE_MAP""", """XLMConfig""", """XLMOnnxConfig"""],
    """tokenization_xlm""": ["""XLMTokenizer"""],
}

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _UpperCAmelCase     =			[
       """XLM_PRETRAINED_MODEL_ARCHIVE_LIST""",
       """XLMForMultipleChoice""",
       """XLMForQuestionAnswering""",
       """XLMForQuestionAnsweringSimple""",
       """XLMForSequenceClassification""",
       """XLMForTokenClassification""",
       """XLMModel""",
       """XLMPreTrainedModel""",
       """XLMWithLMHeadModel""",
   ]

try:
   if not is_tf_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   _UpperCAmelCase     =			[
       """TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST""",
       """TFXLMForMultipleChoice""",
       """TFXLMForQuestionAnsweringSimple""",
       """TFXLMForSequenceClassification""",
       """TFXLMForTokenClassification""",
       """TFXLMMainLayer""",
       """TFXLMModel""",
       """TFXLMPreTrainedModel""",
       """TFXLMWithLMHeadModel""",
   ]


if TYPE_CHECKING:
   from .configuration_xlm import XLM_PRETRAINED_CONFIG_ARCHIVE_MAP, XLMConfig, XLMOnnxConfig
   from .tokenization_xlm import XLMTokenizer

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_xlm import (
          XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
          XLMForMultipleChoice,
          XLMForQuestionAnswering,
          XLMForQuestionAnsweringSimple,
          XLMForSequenceClassification,
          XLMForTokenClassification,
          XLMModel,
          XLMPreTrainedModel,
          XLMWithLMHeadModel,
      )

   try:
      if not is_tf_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_tf_xlm import (
          TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFXLMForMultipleChoice,
          TFXLMForQuestionAnsweringSimple,
          TFXLMForSequenceClassification,
          TFXLMForTokenClassification,
          TFXLMMainLayer,
          TFXLMModel,
          TFXLMPreTrainedModel,
          TFXLMWithLMHeadModel,
      )

else:
   import sys

   _UpperCAmelCase     =			_LazyModule(__name__, globals()["""__file__"""], _import_structure, module_spec=__spec__)



