from argparse import ArgumentParser

from . import BaseTransformersCLICommand




def       __lowerCAmelCase      (						A_   : Union[str, Any]					)    ->  List[str]:
  return DownloadCommand(args.model						,	args.cache_dir						,	args.force						,	args.trust_remote_code					)





class 			UpperCAmelCase__    (		snake_case					):



  """simple docstring"""






  @staticmethod
  def 	_UpperCAmelCase				(				__lowerCAmelCase:			ArgumentParser   ) -> Optional[int]:




    '''simple docstring'''




    __UpperCAmelCase						    =       parser.add_parser("download"   )
    download_parser.add_argument(
        "--cache-dir"				,				type=__lowerCAmelCase				,				default=__lowerCAmelCase				,				help="Path to location to store the models"   )
    download_parser.add_argument(
        "--force"				,				action="store_true"				,				help="Force the model to be download even if already in cache-dir"   )
    download_parser.add_argument(
        "--trust-remote-code"				,				action="store_true"				,				help="Whether or not to allow for custom models defined on the Hub in their own modeling files. Use only if you've reviewed the code as it will execute on your local machine"				,				)
    download_parser.add_argument("model"				,				type=__lowerCAmelCase				,				help="Name of the model to download"   )
    download_parser.set_defaults(func=__lowerCAmelCase   )






  def __init__(				self:			Dict				,				__lowerCAmelCase:			str				,				__lowerCAmelCase:			str				,				__lowerCAmelCase:			bool				,				__lowerCAmelCase:			bool   ) -> Optional[Any]:




    '''simple docstring'''




    __UpperCAmelCase						    =       model
    __UpperCAmelCase						    =       cache
    __UpperCAmelCase						    =       force
    __UpperCAmelCase						    =       trust_remote_code






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




    '''simple docstring'''




    from ..models.auto import AutoModel, AutoTokenizer

    AutoModel.from_pretrained(
        self._model				,				cache_dir=self._cache				,				force_download=self._force				,				trust_remote_code=self._trust_remote_code   )
    AutoTokenizer.from_pretrained(
        self._model				,				cache_dir=self._cache				,				force_download=self._force				,				trust_remote_code=self._trust_remote_code   )




import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss

from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
from transformers.models.bert.modeling_bert import (
    BERT_INPUTS_DOCSTRING,
    BERT_START_DOCSTRING,
    BertEmbeddings,
    BertLayer,
    BertPooler,
    BertPreTrainedModel,
)




def       __lowerCAmelCase      (						A_   : Optional[Any]					)    ->  List[str]:
  __UpperCAmelCase						    =       torch.exp(A_					)
  __UpperCAmelCase						    =       torch.sum(A_						,	dim=1					)  # sum of exp(x_i)
  __UpperCAmelCase						    =       torch.sum(x * exp_x						,	dim=1					)  # sum of x_i * exp(x_i)
  return torch.log(A_					) - B / A


