'''simple docstring'''

import shutil
import tempfile
import unittest

import numpy as np

from transformers.testing_utils import (
    is_pt_tf_cross_test,
    require_tf,
    require_torch,
    require_torchvision,
    require_vision,
)
from transformers.utils import is_tf_available, is_torch_available, is_vision_available


if is_vision_available():
  from PIL import Image

  from transformers import AutoProcessor, SamImageProcessor, SamProcessor

if is_torch_available():
  import torch

if is_tf_available():
  import tensorflow as tf


@require_vision
@require_torchvision
class 					_lowercase						(   unittest.TestCase ):
      '''simple docstring'''


      def 		lowerCAmelCase__       (					self				)      ->      Union[str, Any]:



         '''simple docstring'''



         UpperCAmelCase__				:					Optional[Any] 					=	tempfile.mkdtemp()
         UpperCAmelCase__				:					Tuple 					=	SamImageProcessor()
         UpperCAmelCase__				:					Dict 					=	SamProcessor(lowerCamelCase_				)
         processor.save_pretrained(self.tmpdirname				)




      def 		lowerCAmelCase__       (					self       ,**lowerCamelCase_				)      ->      Optional[Any]:



         '''simple docstring'''



         return AutoProcessor.from_pretrained(self.tmpdirname       ,**lowerCamelCase_				).image_processor




      def 		lowerCAmelCase__       (					self				)      ->      List[Any]:



         '''simple docstring'''



         shutil.rmtree(self.tmpdirname				)




      def 		lowerCAmelCase__       (					self				)      ->      Tuple:



         '''simple docstring'''



         UpperCAmelCase__				:					int 					=	[np.random.randint(255       ,size=(3, 30, 400)       ,dtype=np.uinta				)]

         UpperCAmelCase__				:					Any 					=	[Image.fromarray(np.moveaxis(lowerCamelCase_       ,0       ,-1				)				) for x in image_inputs]

         return image_inputs




      def 		lowerCAmelCase__       (					self				)      ->      int:



         '''simple docstring'''



         UpperCAmelCase__				:					Optional[Any] 					=	SamProcessor(image_processor=self.get_image_processor()				)
         processor.save_pretrained(self.tmpdirname				)

         UpperCAmelCase__				:					Dict 					=	self.get_image_processor(do_normalize=lowerCamelCase_       ,padding_value=1.0				)

         UpperCAmelCase__				:					Dict 					=	SamProcessor.from_pretrained(self.tmpdirname       ,do_normalize=lowerCamelCase_       ,padding_value=1.0				)

         self.assertEqual(processor.image_processor.to_json_string()       ,image_processor_add_kwargs.to_json_string()				)
         self.assertIsInstance(processor.image_processor       ,lowerCamelCase_				)




      def 		lowerCAmelCase__       (					self				)      ->      Optional[int]:



         '''simple docstring'''



         UpperCAmelCase__				:					int 					=	self.get_image_processor()

         UpperCAmelCase__				:					Dict 					=	SamProcessor(image_processor=lowerCamelCase_				)

         UpperCAmelCase__				:					Tuple 					=	self.prepare_image_inputs()

         UpperCAmelCase__				:					str 					=	image_processor(lowerCamelCase_       ,return_tensors='''np'''				)
         UpperCAmelCase__				:					Tuple 					=	processor(images=lowerCamelCase_       ,return_tensors='''np'''				)

         input_feat_extract.pop('''original_sizes'''				)  # pop original_sizes as it is popped in the processor
         input_feat_extract.pop('''reshaped_input_sizes'''				)  # pop original_sizes as it is popped in the processor

         for key in input_feat_extract.keys():
            self.assertAlmostEqual(input_feat_extract[key].sum()       ,input_processor[key].sum()       ,delta=1e-2				)




      @require_torch
      def 		lowerCAmelCase__       (					self				)      ->      Union[str, Any]:



         '''simple docstring'''



         UpperCAmelCase__				:					Any 					=	self.get_image_processor()

         UpperCAmelCase__				:					int 					=	SamProcessor(image_processor=lowerCamelCase_				)
         UpperCAmelCase__				:					Tuple 					=	[torch.ones((1, 3, 5, 5)				)]

         UpperCAmelCase__				:					Any 					=	[[1764, 2646]]

         UpperCAmelCase__				:					Optional[int] 					=	[[683, 1024]]
         UpperCAmelCase__				:					Optional[int] 					=	processor.post_process_masks(lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_				)
         self.assertEqual(masks[0].shape       ,(1, 3, 1764, 2646)				)

         UpperCAmelCase__				:					Optional[int] 					=	processor.post_process_masks(
             lowerCamelCase_       ,torch.tensor(lowerCamelCase_				)       ,torch.tensor(lowerCamelCase_				)				)
         self.assertEqual(masks[0].shape       ,(1, 3, 1764, 2646)				)

         # should also work with np
         UpperCAmelCase__				:					Union[str, Any] 					=	[np.ones((1, 3, 5, 5)				)]
         UpperCAmelCase__				:					List[Any] 					=	processor.post_process_masks(lowerCamelCase_       ,np.array(lowerCamelCase_				)       ,np.array(lowerCamelCase_				)				)

         self.assertEqual(masks[0].shape       ,(1, 3, 1764, 2646)				)

         UpperCAmelCase__				:					Dict 					=	[[1, 0], [0, 1]]
         with self.assertRaises(lowerCamelCase_				):
            UpperCAmelCase__				:					Union[str, Any] 					=	processor.post_process_masks(lowerCamelCase_       ,np.array(lowerCamelCase_				)       ,np.array(lowerCamelCase_				)				)




