




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). '''	,      UpperCAmelCase__	,      )
class        UpperCamelCase_					(				UpperCAmelCase__  ):

    '''simple docstring'''






    UpperCAmelCase__													=   RobertaConfig
    UpperCAmelCase__													=   '''roberta'''


    def __init__(    self   :			Any      ,				UpperCAmelCase__   :			Any)		->str:





           '''simple docstring'''

           super().__init__(UpperCAmelCase__)

           A__   				=					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. '''	,      UpperCAmelCase__	,      )
class        UpperCamelCase_					(				UpperCAmelCase__  ):

    '''simple docstring'''






    UpperCAmelCase__													=   RobertaConfig
    UpperCAmelCase__													=   '''roberta'''


    def __init__(    self   :			Tuple      ,				UpperCAmelCase__   :			Optional[int])		->Optional[Any]:





           '''simple docstring'''

           super().__init__(UpperCAmelCase__)
           A__   				=					config.num_labels
           A__   				=					config.num_hidden_layers

           A__   				=					DeeRobertaModel(UpperCAmelCase__)
           A__   				=					nn.Dropout(config.hidden_dropout_prob)
           A__   				=					nn.Linear(config.hidden_size      ,				self.config.num_labels)





    @add_start_docstrings_to_model_forward(UpperCAmelCase__)
    def  SCREAMING_SNAKE_CASE							(    self   :			Any      ,				UpperCAmelCase__   :			int=None      ,				UpperCAmelCase__   :			Dict=None      ,				UpperCAmelCase__   :			List[str]=None      ,				UpperCAmelCase__   :			Tuple=None      ,				UpperCAmelCase__   :			List[str]=None      ,				UpperCAmelCase__   :			Optional[int]=None      ,				UpperCAmelCase__   :			Any=None      ,				UpperCAmelCase__   :			str=-1      ,				UpperCAmelCase__   :			Optional[int]=False      ,				)		->Dict:





           '''simple docstring'''

           A__   				=					self.num_layers
           try:
                  A__   				=					self.roberta(
                      UpperCAmelCase__      ,				attention_mask=UpperCAmelCase__      ,				token_type_ids=UpperCAmelCase__      ,				position_ids=UpperCAmelCase__      ,				head_mask=UpperCAmelCase__      ,				inputs_embeds=UpperCAmelCase__      ,				)

                  A__   				=					outputs[1]

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

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

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

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

import numpy as np
import PIL.Image

from ...image_processing_utils import BaseImageProcessor, BatchFeature
from ...image_transforms import rescale, resize, to_channel_dimension_format
from ...image_utils import (
    ChannelDimension,
    PILImageResampling,
    get_image_size,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, logging


_lowerCamelCase			:	str							 =     logging.get_logger(__name__)


class        UpperCamelCase_					(				UpperCAmelCase__  ):

    '''simple docstring'''






    UpperCAmelCase__													=   ['''pixel_values''']


    def __init__(    self   :			List[Any]      ,				UpperCAmelCase__   :			bool = True      ,				UpperCAmelCase__   :			int = 32      ,				UpperCAmelCase__   :			Dict=PILImageResampling.BILINEAR      ,				UpperCAmelCase__   :			bool = True      ,				**UpperCAmelCase__   :			Optional[int]      ,				)		->None:





           '''simple docstring'''

           A__   				=					do_resize
           A__   				=					do_rescale
           A__   				=					size_divisor
           A__   				=					resample
           super().__init__(**UpperCAmelCase__)


    def  SCREAMING_SNAKE_CASE							(    self   :			int      ,				UpperCAmelCase__   :			np.ndarray      ,				UpperCAmelCase__   :			int      ,				UpperCAmelCase__   :			List[str]      ,				UpperCAmelCase__   :			Optional[ChannelDimension] = None      ,				**UpperCAmelCase__   :			Dict)		->np.ndarray:





           '''simple docstring'''

           A__				,							A__   				=					get_image_size(UpperCAmelCase__)
           # Rounds the height and width down to the closest multiple of size_divisor
           A__   				=					height // size_divisor * size_divisor
           A__   				=					width // size_divisor * size_divisor
           A__   				=					resize(UpperCAmelCase__      ,				(new_h, new_w)      ,				resample=UpperCAmelCase__      ,				data_format=UpperCAmelCase__      ,				**UpperCAmelCase__)
           return image


    def  SCREAMING_SNAKE_CASE							(    self   :			List[Any]      ,				UpperCAmelCase__   :			np.ndarray      ,				UpperCAmelCase__   :			float      ,				UpperCAmelCase__   :			Optional[ChannelDimension] = None      ,				**UpperCAmelCase__   :			List[str])		->np.ndarray:





           '''simple docstring'''

           return rescale(image=UpperCAmelCase__      ,				scale=UpperCAmelCase__      ,				data_format=UpperCAmelCase__      ,				**UpperCAmelCase__)





    def  SCREAMING_SNAKE_CASE							(    self   :			Optional[int]      ,				UpperCAmelCase__   :			Union["PIL.Image.Image", TensorType, List["PIL.Image.Image"], List[TensorType]]      ,				UpperCAmelCase__   :			Optional[bool] = None      ,				UpperCAmelCase__   :			Optional[int] = None      ,				UpperCAmelCase__   :			Dict=None      ,				UpperCAmelCase__   :			Optional[bool] = None      ,				UpperCAmelCase__   :			Optional[Union[TensorType, str]] = None      ,				UpperCAmelCase__   :			ChannelDimension = ChannelDimension.FIRST      ,				**UpperCAmelCase__   :			List[str]      ,				)		->BatchFeature:





           '''simple docstring'''

           A__   				=					do_resize if do_resize is not None else self.do_resize
           A__   				=					do_rescale if do_rescale is not None else self.do_rescale
           A__   				=					size_divisor if size_divisor is not None else self.size_divisor
           A__   				=					resample if resample is not None else self.resample

           if do_resize and size_divisor is None:
                  raise ValueError('''size_divisor is required for resizing''')

           A__   				=					make_list_of_images(UpperCAmelCase__)

           if not valid_images(UpperCAmelCase__):
                  raise ValueError('''Invalid image(s)''')

           # All transformations expect numpy arrays.
           A__   				=					[to_numpy_array(UpperCAmelCase__) for img in images]

           if do_resize:
                  A__   				=					[self.resize(UpperCAmelCase__      ,				size_divisor=UpperCAmelCase__      ,				resample=UpperCAmelCase__) for image in images]

           if do_rescale:
                  A__   				=					[self.rescale(UpperCAmelCase__      ,				scale=1 / 255) for image in images]

           A__   				=					[to_channel_dimension_format(UpperCAmelCase__      ,				UpperCAmelCase__) for image in images]

           A__   				=					{'''pixel_values''': images}
           return BatchFeature(data=UpperCAmelCase__      ,				tensor_type=UpperCAmelCase__)


