




'''simple docstring'''





from collections import OrderedDict
from typing import Any, List, Mapping, Optional

from ... import PreTrainedTokenizer, TensorType, is_torch_available
from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast, PatchingSpec
from ...utils import logging


UpperCamelCase_		:     Dict          =		logging.get_logger(__name__)


UpperCamelCase_		:     Dict          =		{
    '''Salesforce/codegen-350M-nl''': '''https://huggingface.co/Salesforce/codegen-350M-nl/resolve/main/config.json''',
    '''Salesforce/codegen-350M-multi''': '''https://huggingface.co/Salesforce/codegen-350M-multi/resolve/main/config.json''',
    '''Salesforce/codegen-350M-mono''': '''https://huggingface.co/Salesforce/codegen-350M-mono/resolve/main/config.json''',
    '''Salesforce/codegen-2B-nl''': '''https://huggingface.co/Salesforce/codegen-2B-nl/resolve/main/config.json''',
    '''Salesforce/codegen-2B-multi''': '''https://huggingface.co/Salesforce/codegen-2B-multi/resolve/main/config.json''',
    '''Salesforce/codegen-2B-mono''': '''https://huggingface.co/Salesforce/codegen-2B-mono/resolve/main/config.json''',
    '''Salesforce/codegen-6B-nl''': '''https://huggingface.co/Salesforce/codegen-6B-nl/resolve/main/config.json''',
    '''Salesforce/codegen-6B-multi''': '''https://huggingface.co/Salesforce/codegen-6B-multi/resolve/main/config.json''',
    '''Salesforce/codegen-6B-mono''': '''https://huggingface.co/Salesforce/codegen-6B-mono/resolve/main/config.json''',
    '''Salesforce/codegen-16B-nl''': '''https://huggingface.co/Salesforce/codegen-16B-nl/resolve/main/config.json''',
    '''Salesforce/codegen-16B-multi''': '''https://huggingface.co/Salesforce/codegen-16B-multi/resolve/main/config.json''',
    '''Salesforce/codegen-16B-mono''': '''https://huggingface.co/Salesforce/codegen-16B-mono/resolve/main/config.json''',
}

class   _a							(		__lowerCAmelCase				):
     SCREAMING_SNAKE_CASE_				:       Optional[Any]       							=		"""codegen"""
     SCREAMING_SNAKE_CASE_				:       int       							=		{
         """max_position_embeddings""": """n_positions""",
         """hidden_size""": """n_embd""",
         """num_attention_heads""": """n_head""",
         """num_hidden_layers""": """n_layer""",
     }






     def __init__(					self   ,_SCREAMING_SNAKE_CASE=50_400   ,_SCREAMING_SNAKE_CASE=2_048   ,_SCREAMING_SNAKE_CASE=2_048   ,_SCREAMING_SNAKE_CASE=4_096   ,_SCREAMING_SNAKE_CASE=28   ,_SCREAMING_SNAKE_CASE=16   ,_SCREAMING_SNAKE_CASE=64   ,_SCREAMING_SNAKE_CASE=None   ,_SCREAMING_SNAKE_CASE="gelu_new"   ,_SCREAMING_SNAKE_CASE=0.0   ,_SCREAMING_SNAKE_CASE=0.0   ,_SCREAMING_SNAKE_CASE=0.0   ,_SCREAMING_SNAKE_CASE=1e-5   ,_SCREAMING_SNAKE_CASE=0.0_2   ,_SCREAMING_SNAKE_CASE=True   ,_SCREAMING_SNAKE_CASE=50_256   ,_SCREAMING_SNAKE_CASE=50_256   ,_SCREAMING_SNAKE_CASE=False   ,**_SCREAMING_SNAKE_CASE   ,)				->     Union[str, Any]:
          _snake_case				= vocab_size
          _snake_case				= n_ctx
          _snake_case				= n_positions
          _snake_case				= n_embd
          _snake_case				= n_layer
          _snake_case				= n_head
          _snake_case				= n_inner
          _snake_case				= rotary_dim
          _snake_case				= activation_function
          _snake_case				= resid_pdrop
          _snake_case				= embd_pdrop
          _snake_case				= attn_pdrop
          _snake_case				= layer_norm_epsilon
          _snake_case				= initializer_range
          _snake_case				= use_cache

          _snake_case				= bos_token_id
          _snake_case				= eos_token_id

          super().__init__(
              bos_token_id=_SCREAMING_SNAKE_CASE   ,eos_token_id=_SCREAMING_SNAKE_CASE   ,tie_word_embeddings=_SCREAMING_SNAKE_CASE   ,**_SCREAMING_SNAKE_CASE       )







