




from maths.prime_check import is_prime




def       __lowerCAmelCase      (						A_   : int					)    ->  int:
  if not isinstance(A_						,	A_					):
    __UpperCAmelCase						    =       F'''Input value of [number={number}] must be an integer'''
    raise TypeError(A_					)
  if is_prime(A_					) and is_prime(number + 2					):
    return number + 2
  else:
    return -1


if __name__ == "__main__":
       import doctest

       doctest.testmod()




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)