class 			UpperCAmelCase__    (		nn.Module					):



  """simple docstring"""






  def __init__(				self:			Optional[Any]				,				__lowerCAmelCase:			int   ) -> List[Any]:




    '''simple docstring'''




    super().__init__()
    __UpperCAmelCase						    =       config.output_attentions
    __UpperCAmelCase						    =       config.output_hidden_states
    __UpperCAmelCase						    =       nn.ModuleList([BertLayer(__lowerCAmelCase   ) for _ in range(config.num_hidden_layers   )]   )
    __UpperCAmelCase						    =       nn.ModuleList([BertHighway(__lowerCAmelCase   ) for _ in range(config.num_hidden_layers   )]   )

    __UpperCAmelCase						    =       [-1 for _ in range(config.num_hidden_layers   )]






  def 	_UpperCAmelCase				(				self:			Optional[Any]				,				__lowerCAmelCase:			List[str]   ) -> Optional[Any]:




    '''simple docstring'''




    if (type(__lowerCAmelCase   ) is float) or (type(__lowerCAmelCase   ) is int):
      for i in range(len(self.early_exit_entropy   )   ):
        __UpperCAmelCase						    =       x
    else:
      __UpperCAmelCase						    =       x






  def 	_UpperCAmelCase				(				self:			Union[str, Any]				,				__lowerCAmelCase:			Optional[int]   ) -> Optional[Any]:




    '''simple docstring'''




    __UpperCAmelCase						    =       pooler.state_dict()
    for highway in self.highway:
      for name, param in highway.pooler.state_dict().items():
        param.copy_(loaded_model[name]   )






  def 	_UpperCAmelCase				(				self:			Dict				,				__lowerCAmelCase:			Optional[int]				,				__lowerCAmelCase:			Union[str, Any]=None				,				__lowerCAmelCase:			Any=None				,				__lowerCAmelCase:			Optional[int]=None				,				__lowerCAmelCase:			Optional[Any]=None				,				) -> List[str]:




    '''simple docstring'''




    __UpperCAmelCase						    =       ()
    __UpperCAmelCase						    =       ()
    __UpperCAmelCase						    =       ()
    for i, layer_module in enumerate(self.layer   ):
      if self.output_hidden_states:
        __UpperCAmelCase						    =       all_hidden_states + (hidden_states,)

      __UpperCAmelCase						    =       layer_module(
          __lowerCAmelCase				,				__lowerCAmelCase				,				head_mask[i]				,				__lowerCAmelCase				,				__lowerCAmelCase   )
      __UpperCAmelCase						    =       layer_outputs[0]

      if self.output_attentions:
        __UpperCAmelCase						    =       all_attentions + (layer_outputs[1],)

      __UpperCAmelCase						    =       (hidden_states,)
      if self.output_hidden_states:
        __UpperCAmelCase						    =       current_outputs + (all_hidden_states,)
      if self.output_attentions:
        __UpperCAmelCase						    =       current_outputs + (all_attentions,)

      __UpperCAmelCase						    =       self.highway[i](__lowerCAmelCase   )
      # logits, pooled_output

      if not self.training:
        __UpperCAmelCase						    =       highway_exit[0]
        __UpperCAmelCase						    =       entropy(__lowerCAmelCase   )
        __UpperCAmelCase						    =       highway_exit + (highway_entropy,)  # logits, hidden_states(?), entropy
        __UpperCAmelCase						    =       all_highway_exits + (highway_exit,)

        if highway_entropy < self.early_exit_entropy[i]:
          __UpperCAmelCase						    =       (highway_logits,) + current_outputs[1:] + (all_highway_exits,)
          raise HighwayException(__lowerCAmelCase				,				i + 1   )
      else:
        __UpperCAmelCase						    =       all_highway_exits + (highway_exit,)

        # Add last layer
    if self.output_hidden_states:
      __UpperCAmelCase						    =       all_hidden_states + (hidden_states,)

    __UpperCAmelCase						    =       (hidden_states,)
    if self.output_hidden_states:
      __UpperCAmelCase						    =       outputs + (all_hidden_states,)
    if self.output_attentions:
      __UpperCAmelCase						    =       outputs + (all_attentions,)

    __UpperCAmelCase						    =       outputs + (all_highway_exits,)
    return outputs  # last-layer hidden state, (all hidden states), (all attentions), all highway exits






@add_start_docstrings(
    'The Bert Model transformer with early exiting (DeeBERT). '		,	snake_case		,	)
