





'''simple docstring'''






import inspect
import unittest

import numpy as np

from transformers import ViTConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor


if is_flax_available():
      import jax

      from transformers.models.vit.modeling_flax_vit import FlaxViTForImageClassification, FlaxViTModel







class 		UpperCAmelCase   (				unittest.TestCase):


  """simple docstring"""







  def __init__(				self    :					int							,	UpperCamelCase__    :					Dict							,	UpperCamelCase__    :					Union[str, Any]=13							,	UpperCamelCase__    :					Optional[Any]=30							,	UpperCamelCase__    :					List[str]=2							,	UpperCamelCase__    :					Tuple=3							,	UpperCamelCase__    :					Dict=True							,	UpperCamelCase__    :					Optional[Any]=True							,	UpperCamelCase__    :					str=32							,	UpperCamelCase__    :					Union[str, Any]=5							,	UpperCamelCase__    :					List[str]=4							,	UpperCamelCase__    :					Union[str, Any]=37							,	UpperCamelCase__    :					Optional[Any]="gelu"							,	UpperCamelCase__    :					Any=0.1							,	UpperCamelCase__    :					List[str]=0.1							,	UpperCamelCase__    :					int=10							,	UpperCamelCase__    :					str=0.02							,	)     ->						int:
    _UpperCamelCase       			=parent
    _UpperCamelCase       			=batch_size
    _UpperCamelCase       			=image_size
    _UpperCamelCase       			=patch_size
    _UpperCamelCase       			=num_channels
    _UpperCamelCase       			=is_training
    _UpperCamelCase       			=use_labels
    _UpperCamelCase       			=hidden_size
    _UpperCamelCase       			=num_hidden_layers
    _UpperCamelCase       			=num_attention_heads
    _UpperCamelCase       			=intermediate_size
    _UpperCamelCase       			=hidden_act
    _UpperCamelCase       			=hidden_dropout_prob
    _UpperCamelCase       			=attention_probs_dropout_prob
    _UpperCamelCase       			=type_sequence_label_size
    _UpperCamelCase       			=initializer_range

    # in ViT, the seq length equals the number of patches + 1 (we add 1 for the [CLS] token)
    _UpperCamelCase       			=(image_size // patch_size) ** 2
    _UpperCamelCase       			=num_patches + 1


  def   UpperCamelCase__      (				self    :					int  )     ->						str:
    _UpperCamelCase       			=floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]  )

    _UpperCamelCase       			=ViTConfig(
        image_size=self.image_size							,	patch_size=self.patch_size							,	num_channels=self.num_channels							,	hidden_size=self.hidden_size							,	num_hidden_layers=self.num_hidden_layers							,	num_attention_heads=self.num_attention_heads							,	intermediate_size=self.intermediate_size							,	hidden_act=self.hidden_act							,	hidden_dropout_prob=self.hidden_dropout_prob							,	attention_probs_dropout_prob=self.attention_probs_dropout_prob							,	is_decoder=UpperCamelCase__							,	initializer_range=self.initializer_range							,	)

    return config, pixel_values


  def   UpperCamelCase__      (				self    :					str							,	UpperCamelCase__    :					List[str]							,	UpperCamelCase__    :					str  )     ->						Union[str, Any]:
    _UpperCamelCase       			=FlaxViTModel(config=UpperCamelCase__  )
    _UpperCamelCase       			=model(UpperCamelCase__  )
    # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token)
    _UpperCamelCase       			=(self.image_size, self.image_size)
    _UpperCamelCase       			=(self.patch_size, self.patch_size)
    _UpperCamelCase       			=(image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
    self.parent.assertEqual(result.last_hidden_state.shape							,	(self.batch_size, num_patches + 1, self.hidden_size)  )


  def   UpperCamelCase__      (				self    :					Optional[Any]							,	UpperCamelCase__    :					Any							,	UpperCamelCase__    :					Optional[int]  )     ->						Any:
    _UpperCamelCase       			=self.type_sequence_label_size
    _UpperCamelCase       			=FlaxViTForImageClassification(config=UpperCamelCase__  )
    _UpperCamelCase       			=model(UpperCamelCase__  )
    self.parent.assertEqual(result.logits.shape							,	(self.batch_size, self.type_sequence_label_size)  )

    # test greyscale images
    _UpperCamelCase       			=1
    _UpperCamelCase       			=FlaxViTForImageClassification(UpperCamelCase__  )

    _UpperCamelCase       			=floats_tensor([self.batch_size, 1, self.image_size, self.image_size]  )
    _UpperCamelCase       			=model(UpperCamelCase__  )




  def   UpperCamelCase__      (				self    :					Optional[Any]  )     ->						int:
    _UpperCamelCase       			=self.prepare_config_and_inputs()
    (
        (
        _UpperCamelCase
    )			,     (
        _UpperCamelCase
    )			,     
    )       			=config_and_inputs
    _UpperCamelCase       			={'''pixel_values''': pixel_values}
    return config, inputs_dict







@require_flax
class 		UpperCAmelCase   (				lowercase_   ,     unittest.TestCase):


  """simple docstring"""







  lowerCAmelCase_					=				(FlaxViTModel, FlaxViTForImageClassification) if is_flax_available() else ()


  def   UpperCamelCase__      (				self    :					Dict  )     ->						None:
    _UpperCamelCase       			=FlaxViTModelTester(self  )
    _UpperCamelCase       			=ConfigTester(self							,	config_class=UpperCamelCase__							,	has_text_modality=UpperCamelCase__							,	hidden_size=37  )


  def   UpperCamelCase__      (				self    :					Dict  )     ->						int:
    self.config_tester.run_common_tests()


  def   UpperCamelCase__      (				self    :					Tuple  )     ->						Union[str, Any]:
    _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_model(*UpperCamelCase__  )


  def   UpperCamelCase__      (				self    :					List[str]  )     ->						Any:
    _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs()
    self.model_tester.create_and_check_for_image_classification(*UpperCamelCase__  )


  def   UpperCamelCase__      (				self    :					Any  )     ->						Optional[Any]:
    _UpperCamelCase			,     _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs_for_common()

    for model_class in self.all_model_classes:
      _UpperCamelCase       			=model_class(UpperCamelCase__  )
      _UpperCamelCase       			=inspect.signature(model.__call__  )
      # signature.parameters is an OrderedDict => so arg_names order is deterministic
      _UpperCamelCase       			=[*signature.parameters.keys()]

      _UpperCamelCase       			=['''pixel_values''']
      self.assertListEqual(arg_names[:1]							,	UpperCamelCase__  )


  def   UpperCamelCase__      (				self    :					int  )     ->						List[Any]:
    _UpperCamelCase			,     _UpperCamelCase       			=self.model_tester.prepare_config_and_inputs_for_common()

    for model_class in self.all_model_classes:
      with self.subTest(model_class.__name__  ):
        _UpperCamelCase       			=self._prepare_for_class(UpperCamelCase__							,	UpperCamelCase__  )
        _UpperCamelCase       			=model_class(UpperCamelCase__  )

        @jax.jit
        def model_jitted(UpperCamelCase__    :					str							,	**UpperCamelCase__    :					List[Any]  ):
          return model(pixel_values=UpperCamelCase__							,	**UpperCamelCase__  )

        with self.subTest('''JIT Enabled'''  ):
          _UpperCamelCase       			=model_jitted(**UpperCamelCase__  ).to_tuple()

        with self.subTest('''JIT Disabled'''  ):
          with jax.disable_jit():
            _UpperCamelCase       			=model_jitted(**UpperCamelCase__  ).to_tuple()

        self.assertEqual(len(UpperCamelCase__  )							,	len(UpperCamelCase__  )  )
        for jitted_output, output in zip(UpperCamelCase__							,	UpperCamelCase__  ):
          self.assertEqual(jitted_output.shape							,	output.shape  )




  @slow
  def   UpperCamelCase__      (				self    :					Tuple  )     ->						Any:
    for model_class_name in self.all_model_classes:
      _UpperCamelCase       			=model_class_name.from_pretrained('''google/vit-base-patch16-224'''  )
      _UpperCamelCase       			=model(np.ones((1, 3, 224, 224)  )  )
      self.assertIsNotNone(UpperCamelCase__  )













'''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