@require_vision
@require_tf
class 					_lowercase						(   unittest.TestCase ):
      '''simple docstring'''


      def 		lowerCAmelCase__       (					self				)      ->      Union[str, Any]:



         '''simple docstring'''



         UpperCAmelCase__				:					Dict 					=	tempfile.mkdtemp()
         UpperCAmelCase__				:					List[Any] 					=	SamImageProcessor()
         UpperCAmelCase__				:					Dict 					=	SamProcessor(lowerCamelCase_				)
         processor.save_pretrained(self.tmpdirname				)




      def 		lowerCAmelCase__       (					self       ,**lowerCamelCase_				)      ->      Union[str, Any]:



         '''simple docstring'''



         return AutoProcessor.from_pretrained(self.tmpdirname       ,**lowerCamelCase_				).image_processor




      def 		lowerCAmelCase__       (					self				)      ->      Tuple:



         '''simple docstring'''



         shutil.rmtree(self.tmpdirname				)




      def 		lowerCAmelCase__       (					self				)      ->      Union[str, Any]:



         '''simple docstring'''



         UpperCAmelCase__				:					List[Any] 					=	[np.random.randint(255       ,size=(3, 30, 400)       ,dtype=np.uinta				)]

         UpperCAmelCase__				:					Optional[Any] 					=	[Image.fromarray(np.moveaxis(lowerCamelCase_       ,0       ,-1				)				) for x in image_inputs]

         return image_inputs




      def 		lowerCAmelCase__       (					self				)      ->      List[Any]:



         '''simple docstring'''



         UpperCAmelCase__				:					Union[str, Any] 					=	SamProcessor(image_processor=self.get_image_processor()				)
         processor.save_pretrained(self.tmpdirname				)

         UpperCAmelCase__				:					Optional[Any] 					=	self.get_image_processor(do_normalize=lowerCamelCase_       ,padding_value=1.0				)

         UpperCAmelCase__				:					Tuple 					=	SamProcessor.from_pretrained(self.tmpdirname       ,do_normalize=lowerCamelCase_       ,padding_value=1.0				)

         self.assertEqual(processor.image_processor.to_json_string()       ,image_processor_add_kwargs.to_json_string()				)
         self.assertIsInstance(processor.image_processor       ,lowerCamelCase_				)




      def 		lowerCAmelCase__       (					self				)      ->      int:



         '''simple docstring'''



         UpperCAmelCase__				:					Any 					=	self.get_image_processor()

         UpperCAmelCase__				:					Dict 					=	SamProcessor(image_processor=lowerCamelCase_				)

         UpperCAmelCase__				:					Dict 					=	self.prepare_image_inputs()

         UpperCAmelCase__				:					Optional[int] 					=	image_processor(lowerCamelCase_       ,return_tensors='''np'''				)
         UpperCAmelCase__				:					int 					=	processor(images=lowerCamelCase_       ,return_tensors='''np'''				)

         input_feat_extract.pop('''original_sizes'''				)  # pop original_sizes as it is popped in the processor
         input_feat_extract.pop('''reshaped_input_sizes'''				)  # pop reshaped_input_sizes as it is popped in the processor

         for key in input_feat_extract.keys():
            self.assertAlmostEqual(input_feat_extract[key].sum()       ,input_processor[key].sum()       ,delta=1e-2				)




      @require_tf
      def 		lowerCAmelCase__       (					self				)      ->      Optional[Any]:



         '''simple docstring'''



         UpperCAmelCase__				:					List[str] 					=	self.get_image_processor()

         UpperCAmelCase__				:					List[str] 					=	SamProcessor(image_processor=lowerCamelCase_				)
         UpperCAmelCase__				:					Dict 					=	[tf.ones((1, 3, 5, 5)				)]

         UpperCAmelCase__				:					Dict 					=	[[1764, 2646]]

         UpperCAmelCase__				:					Union[str, Any] 					=	[[683, 1024]]
         UpperCAmelCase__				:					Optional[int] 					=	processor.post_process_masks(lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,return_tensors='''tf'''				)
         self.assertEqual(masks[0].shape       ,(1, 3, 1764, 2646)				)

         UpperCAmelCase__				:					Optional[int] 					=	processor.post_process_masks(
             lowerCamelCase_       ,tf.convert_to_tensor(lowerCamelCase_				)       ,tf.convert_to_tensor(lowerCamelCase_				)       ,return_tensors='''tf'''       ,)
         self.assertEqual(masks[0].shape       ,(1, 3, 1764, 2646)				)

         # should also work with np
         UpperCAmelCase__				:					int 					=	[np.ones((1, 3, 5, 5)				)]
         UpperCAmelCase__				:					Optional[Any] 					=	processor.post_process_masks(
             lowerCamelCase_       ,np.array(lowerCamelCase_				)       ,np.array(lowerCamelCase_				)       ,return_tensors='''tf'''				)

         self.assertEqual(masks[0].shape       ,(1, 3, 1764, 2646)				)

         UpperCAmelCase__				:					str 					=	[[1, 0], [0, 1]]
         with self.assertRaises(tf.errors.InvalidArgumentError				):
            UpperCAmelCase__				:					int 					=	processor.post_process_masks(
                lowerCamelCase_       ,np.array(lowerCamelCase_				)       ,np.array(lowerCamelCase_				)       ,return_tensors='''tf'''				)