class 			UpperCAmelCase__    (		snake_case					):



  """simple docstring"""






  def __init__(				self:			Union[str, Any]				,				__lowerCAmelCase:			Optional[Any]   ) -> int:




    '''simple docstring'''




    super().__init__(__lowerCAmelCase   )
    __UpperCAmelCase						    =       config

    __UpperCAmelCase						    =       BertEmbeddings(__lowerCAmelCase   )
    __UpperCAmelCase						    =       DeeBertEncoder(__lowerCAmelCase   )
    __UpperCAmelCase						    =       BertPooler(__lowerCAmelCase   )

    self.init_weights()






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




    '''simple docstring'''




    self.encoder.init_highway_pooler(self.pooler   )






  def 	_UpperCAmelCase				(				self:			List[str]   ) -> Tuple:




    '''simple docstring'''




    return self.embeddings.word_embeddings






  def 	_UpperCAmelCase				(				self:			Optional[Any]				,				__lowerCAmelCase:			List[str]   ) -> Dict:




    '''simple docstring'''




    __UpperCAmelCase						    =       value






  def 	_UpperCAmelCase				(				self:			Union[str, Any]				,				__lowerCAmelCase:			Optional[int]   ) -> Any:




    '''simple docstring'''




    for layer, heads in heads_to_prune.items():
      self.encoder.layer[layer].attention.prune_heads(__lowerCAmelCase   )






  @add_start_docstrings_to_model_forward(__lowerCAmelCase   )
  def 	_UpperCAmelCase				(				self:			Optional[int]				,				__lowerCAmelCase:			Dict=None				,				__lowerCAmelCase:			Tuple=None				,				__lowerCAmelCase:			List[str]=None				,				__lowerCAmelCase:			Union[str, Any]=None				,				__lowerCAmelCase:			str=None				,				__lowerCAmelCase:			Dict=None				,				__lowerCAmelCase:			Optional[int]=None				,				__lowerCAmelCase:			Tuple=None				,				) -> List[str]:




    '''simple docstring'''




    if input_ids is not None and inputs_embeds is not None:
      raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time"   )
    elif input_ids is not None:
      __UpperCAmelCase						    =       input_ids.size()
    elif inputs_embeds is not None:
      __UpperCAmelCase						    =       inputs_embeds.size()[:-1]
    else:
      raise ValueError("You have to specify either input_ids or inputs_embeds"   )

    __UpperCAmelCase						    =       input_ids.device if input_ids is not None else inputs_embeds.device

    if attention_mask is None:
      __UpperCAmelCase						    =       torch.ones(__lowerCAmelCase				,				device=__lowerCAmelCase   )
    if encoder_attention_mask is None:
      __UpperCAmelCase						    =       torch.ones(__lowerCAmelCase				,				device=__lowerCAmelCase   )
    if token_type_ids is None:
      __UpperCAmelCase						    =       torch.zeros(__lowerCAmelCase				,				dtype=torch.long				,				device=__lowerCAmelCase   )

    # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
    # ourselves in which case we just need to make it broadcastable to all heads.
    __UpperCAmelCase						    =       self.get_extended_attention_mask(__lowerCAmelCase				,				__lowerCAmelCase				,				__lowerCAmelCase   )

    # If a 2D ou 3D attention mask is provided for the cross-attention
    # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
    if encoder_attention_mask.dim() == 3:
      __UpperCAmelCase						    =       encoder_attention_mask[:, None, :, :]
    if encoder_attention_mask.dim() == 2:
      __UpperCAmelCase						    =       encoder_attention_mask[:, None, None, :]

    __UpperCAmelCase						    =       encoder_extended_attention_mask.to(
        dtype=next(self.parameters()   ).dtype   )  # fp16 compatibility
    __UpperCAmelCase						    =       (1.0 - encoder_extended_attention_mask) * -10000.0

    # Prepare head mask if needed
    # 1.0 in head_mask indicate we keep the head
    # attention_probs has shape bsz x n_heads x N x N
    # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
    # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
    __UpperCAmelCase						    =       self.get_head_mask(__lowerCAmelCase				,				self.config.num_hidden_layers   )

    __UpperCAmelCase						    =       self.embeddings(
        input_ids=__lowerCAmelCase				,				position_ids=__lowerCAmelCase				,				token_type_ids=__lowerCAmelCase				,				inputs_embeds=__lowerCAmelCase   )
    __UpperCAmelCase						    =       self.encoder(
        __lowerCAmelCase				,				attention_mask=__lowerCAmelCase				,				head_mask=__lowerCAmelCase				,				encoder_hidden_states=__lowerCAmelCase				,				encoder_attention_mask=__lowerCAmelCase				,				)
    __UpperCAmelCase						    =       encoder_outputs[0]
    __UpperCAmelCase						    =       self.pooler(__lowerCAmelCase   )

    __UpperCAmelCase						    =       (
        sequence_output,
        pooled_output,
    ) + encoder_outputs[
        1:
    ]  # add hidden_states and attentions if they are here
    return outputs  # sequence_output, pooled_output, (hidden_states), (attentions), highway exits