class   _a							(		__lowerCAmelCase				):






     def __init__(					self   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE = "default"   ,_SCREAMING_SNAKE_CASE = None   ,_SCREAMING_SNAKE_CASE = False   ,)				->     Tuple:
          super().__init__(_SCREAMING_SNAKE_CASE   ,task=_SCREAMING_SNAKE_CASE   ,patching_specs=_SCREAMING_SNAKE_CASE   ,use_past=_SCREAMING_SNAKE_CASE       )
          if not getattr(self._config   ,"pad_token_id"   ,_SCREAMING_SNAKE_CASE       ):
               # TODO: how to do that better?
               _snake_case				= 0






     @property
     def 							_lowercase      (					self       )				->     Mapping[str, Mapping[int, str]]:
          _snake_case				= OrderedDict({"input_ids": {0: "batch", 1: "sequence"}}       )
          if self.use_past:
               self.fill_with_past_key_values_(_SCREAMING_SNAKE_CASE   ,direction="inputs"       )
               _snake_case				= {0: "batch", 1: "past_sequence + sequence"}
          else:
               _snake_case				= {0: "batch", 1: "sequence"}

          return common_inputs






     @property
     def 							_lowercase      (					self       )				->     int:
          return self._config.n_layer






     @property
     def 							_lowercase      (					self       )				->     int:
          return self._config.n_head






     def 							_lowercase      (					self   ,_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE = -1   ,_SCREAMING_SNAKE_CASE = -1   ,_SCREAMING_SNAKE_CASE = False   ,_SCREAMING_SNAKE_CASE = None   ,)				->     Mapping[str, Any]:
          _snake_case				= super(_SCREAMING_SNAKE_CASE   ,self       ).generate_dummy_inputs(
              _SCREAMING_SNAKE_CASE   ,batch_size=_SCREAMING_SNAKE_CASE   ,seq_length=_SCREAMING_SNAKE_CASE   ,is_pair=_SCREAMING_SNAKE_CASE   ,framework=_SCREAMING_SNAKE_CASE       )

          # We need to order the input in the way they appears in the forward()
          _snake_case				= OrderedDict({"input_ids": common_inputs["input_ids"]}       )

          # Need to add the past_keys
          if self.use_past:
               if not is_torch_available():
                    raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed."       )
               else:
                    import torch

                    _snake_case				,   _snake_case				= common_inputs["input_ids"].shape
                    # Not using the same length for past_key_values
                    _snake_case				= seqlen + 2
                    _snake_case				= (
                        batch,
                        self.num_attention_heads,
                        past_key_values_length,
                        self._config.hidden_size // self.num_attention_heads,
                    )
                    _snake_case				= [
                        (torch.zeros(_SCREAMING_SNAKE_CASE       ), torch.zeros(_SCREAMING_SNAKE_CASE       )) for _ in range(self.num_layers       )
                    ]

          _snake_case				= common_inputs["attention_mask"]
          if self.use_past:
               _snake_case				= ordered_inputs["attention_mask"].dtype
               _snake_case				= torch.cat(
                   [ordered_inputs["attention_mask"], torch.ones(_SCREAMING_SNAKE_CASE   ,_SCREAMING_SNAKE_CASE   ,dtype=_SCREAMING_SNAKE_CASE       )]   ,dim=1       )

          return ordered_inputs






     @property
     def 							_lowercase      (					self       )				->     int:
          return 13








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