@require_vision
@require_torchvision
class 					_lowercase						(   unittest.TestCase ):
      '''simple docstring'''


      def 		lowerCAmelCase__       (					self				)      ->      List[str]:



         '''simple docstring'''



         UpperCAmelCase__				:					List[Any] 					=	tempfile.mkdtemp()
         UpperCAmelCase__				:					Dict 					=	SamImageProcessor()
         UpperCAmelCase__				:					Tuple 					=	SamProcessor(lowerCamelCase_				)
         processor.save_pretrained(self.tmpdirname				)




      def 		lowerCAmelCase__       (					self       ,**lowerCamelCase_				)      ->      Union[str, Any]:



         '''simple docstring'''



         return AutoProcessor.from_pretrained(self.tmpdirname       ,**lowerCamelCase_				).image_processor




      def 		lowerCAmelCase__       (					self				)      ->      Union[str, Any]:



         '''simple docstring'''



         shutil.rmtree(self.tmpdirname				)




      def 		lowerCAmelCase__       (					self				)      ->      Optional[Any]:



         '''simple docstring'''



         UpperCAmelCase__				:					Tuple 					=	[np.random.randint(255       ,size=(3, 30, 400)       ,dtype=np.uinta				)]

         UpperCAmelCase__				:					Any 					=	[Image.fromarray(np.moveaxis(lowerCamelCase_       ,0       ,-1				)				) for x in image_inputs]

         return image_inputs




      @is_pt_tf_cross_test
      def 		lowerCAmelCase__       (					self				)      ->      int:



         '''simple docstring'''



         UpperCAmelCase__				:					Optional[int] 					=	self.get_image_processor()

         UpperCAmelCase__				:					Dict 					=	SamProcessor(image_processor=lowerCamelCase_				)
         UpperCAmelCase__				:					Optional[int] 					=	np.random.randint(0       ,2       ,size=(1, 3, 5, 5)				).astype(np.floataa				)
         UpperCAmelCase__				:					Dict 					=	[tf.convert_to_tensor(lowerCamelCase_				)]
         UpperCAmelCase__				:					Any 					=	[torch.tensor(lowerCamelCase_				)]

         UpperCAmelCase__				:					List[str] 					=	[[1764, 2646]]

         UpperCAmelCase__				:					Dict 					=	[[683, 1024]]
         UpperCAmelCase__				:					int 					=	processor.post_process_masks(
             lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,return_tensors='''tf'''				)
         UpperCAmelCase__				:					Optional[Any] 					=	processor.post_process_masks(
             lowerCamelCase_       ,lowerCamelCase_       ,lowerCamelCase_       ,return_tensors='''pt'''				)

         self.assertTrue(np.all(tf_masks[0].numpy() == pt_masks[0].numpy()				)				)




      @is_pt_tf_cross_test
      def 		lowerCAmelCase__       (					self				)      ->      str:



         '''simple docstring'''



         UpperCAmelCase__				:					Union[str, Any] 					=	self.get_image_processor()

         UpperCAmelCase__				:					str 					=	SamProcessor(image_processor=lowerCamelCase_				)

         UpperCAmelCase__				:					str 					=	self.prepare_image_inputs()

         UpperCAmelCase__				:					int 					=	image_processor(lowerCamelCase_       ,return_tensors='''pt'''				)['''pixel_values'''].numpy()
         UpperCAmelCase__				:					Optional[int] 					=	processor(images=lowerCamelCase_       ,return_tensors='''pt'''				)['''pixel_values'''].numpy()

         UpperCAmelCase__				:					List[str] 					=	image_processor(lowerCamelCase_       ,return_tensors='''tf'''				)['''pixel_values'''].numpy()
         UpperCAmelCase__				:					Tuple 					=	processor(images=lowerCamelCase_       ,return_tensors='''tf'''				)['''pixel_values'''].numpy()

         self.assertTrue(np.allclose(lowerCamelCase_       ,lowerCamelCase_				)				)
         self.assertTrue(np.allclose(lowerCamelCase_       ,lowerCamelCase_				)				)
         self.assertTrue(np.allclose(lowerCamelCase_       ,lowerCamelCase_				)				)