class 			UpperCAmelCase__    (		snake_case					):



  """simple docstring"""






  def __init__(				self:			Dict				,				__lowerCAmelCase:			Dict				,				__lowerCAmelCase:			List[str]   ) -> Any:




    '''simple docstring'''




    __UpperCAmelCase						    =       message
    __UpperCAmelCase						    =       exit_layer  # start from 1!






class 			UpperCAmelCase__    (		nn.Module					):



  """simple docstring"""






  def __init__(				self:			Union[str, Any]				,				__lowerCAmelCase:			int   ) -> Optional[Any]:




    '''simple docstring'''




    super().__init__()
    __UpperCAmelCase						    =       BertPooler(__lowerCAmelCase   )
    __UpperCAmelCase						    =       nn.Dropout(config.hidden_dropout_prob   )
    __UpperCAmelCase						    =       nn.Linear(config.hidden_size				,				config.num_labels   )






  def 	_UpperCAmelCase				(				self:			Optional[int]				,				__lowerCAmelCase:			Any   ) -> Tuple:




    '''simple docstring'''




    __UpperCAmelCase						    =       encoder_outputs[0]
    __UpperCAmelCase						    =       self.pooler(__lowerCAmelCase   )
    # "return" pooler_output

    # BertModel
    __UpperCAmelCase						    =       (pooler_input, pooler_output) + encoder_outputs[1:]
    # "return" bmodel_output

    # Dropout and classification
    __UpperCAmelCase						    =       bmodel_output[1]

    __UpperCAmelCase						    =       self.dropout(__lowerCAmelCase   )
    __UpperCAmelCase						    =       self.classifier(__lowerCAmelCase   )

    return logits, pooled_output






@add_start_docstrings(
    'Bert Model (with early exiting - DeeBERT) with a classifier on top,\n    also takes care of multi-layer training. '		,	snake_case		,	)
class 			UpperCAmelCase__    (		snake_case					):



  """simple docstring"""






  def __init__(				self:			Optional[Any]				,				__lowerCAmelCase:			List[str]   ) -> str:




    '''simple docstring'''




    super().__init__(__lowerCAmelCase   )
    __UpperCAmelCase						    =       config.num_labels
    __UpperCAmelCase						    =       config.num_hidden_layers

    __UpperCAmelCase						    =       DeeBertModel(__lowerCAmelCase   )
    __UpperCAmelCase						    =       nn.Dropout(config.hidden_dropout_prob   )
    __UpperCAmelCase						    =       nn.Linear(config.hidden_size				,				self.config.num_labels   )

    self.init_weights()






  @add_start_docstrings_to_model_forward(__lowerCAmelCase   )
  def 	_UpperCAmelCase				(				self:			Any				,				__lowerCAmelCase:			Any=None				,				__lowerCAmelCase:			Optional[int]=None				,				__lowerCAmelCase:			Optional[int]=None				,				__lowerCAmelCase:			List[str]=None				,				__lowerCAmelCase:			Dict=None				,				__lowerCAmelCase:			int=None				,				__lowerCAmelCase:			str=None				,				__lowerCAmelCase:			Optional[Any]=-1				,				__lowerCAmelCase:			str=False				,				) -> Optional[int]:




    '''simple docstring'''




    __UpperCAmelCase						    =       self.num_layers
    try:
      __UpperCAmelCase						    =       self.bert(
          __lowerCAmelCase				,				attention_mask=__lowerCAmelCase				,				token_type_ids=__lowerCAmelCase				,				position_ids=__lowerCAmelCase				,				head_mask=__lowerCAmelCase				,				inputs_embeds=__lowerCAmelCase				,				)
      # sequence_output, pooled_output, (hidden_states), (attentions), highway exits

      __UpperCAmelCase						    =       outputs[1]

      __UpperCAmelCase						    =       self.dropout(__lowerCAmelCase   )
      __UpperCAmelCase						    =       self.classifier(__lowerCAmelCase   )
      __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(__lowerCAmelCase   )
      __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(__lowerCAmelCase   )
          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(__lowerCAmelCase   )

      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), (highway_exits)



