





'''simple docstring'''






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_tokenizers_available,
    is_torch_available,
)


__lowerCamelCase						: List[str]		=   {"""configuration_plbart""": ["""PLBART_PRETRAINED_CONFIG_ARCHIVE_MAP""", """PLBartConfig"""]}

try:
      if not is_sentencepiece_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __lowerCamelCase						: int		=   ["""PLBartTokenizer"""]

try:
      if not is_torch_available():
            raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
      pass
else:
      __lowerCamelCase						: Optional[int]		=   [
          """PLBART_PRETRAINED_MODEL_ARCHIVE_LIST""",
          """PLBartForCausalLM""",
          """PLBartForConditionalGeneration""",
          """PLBartForSequenceClassification""",
          """PLBartModel""",
          """PLBartPreTrainedModel""",
      ]


if TYPE_CHECKING:
      from .configuration_plbart import PLBART_PRETRAINED_CONFIG_ARCHIVE_MAP, PLBartConfig

      try:
            if not is_sentencepiece_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .tokenization_plbart import PLBartTokenizer

      try:
            if not is_torch_available():
                  raise OptionalDependencyNotAvailable()
      except OptionalDependencyNotAvailable:
            pass
      else:
            from .modeling_plbart import (
                PLBART_PRETRAINED_MODEL_ARCHIVE_LIST,
                PLBartForCausalLM,
                PLBartForConditionalGeneration,
                PLBartForSequenceClassification,
                PLBartModel,
                PLBartPreTrainedModel,
            )


else:
      import sys

      __lowerCamelCase						: Any		=   _LazyModule(__name__, globals()['__file__'], _import_structure)











'''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 		UpperCAmelCase   (				lowercase_):


  """simple docstring"""







  lowerCAmelCase_					=				RobertaConfig
  lowerCAmelCase_					=				"""roberta"""


  def __init__(				self    :					Dict							,	UpperCamelCase__    :					List[str]  )     ->						List[Any]:
    super().__init__(UpperCamelCase__  )

    _UpperCamelCase       			=RobertaEmbeddings(UpperCamelCase__  )
    self.init_weights()







@add_start_docstrings(
    """RoBERTa Model (with early exiting - DeeRoBERTa) with a classifier on top,
    also takes care of multi-layer training. """   ,     lowercase_   ,     )
class 		UpperCAmelCase   (				lowercase_):


  """simple docstring"""







  lowerCAmelCase_					=				RobertaConfig
  lowerCAmelCase_					=				"""roberta"""


  def __init__(				self    :					Tuple							,	UpperCamelCase__    :					Any  )     ->						Dict:
    super().__init__(UpperCamelCase__  )
    _UpperCamelCase       			=config.num_labels
    _UpperCamelCase       			=config.num_hidden_layers

    _UpperCamelCase       			=DeeRobertaModel(UpperCamelCase__  )
    _UpperCamelCase       			=nn.Dropout(config.hidden_dropout_prob  )
    _UpperCamelCase       			=nn.Linear(config.hidden_size							,	self.config.num_labels  )




  @add_start_docstrings_to_model_forward(UpperCamelCase__  )
  def   UpperCamelCase__      (				self    :					List[Any]							,	UpperCamelCase__    :					Optional[int]=None							,	UpperCamelCase__    :					Union[str, Any]=None							,	UpperCamelCase__    :					List[str]=None							,	UpperCamelCase__    :					Optional[Any]=None							,	UpperCamelCase__    :					int=None							,	UpperCamelCase__    :					Union[str, Any]=None							,	UpperCamelCase__    :					Tuple=None							,	UpperCamelCase__    :					Union[str, Any]=-1							,	UpperCamelCase__    :					Dict=False							,	)     ->						Union[str, Any]:
    _UpperCamelCase       			=self.num_layers
    try:
      _UpperCamelCase       			=self.roberta(
          UpperCamelCase__							,	attention_mask=UpperCamelCase__							,	token_type_ids=UpperCamelCase__							,	position_ids=UpperCamelCase__							,	head_mask=UpperCamelCase__							,	inputs_embeds=UpperCamelCase__							,	)

      _UpperCamelCase       			=outputs[1]

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

    if not self.training:
      _UpperCamelCase       			=entropy(UpperCamelCase__  )
      _UpperCamelCase       			=[]
      _UpperCamelCase       			=[]
    if labels is not None:
      if self.num_labels == 1:
        #  We are doing regression
        _UpperCamelCase       			=MSELoss()
        _UpperCamelCase       			=loss_fct(logits.view(-1  )							,	labels.view(-1  )  )
      else:
        _UpperCamelCase       			=CrossEntropyLoss()
        _UpperCamelCase       			=loss_fct(logits.view(-1							,	self.num_labels  )							,	labels.view(-1  )  )

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

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

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