'''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 					_lowercase						(   lowerCAmelCase ):
      '''simple docstring'''


      UpperCAmelCase_   :   int   				= RobertaConfig
      UpperCAmelCase_   :   Tuple   				= '''roberta'''




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



         '''simple docstring'''



         super().__init__(lowerCamelCase_				)

         UpperCAmelCase__				:					Optional[Any] 					=	RobertaEmbeddings(lowerCamelCase_				)
         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 					_lowercase						(   lowerCAmelCase ):
      '''simple docstring'''


      UpperCAmelCase_   :   Tuple   				= RobertaConfig
      UpperCAmelCase_   :   Optional[Any]   				= '''roberta'''




      def __init__(					self       ,lowerCamelCase_				)      ->      str:



         '''simple docstring'''



         super().__init__(lowerCamelCase_				)
         UpperCAmelCase__				:					Any 					=	config.num_labels
         UpperCAmelCase__				:					Union[str, Any] 					=	config.num_hidden_layers

         UpperCAmelCase__				:					str 					=	DeeRobertaModel(lowerCamelCase_				)
         UpperCAmelCase__				:					Optional[int] 					=	nn.Dropout(config.hidden_dropout_prob				)
         UpperCAmelCase__				:					Union[str, Any] 					=	nn.Linear(config.hidden_size       ,self.config.num_labels				)




      @add_start_docstrings_to_model_forward(lowerCamelCase_				)
      def 		lowerCAmelCase__       (					self       ,lowerCamelCase_=None       ,lowerCamelCase_=None       ,lowerCamelCase_=None       ,lowerCamelCase_=None       ,lowerCamelCase_=None       ,lowerCamelCase_=None       ,lowerCamelCase_=None       ,lowerCamelCase_=-1       ,lowerCamelCase_=False       ,)      ->      Optional[Any]:



         '''simple docstring'''



         UpperCAmelCase__				:					Tuple 					=	self.num_layers
         try:
            UpperCAmelCase__				:					Union[str, Any] 					=	self.roberta(
                lowerCamelCase_       ,attention_mask=lowerCamelCase_       ,token_type_ids=lowerCamelCase_       ,position_ids=lowerCamelCase_       ,head_mask=lowerCamelCase_       ,inputs_embeds=lowerCamelCase_       ,)

            UpperCAmelCase__				:					List[Any] 					=	outputs[1]

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

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

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

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

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




