


'''simple docstring'''







from typing import Dict, List, Optional, Tuple, Union

import numpy as np

from ...image_processing_utils import BaseImageProcessor, BatchFeature, get_size_dict
from ...image_transforms import (
    center_crop,
    get_resize_output_image_size,
    normalize,
    rescale,
    resize,
    to_channel_dimension_format,
)
from ...image_utils import (
    IMAGENET_STANDARD_MEAN,
    IMAGENET_STANDARD_STD,
    ChannelDimension,
    ImageInput,
    PILImageResampling,
    make_list_of_images,
    to_numpy_array,
    valid_images,
)
from ...utils import TensorType, is_torch_available, is_torch_tensor, logging


if is_torch_available():
		import torch


_SCREAMING_SNAKE_CASE							: Union[str, Any]										=			logging.get_logger(__name__)






class 		_snake_case     (					__lowerCAmelCase		):






						'''simple docstring'''







						__snake_case				=   ["""pixel_values"""]







						def __init__(	self:						Any   ,				__UpperCamelCase:						Union[str, Any] = True   ,				__UpperCamelCase:						int = None   ,				__UpperCamelCase:						Tuple = PILImageResampling.BILINEAR   ,				__UpperCamelCase:						List[Any] = True   ,				__UpperCamelCase:						Tuple = None   ,				__UpperCamelCase:						Optional[Any] = True   ,				__UpperCamelCase:						List[Any] = 1 / 255   ,				__UpperCamelCase:						Union[str, Any] = True   ,				__UpperCamelCase:						Tuple = None   ,				__UpperCamelCase:						Any = None   ,				**__UpperCamelCase:						Any   ,				)				->						None:
										super().__init__(**_SCREAMING_SNAKE_CASE  )
										__magic_name__					:       Union[str, Any]    							=							size if size is not None else {"shortest_edge": 256}
										__magic_name__					:       Dict    							=							get_size_dict(_SCREAMING_SNAKE_CASE   ,				default_to_square=_SCREAMING_SNAKE_CASE  )
										__magic_name__					:       Optional[int]    							=							crop_size if crop_size is not None else {"height": 224, "width": 224}
										__magic_name__					:       Optional[Any]    							=							get_size_dict(_SCREAMING_SNAKE_CASE   ,				param_name="crop_size"  )
										__magic_name__					:       str    							=							do_resize
										__magic_name__					:       Any    							=							size
										__magic_name__					:       int    							=							resample
										__magic_name__					:       int    							=							do_center_crop
										__magic_name__					:       int    							=							crop_size
										__magic_name__					:       Union[str, Any]    							=							do_rescale
										__magic_name__					:       int    							=							rescale_factor
										__magic_name__					:       Optional[int]    							=							do_normalize
										__magic_name__					:       List[Any]    							=							image_mean if image_mean is not None else IMAGENET_STANDARD_MEAN
										__magic_name__					:       Tuple    							=							image_std if image_std is not None else IMAGENET_STANDARD_STD







						def lowerCAmelCase__			(	self:						Any   ,				__UpperCamelCase:						List[Any]   ,				__UpperCamelCase:						List[str]   ,				__UpperCamelCase:						int = PILImageResampling.BICUBIC   ,				__UpperCamelCase:						Optional[int] = None   ,				**__UpperCamelCase:						List[Any]   ,				)				->						np.ndarray:
										__magic_name__					:       Optional[Any]    							=							get_size_dict(_SCREAMING_SNAKE_CASE   ,				default_to_square=_SCREAMING_SNAKE_CASE  )
										if "shortest_edge" not in size:
														raise ValueError(f"""The `size` parameter must contain the key `shortest_edge`. Got {size.keys()}"""  )
										__magic_name__					:       List[Any]    							=							get_resize_output_image_size(_SCREAMING_SNAKE_CASE   ,				size=size["shortest_edge"]   ,				default_to_square=_SCREAMING_SNAKE_CASE  )
										return resize(_SCREAMING_SNAKE_CASE   ,				size=_SCREAMING_SNAKE_CASE   ,				resample=_SCREAMING_SNAKE_CASE   ,				data_format=_SCREAMING_SNAKE_CASE   ,				**_SCREAMING_SNAKE_CASE  )







						def lowerCAmelCase__			(	self:						Union[str, Any]   ,				__UpperCamelCase:						str   ,				__UpperCamelCase:						int   ,				__UpperCamelCase:						List[Any] = None   ,				**__UpperCamelCase:						str   ,				)				->						np.ndarray:
										__magic_name__					:       int    							=							get_size_dict(_SCREAMING_SNAKE_CASE  )
										if "height" not in size or "width" not in size:
														raise ValueError(f"""The `size` parameter must contain the keys `height` and `width`. Got {size.keys()}"""  )
										return center_crop(_SCREAMING_SNAKE_CASE   ,				size=(size["height"], size["width"])   ,				data_format=_SCREAMING_SNAKE_CASE   ,				**_SCREAMING_SNAKE_CASE  )







						def lowerCAmelCase__			(	self:						str   ,				__UpperCamelCase:						int   ,				__UpperCamelCase:						Optional[int]   ,				__UpperCamelCase:						Optional[int] = None   ,				**__UpperCamelCase:						List[Any]  )				->						np.ndarray:
										return rescale(_SCREAMING_SNAKE_CASE   ,				scale=_SCREAMING_SNAKE_CASE   ,				data_format=_SCREAMING_SNAKE_CASE   ,				**_SCREAMING_SNAKE_CASE  )







						def lowerCAmelCase__			(	self:						int   ,				__UpperCamelCase:						Union[str, Any]   ,				__UpperCamelCase:						Any   ,				__UpperCamelCase:						List[Any]   ,				__UpperCamelCase:						List[str] = None   ,				**__UpperCamelCase:						Optional[Any]   ,				)				->						np.ndarray:
										return normalize(_SCREAMING_SNAKE_CASE   ,				mean=_SCREAMING_SNAKE_CASE   ,				std=_SCREAMING_SNAKE_CASE   ,				data_format=_SCREAMING_SNAKE_CASE   ,				**_SCREAMING_SNAKE_CASE  )







						def lowerCAmelCase__			(	self:						Any   ,				__UpperCamelCase:						Any   ,				__UpperCamelCase:						Union[str, Any] = None   ,				__UpperCamelCase:						str = None   ,				__UpperCamelCase:						Tuple = None   ,				__UpperCamelCase:						Dict = None   ,				__UpperCamelCase:						int = None   ,				__UpperCamelCase:						Optional[int] = None   ,				__UpperCamelCase:						List[Any] = None   ,				__UpperCamelCase:						Optional[int] = None   ,				__UpperCamelCase:						Any = None   ,				__UpperCamelCase:						Tuple = None   ,				__UpperCamelCase:						Union[str, Any] = None   ,				__UpperCamelCase:						List[Any] = ChannelDimension.FIRST   ,				**__UpperCamelCase:						Optional[int]   ,				)				->						Optional[Any]:
										__magic_name__					:       Optional[int]    							=							do_resize if do_resize is not None else self.do_resize
										__magic_name__					:       List[str]    							=							size if size is not None else self.size
										__magic_name__					:       List[Any]    							=							get_size_dict(_SCREAMING_SNAKE_CASE   ,				default_to_square=_SCREAMING_SNAKE_CASE  )
										__magic_name__					:       int    							=							resample if resample is not None else self.resample
										__magic_name__					:       Union[str, Any]    							=							do_center_crop if do_center_crop is not None else self.do_center_crop
										__magic_name__					:       Tuple    							=							crop_size if crop_size is not None else self.crop_size
										__magic_name__					:       Any    							=							get_size_dict(_SCREAMING_SNAKE_CASE   ,				param_name="crop_size"  )
										__magic_name__					:       str    							=							do_rescale if do_rescale is not None else self.do_rescale
										__magic_name__					:       Any    							=							rescale_factor if rescale_factor is not None else self.rescale_factor
										__magic_name__					:       Any    							=							do_normalize if do_normalize is not None else self.do_normalize
										__magic_name__					:       int    							=							image_mean if image_mean is not None else self.image_mean
										__magic_name__					:       str    							=							image_std if image_std is not None else self.image_std

										__magic_name__					:       Optional[int]    							=							make_list_of_images(_SCREAMING_SNAKE_CASE  )

										if not valid_images(_SCREAMING_SNAKE_CASE  ):
														raise ValueError(
														    "Invalid image type. Must be of type PIL.Image.Image, numpy.ndarray, "
														    "torch.Tensor, tf.Tensor or jax.ndarray."  )

										if do_resize and size is None:
														raise ValueError("Size must be specified if do_resize is True."  )

										if do_center_crop and crop_size is None:
														raise ValueError("Crop size must be specified if do_center_crop is True."  )

										if do_rescale and rescale_factor is None:
														raise ValueError("Rescale factor must be specified if do_rescale is True."  )

										if do_normalize and (image_mean is None or image_std is None):
														raise ValueError("Image mean and std must be specified if do_normalize is True."  )

										# All transformations expect numpy arrays.
										__magic_name__					:       int    							=							[to_numpy_array(_SCREAMING_SNAKE_CASE  ) for image in images]

										if do_resize:
														__magic_name__					:       Union[str, Any]    							=							[self.resize(image=_SCREAMING_SNAKE_CASE   ,				size=_SCREAMING_SNAKE_CASE   ,				resample=_SCREAMING_SNAKE_CASE  ) for image in images]

										if do_center_crop:
														__magic_name__					:       List[Any]    							=							[self.center_crop(image=_SCREAMING_SNAKE_CASE   ,				size=_SCREAMING_SNAKE_CASE  ) for image in images]

										if do_rescale:
														__magic_name__					:       Optional[int]    							=							[self.rescale(image=_SCREAMING_SNAKE_CASE   ,				scale=_SCREAMING_SNAKE_CASE  ) for image in images]

										if do_normalize:
														__magic_name__					:       str    							=							[self.normalize(image=_SCREAMING_SNAKE_CASE   ,				mean=_SCREAMING_SNAKE_CASE   ,				std=_SCREAMING_SNAKE_CASE  ) for image in images]

										__magic_name__					:       Tuple    							=							[to_channel_dimension_format(_SCREAMING_SNAKE_CASE   ,				_SCREAMING_SNAKE_CASE  ) for image in images]

										__magic_name__					:       str    							=							{"pixel_values": images}
										return BatchFeature(data=_SCREAMING_SNAKE_CASE   ,				tensor_type=_SCREAMING_SNAKE_CASE  )







						def lowerCAmelCase__			(	self:						Optional[Any]   ,				__UpperCamelCase:						List[str]   ,				__UpperCamelCase:						Union[str, Any] = None  )				->						Optional[Any]:
										__magic_name__					:       str    							=							outputs.logits

										# Resize logits and compute semantic segmentation maps
										if target_sizes is not None:
														if len(_SCREAMING_SNAKE_CASE  ) != len(_SCREAMING_SNAKE_CASE  ):
																		raise ValueError(
																		    "Make sure that you pass in as many target sizes as the batch dimension of the logits"  )

														if is_torch_tensor(_SCREAMING_SNAKE_CASE  ):
																		__magic_name__					:       Any    							=							target_sizes.numpy()

														__magic_name__					:       Union[str, Any]    							=							[]

														for idx in range(len(_SCREAMING_SNAKE_CASE  )  ):
																		__magic_name__					:       Optional[Any]    							=							torch.nn.functional.interpolate(
																		    logits[idx].unsqueeze(dim=0  )   ,				size=target_sizes[idx]   ,				mode="bilinear"   ,				align_corners=_SCREAMING_SNAKE_CASE  )
																		__magic_name__					:       List[Any]    							=							resized_logits[0].argmax(dim=0  )
																		semantic_segmentation.append(_SCREAMING_SNAKE_CASE  )
										else:
														__magic_name__					:       int    							=							logits.argmax(dim=1  )
														__magic_name__					:       Optional[int]    							=							[semantic_segmentation[i] for i in range(semantic_segmentation.shape[0]  )]

										return semantic_segmentation





