




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). '''							,		_a							,		)
class     __a(	_a  ):






    """simple docstring"""





    lowerCAmelCase         = RobertaConfig
    lowerCAmelCase         = '''roberta'''





    def __init__( self  ,_SCREAMING_SNAKE_CASE    )					->						Dict:
         super().__init__(_SCREAMING_SNAKE_CASE    )

         UpperCAmelCase_     :				Tuple  =     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. '''							,		_a							,		)
class     __a(	_a  ):






    """simple docstring"""





    lowerCAmelCase         = RobertaConfig
    lowerCAmelCase         = '''roberta'''





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

         UpperCAmelCase_     :				List[Any]  =     DeeRobertaModel(_SCREAMING_SNAKE_CASE    )
         UpperCAmelCase_     :				int  =     nn.Dropout(config.hidden_dropout_prob    )
         UpperCAmelCase_     :				Tuple  =     nn.Linear(config.hidden_size  ,self.config.num_labels    )





    @add_start_docstrings_to_model_forward(_SCREAMING_SNAKE_CASE    )
    def 				a__ ( 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  ,)					->						int:
         UpperCAmelCase_     :				Union[str, Any]  =     self.num_layers
         try:
              UpperCAmelCase_     :				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_     :				Optional[Any]  =     outputs[1]

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

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

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

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





import warnings

from ...utils import logging
from .image_processing_clip import CLIPImageProcessor


__a							     =					logging.get_logger(__name__)

class     __a(	_a  ):






    """simple docstring"""





    def __init__( self  ,*_SCREAMING_SNAKE_CASE  ,**_SCREAMING_SNAKE_CASE    )					->						None:
         warnings.warn(
             '''The class CLIPFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please'''
             ''' use CLIPImageProcessor instead.'''  ,_SCREAMING_SNAKE_CASE  ,)
         super().__init__(*_SCREAMING_SNAKE_CASE  ,**_SCREAMING_SNAKE_CASE    )