


'''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). '      ,							UpperCAmelCase      ,							)
class        lowerCAmelCase_  (     UpperCAmelCase	):
       __UpperCAmelCase					=RobertaConfig
       __UpperCAmelCase					='roberta'






       def __init__( self			,     _UpperCamelCase   )->      Optional[int]:
            super().__init__(_UpperCamelCase   )

            _A				=							RobertaEmbeddings(_UpperCamelCase   )
            self.init_weights()






@add_start_docstrings(
    'RoBERTa Model (with early exiting - DeeRoBERTa) with a classifier on top,\n    also takes care of multi-layer training. '      ,							UpperCAmelCase      ,							)
class        lowerCAmelCase_  (     UpperCAmelCase	):
       __UpperCAmelCase					=RobertaConfig
       __UpperCAmelCase					='roberta'






       def __init__( self			,     _UpperCamelCase   )->      List[Any]:
            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   UpperCamelCase						( self			,     _UpperCamelCase=None			,     _UpperCamelCase=None			,     _UpperCamelCase=None			,     _UpperCamelCase=None			,     _UpperCamelCase=None			,     _UpperCamelCase=None			,     _UpperCamelCase=None			,     _UpperCamelCase=-1			,     _UpperCamelCase=False			,     )->      Tuple:
            _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









'''simple docstring'''
import flax.linen as nn
import jax.numpy as jnp

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD






class        lowerCAmelCase_  (     nn.Module	):
       __UpperCAmelCase					=42
       __UpperCAmelCase					=42
       __UpperCAmelCase					=0.0
       __UpperCAmelCase					=1
       __UpperCAmelCase					=1
       __UpperCAmelCase					=True
       __UpperCAmelCase					=False
       __UpperCAmelCase					=False
       __UpperCAmelCase					=False
       __UpperCAmelCase					=jnp.floataa






       def   UpperCamelCase						( self   )->      int:
            _A				=							[]
            _A				=							[]

            for i in range(self.num_layers   ):
                 _A				=							self.in_channels if i == 0 else self.out_channels

                 _A				=							FlaxResnetBlockaD(
                     in_channels=_UpperCamelCase			,     out_channels=self.out_channels			,     dropout_prob=self.dropout			,     dtype=self.dtype			,     )
                 resnets.append(_UpperCamelCase   )

                 _A				=							FlaxTransformeraDModel(
                     in_channels=self.out_channels			,     n_heads=self.num_attention_heads			,     d_head=self.out_channels // self.num_attention_heads			,     depth=1			,     use_linear_projection=self.use_linear_projection			,     only_cross_attention=self.only_cross_attention			,     use_memory_efficient_attention=self.use_memory_efficient_attention			,     dtype=self.dtype			,     )
                 attentions.append(_UpperCamelCase   )

            _A				=							resnets
            _A				=							attentions

            if self.add_downsample:
                 _A				=							FlaxDownsampleaD(self.out_channels			,     dtype=self.dtype   )






       def __call__( self			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase=True   )->      Dict:
            _A				=							()

            for resnet, attn in zip(self.resnets			,     self.attentions   ):
                 _A				=							resnet(_UpperCamelCase			,     _UpperCamelCase			,     deterministic=_UpperCamelCase   )
                 _A				=							attn(_UpperCamelCase			,     _UpperCamelCase			,     deterministic=_UpperCamelCase   )
                 output_states += (hidden_states,)

            if self.add_downsample:
                 _A				=							self.downsamplers_a(_UpperCamelCase   )
                 output_states += (hidden_states,)

            return hidden_states, output_states






class        lowerCAmelCase_  (     nn.Module	):
       __UpperCAmelCase					=42
       __UpperCAmelCase					=42
       __UpperCAmelCase					=0.0
       __UpperCAmelCase					=1
       __UpperCAmelCase					=True
       __UpperCAmelCase					=jnp.floataa






       def   UpperCamelCase						( self   )->      Optional[Any]:
            _A				=							[]

            for i in range(self.num_layers   ):
                 _A				=							self.in_channels if i == 0 else self.out_channels

                 _A				=							FlaxResnetBlockaD(
                     in_channels=_UpperCamelCase			,     out_channels=self.out_channels			,     dropout_prob=self.dropout			,     dtype=self.dtype			,     )
                 resnets.append(_UpperCamelCase   )
            _A				=							resnets

            if self.add_downsample:
                 _A				=							FlaxDownsampleaD(self.out_channels			,     dtype=self.dtype   )






       def __call__( self			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase=True   )->      Union[str, Any]:
            _A				=							()

            for resnet in self.resnets:
                 _A				=							resnet(_UpperCamelCase			,     _UpperCamelCase			,     deterministic=_UpperCamelCase   )
                 output_states += (hidden_states,)

            if self.add_downsample:
                 _A				=							self.downsamplers_a(_UpperCamelCase   )
                 output_states += (hidden_states,)

            return hidden_states, output_states






class        lowerCAmelCase_  (     nn.Module	):
       __UpperCAmelCase					=42
       __UpperCAmelCase					=42
       __UpperCAmelCase					=42
       __UpperCAmelCase					=0.0
       __UpperCAmelCase					=1
       __UpperCAmelCase					=1
       __UpperCAmelCase					=True
       __UpperCAmelCase					=False
       __UpperCAmelCase					=False
       __UpperCAmelCase					=False
       __UpperCAmelCase					=jnp.floataa






       def   UpperCamelCase						( self   )->      str:
            _A				=							[]
            _A				=							[]

            for i in range(self.num_layers   ):
                 _A				=							self.in_channels if (i == self.num_layers - 1) else self.out_channels
                 _A				=							self.prev_output_channel if i == 0 else self.out_channels

                 _A				=							FlaxResnetBlockaD(
                     in_channels=resnet_in_channels + res_skip_channels			,     out_channels=self.out_channels			,     dropout_prob=self.dropout			,     dtype=self.dtype			,     )
                 resnets.append(_UpperCamelCase   )

                 _A				=							FlaxTransformeraDModel(
                     in_channels=self.out_channels			,     n_heads=self.num_attention_heads			,     d_head=self.out_channels // self.num_attention_heads			,     depth=1			,     use_linear_projection=self.use_linear_projection			,     only_cross_attention=self.only_cross_attention			,     use_memory_efficient_attention=self.use_memory_efficient_attention			,     dtype=self.dtype			,     )
                 attentions.append(_UpperCamelCase   )

            _A				=							resnets
            _A				=							attentions

            if self.add_upsample:
                 _A				=							FlaxUpsampleaD(self.out_channels			,     dtype=self.dtype   )






       def __call__( self			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase=True   )->      Dict:
            for resnet, attn in zip(self.resnets			,     self.attentions   ):
                 # pop res hidden states
                 _A				=							res_hidden_states_tuple[-1]
                 _A				=							res_hidden_states_tuple[:-1]
                 _A				=							jnp.concatenate((hidden_states, res_hidden_states)			,     axis=-1   )

                 _A				=							resnet(_UpperCamelCase			,     _UpperCamelCase			,     deterministic=_UpperCamelCase   )
                 _A				=							attn(_UpperCamelCase			,     _UpperCamelCase			,     deterministic=_UpperCamelCase   )

            if self.add_upsample:
                 _A				=							self.upsamplers_a(_UpperCamelCase   )

            return hidden_states






class        lowerCAmelCase_  (     nn.Module	):
       __UpperCAmelCase					=42
       __UpperCAmelCase					=42
       __UpperCAmelCase					=42
       __UpperCAmelCase					=0.0
       __UpperCAmelCase					=1
       __UpperCAmelCase					=True
       __UpperCAmelCase					=jnp.floataa






       def   UpperCamelCase						( self   )->      Optional[Any]:
            _A				=							[]

            for i in range(self.num_layers   ):
                 _A				=							self.in_channels if (i == self.num_layers - 1) else self.out_channels
                 _A				=							self.prev_output_channel if i == 0 else self.out_channels

                 _A				=							FlaxResnetBlockaD(
                     in_channels=resnet_in_channels + res_skip_channels			,     out_channels=self.out_channels			,     dropout_prob=self.dropout			,     dtype=self.dtype			,     )
                 resnets.append(_UpperCamelCase   )

            _A				=							resnets

            if self.add_upsample:
                 _A				=							FlaxUpsampleaD(self.out_channels			,     dtype=self.dtype   )






       def __call__( self			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase=True   )->      Dict:
            for resnet in self.resnets:
                 # pop res hidden states
                 _A				=							res_hidden_states_tuple[-1]
                 _A				=							res_hidden_states_tuple[:-1]
                 _A				=							jnp.concatenate((hidden_states, res_hidden_states)			,     axis=-1   )

                 _A				=							resnet(_UpperCamelCase			,     _UpperCamelCase			,     deterministic=_UpperCamelCase   )

            if self.add_upsample:
                 _A				=							self.upsamplers_a(_UpperCamelCase   )

            return hidden_states






class        lowerCAmelCase_  (     nn.Module	):
       __UpperCAmelCase					=42
       __UpperCAmelCase					=0.0
       __UpperCAmelCase					=1
       __UpperCAmelCase					=1
       __UpperCAmelCase					=False
       __UpperCAmelCase					=False
       __UpperCAmelCase					=jnp.floataa






       def   UpperCamelCase						( self   )->      Dict:
            # there is always at least one resnet
            _A				=							[
                FlaxResnetBlockaD(
                    in_channels=self.in_channels			,     out_channels=self.in_channels			,     dropout_prob=self.dropout			,     dtype=self.dtype			,     )
            ]

            _A				=							[]

            for _ in range(self.num_layers   ):
                 _A				=							FlaxTransformeraDModel(
                     in_channels=self.in_channels			,     n_heads=self.num_attention_heads			,     d_head=self.in_channels // self.num_attention_heads			,     depth=1			,     use_linear_projection=self.use_linear_projection			,     use_memory_efficient_attention=self.use_memory_efficient_attention			,     dtype=self.dtype			,     )
                 attentions.append(_UpperCamelCase   )

                 _A				=							FlaxResnetBlockaD(
                     in_channels=self.in_channels			,     out_channels=self.in_channels			,     dropout_prob=self.dropout			,     dtype=self.dtype			,     )
                 resnets.append(_UpperCamelCase   )

            _A				=							resnets
            _A				=							attentions






       def __call__( self			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase			,     _UpperCamelCase=True   )->      List[Any]:
            _A				=							self.resnets[0](_UpperCamelCase			,     _UpperCamelCase   )
            for attn, resnet in zip(self.attentions			,     self.resnets[1:]   ):
                 _A				=							attn(_UpperCamelCase			,     _UpperCamelCase			,     deterministic=_UpperCamelCase   )
                 _A				=							resnet(_UpperCamelCase			,     _UpperCamelCase			,     deterministic=_UpperCamelCase   )

            return hidden_states