'''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). """   ,   __lowerCAmelCase   ,   )
class   _a							(		__lowerCAmelCase				):
     SCREAMING_SNAKE_CASE_				:       str       							=		RobertaConfig
     SCREAMING_SNAKE_CASE_				:       Dict       							=		"""roberta"""






     def __init__(					self   ,_SCREAMING_SNAKE_CASE       )				->     Optional[Any]:
          super().__init__(_SCREAMING_SNAKE_CASE       )

          _snake_case				= 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. """   ,   __lowerCAmelCase   ,   )
class   _a							(		__lowerCAmelCase				):
     SCREAMING_SNAKE_CASE_				:       List[Any]       							=		RobertaConfig
     SCREAMING_SNAKE_CASE_				:       List[str]       							=		"""roberta"""






     def __init__(					self   ,_SCREAMING_SNAKE_CASE       )				->     str:
          super().__init__(_SCREAMING_SNAKE_CASE       )
          _snake_case				= config.num_labels
          _snake_case				= config.num_hidden_layers

          _snake_case				= DeeRobertaModel(_SCREAMING_SNAKE_CASE       )
          _snake_case				= nn.Dropout(config.hidden_dropout_prob       )
          _snake_case				= nn.Linear(config.hidden_size   ,self.config.num_labels       )






     @add_start_docstrings_to_model_forward(_SCREAMING_SNAKE_CASE       )
     def 							_lowercase      (					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   ,)				->     List[str]:
          _snake_case				= self.num_layers
          try:
               _snake_case				= 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   ,)

               _snake_case				= outputs[1]

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

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

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

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

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


