
from dataclasses import asdict, dataclass
from typing import Optional

from ...configuration_utils import PretrainedConfig
from ...utils import logging


SCREAMING_SNAKE_CASE				:Dict									=							logging.get_logger(__name__)

# TODO Update this
SCREAMING_SNAKE_CASE				:Dict									=							{
    'facebook/esm-1b': 'https://huggingface.co/facebook/esm-1b/resolve/main/config.json',
    # See all ESM models at https://huggingface.co/models?filter=esm
}





class 				UpperCAmelCase (    __SCREAMING_SNAKE_CASE    ):




	'''simple docstring'''





	snake_case_          =	"esm"






	def __init__(			self   :					List[str]				,A   :					Dict=None				,A   :					Tuple=None				,A   :					Any=None				,A   :					Optional[Any]=7_68				,A   :					Tuple=12				,A   :					List[str]=12				,A   :					Tuple=30_72				,A   :					List[str]=0.1				,A   :					List[Any]=0.1				,A   :					int=10_26				,A   :					List[str]=0.02				,A   :					Union[str, Any]=1E-12				,A   :					List[Any]="absolute"				,A   :					List[Any]=True				,A   :					Union[str, Any]=None				,A   :					Optional[int]=False				,A   :					Dict=False				,A   :					Tuple=None				,A   :					Optional[int]=None				,**A   :					List[Any]				,):
					super().__init__(pad_token_id=A				,mask_token_id=A				,**A					)

					__A       =					vocab_size
					__A       =					hidden_size
					__A       =					num_hidden_layers
					__A       =					num_attention_heads
					__A       =					intermediate_size
					__A       =					hidden_dropout_prob
					__A       =					attention_probs_dropout_prob
					__A       =					max_position_embeddings
					__A       =					initializer_range
					__A       =					layer_norm_eps
					__A       =					position_embedding_type
					__A       =					use_cache
					__A       =					emb_layer_norm_before
					__A       =					token_dropout
					__A       =					is_folding_model
					if is_folding_model:
									if esmfold_config is None:
													logger.info("No esmfold_config supplied for folding model, using default values."					)
													__A       =					EsmFoldConfig()
									elif isinstance(A				,A					):
													__A       =					EsmFoldConfig(**A					)
									__A       =					esmfold_config
									if vocab_list is None:
													logger.warning("No vocab_list supplied for folding model, assuming the ESM-2 vocabulary!"					)
													__A       =					get_default_vocab_list()
									else:
													__A       =					vocab_list
					else:
									__A       =					None
									__A       =					None
					if self.esmfold_config is not None and getattr(self.esmfold_config				,"use_esm_attn_map"				,A					):
									raise ValueError("The HuggingFace port of ESMFold does not support use_esm_attn_map at this time!"					)







	def 							UpperCamelCase_    (			self   :					Optional[int]					):
					__A       =					super().to_dict()
					if isinstance(self.esmfold_config				,A					):
									__A       =					self.esmfold_config.to_dict()
					return output





@dataclass
class 				UpperCAmelCase :




	'''simple docstring'''





	snake_case_          =	None
	snake_case_          =	True
	snake_case_          =	False
	snake_case_          =	False
	snake_case_          =	False
	snake_case_          =	0

	snake_case_          =	True
	snake_case_          =	False

	snake_case_          =	128
	snake_case_          =	None






	def 							UpperCamelCase_    (			self   :					List[Any]					):
					if self.trunk is None:
									__A       =					TrunkConfig()
					elif isinstance(self.trunk				,A					):
									__A       =					TrunkConfig(**self.trunk					)







	def 							UpperCamelCase_    (			self   :					Optional[Any]					):
					__A       =					asdict(self					)
					__A       =					self.trunk.to_dict()
					return output





@dataclass
class 				UpperCAmelCase :




	'''simple docstring'''





	snake_case_          =	48
	snake_case_          =	1024
	snake_case_          =	128
	snake_case_          =	32
	snake_case_          =	32
	snake_case_          =	32
	snake_case_          =	0
	snake_case_          =	0
	snake_case_          =	False
	snake_case_          =	4
	snake_case_          =	128
	snake_case_          =	None






	def 							UpperCamelCase_    (			self   :					List[Any]					):
					if self.structure_module is None:
									__A       =					StructureModuleConfig()
					elif isinstance(self.structure_module				,A					):
									__A       =					StructureModuleConfig(**self.structure_module					)

					if self.max_recycles <= 0:
									raise ValueError(f'''`max_recycles` should be positive, got {self.max_recycles}.'''					)
					if self.sequence_state_dim % self.sequence_state_dim != 0:
									raise ValueError(
									    "`sequence_state_dim` should be a round multiple of `sequence_state_dim`, got"
									    f''' {self.sequence_state_dim} and {self.sequence_state_dim}.'''					)
					if self.pairwise_state_dim % self.pairwise_state_dim != 0:
									raise ValueError(
									    "`pairwise_state_dim` should be a round multiple of `pairwise_state_dim`, got"
									    f''' {self.pairwise_state_dim} and {self.pairwise_state_dim}.'''					)

					__A       =					self.sequence_state_dim // self.sequence_head_width
					__A       =					self.pairwise_state_dim // self.pairwise_head_width

					if self.sequence_state_dim != sequence_num_heads * self.sequence_head_width:
									raise ValueError(
									    "`sequence_state_dim` should be equal to `sequence_num_heads * sequence_head_width, got"
									    f''' {self.sequence_state_dim} != {sequence_num_heads} * {self.sequence_head_width}.'''					)
					if self.pairwise_state_dim != pairwise_num_heads * self.pairwise_head_width:
									raise ValueError(
									    "`pairwise_state_dim` should be equal to `pairwise_num_heads * pairwise_head_width, got"
									    f''' {self.pairwise_state_dim} != {pairwise_num_heads} * {self.pairwise_head_width}.'''					)
					if self.pairwise_state_dim % 2 != 0:
									raise ValueError(f'''`pairwise_state_dim` should be even, got {self.pairwise_state_dim}.'''					)

					if self.dropout >= 0.4:
									raise ValueError(f'''`dropout` should not be greater than 0.4, got {self.dropout}.'''					)







	def 							UpperCamelCase_    (			self   :					Tuple					):
					__A       =					asdict(self					)
					__A       =					self.structure_module.to_dict()
					return output





@dataclass
class 				UpperCAmelCase :




	'''simple docstring'''





	snake_case_          =	384
	snake_case_          =	128
	snake_case_          =	16
	snake_case_          =	128
	snake_case_          =	12
	snake_case_          =	4
	snake_case_          =	8
	snake_case_          =	0.1
	snake_case_          =	8
	snake_case_          =	1
	snake_case_          =	2
	snake_case_          =	7
	snake_case_          =	10
	snake_case_          =	1E-8
	snake_case_          =	1E5






	def 							UpperCamelCase_    (			self   :					Union[str, Any]					):
					return asdict(self					)





def 		UpperCAmelCase     (      )						->      int:







				"""simple docstring"""



				return (
				    "<cls>",
				    "<pad>",
				    "<eos>",
				    "<unk>",
				    "L",
				    "A",
				    "G",
				    "V",
				    "S",
				    "E",
				    "R",
				    "T",
				    "I",
				    "D",
				    "P",
				    "K",
				    "Q",
				    "N",
				    "F",
				    "Y",
				    "M",
				    "H",
				    "W",
				    "C",
				    "X",
				    "B",
				    "U",
				    "Z",
				    "O",
				    ".",
				    "-",
				    "<null_1>",
				    "<mask>",
				)








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 		lowercase__  (				_UpperCamelCase) ->			List[str]:
      """simple docstring"""

      UpperCamelCase					       =  torch.exp(_UpperCamelCase)
      UpperCamelCase					       =  torch.sum(_UpperCamelCase       ,							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(_UpperCamelCase) - B / A





class 					A__     (     nn.Module ):

     '''simple docstring'''






     def __init__( self      :    Union[str, Any]   ,						_SCREAMING_SNAKE_CASE      :    Any ):






           """simple docstring"""




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

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







     def   _SCREAMING_SNAKE_CASE     ( self      :    Optional[int]   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any] ):






           """simple docstring"""




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







     def   _SCREAMING_SNAKE_CASE     ( self      :    str   ,						_SCREAMING_SNAKE_CASE      :    List[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   _SCREAMING_SNAKE_CASE     ( self      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    Union[str, Any]   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]=None   ,						_SCREAMING_SNAKE_CASE      :    int=None   ,						_SCREAMING_SNAKE_CASE      :    int=None   ,						_SCREAMING_SNAKE_CASE      :    List[Any]=None   ,						):






           """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(
                     _SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						head_mask[i]   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )
                 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](_SCREAMING_SNAKE_CASE )
                 # logits, pooled_output

                 if not self.training:
                       UpperCamelCase					       =  highway_exit[0]
                       UpperCamelCase					       =  entropy(_SCREAMING_SNAKE_CASE )
                       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(_SCREAMING_SNAKE_CASE   ,						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 					A__     (     __snake_case ):

     '''simple docstring'''






     def __init__( self      :    Any   ,						_SCREAMING_SNAKE_CASE      :    int ):






           """simple docstring"""




           super().__init__(_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  config

           UpperCamelCase					       =  BertEmbeddings(_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  DeeBertEncoder(_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  BertPooler(_SCREAMING_SNAKE_CASE )

           self.init_weights()







     def   _SCREAMING_SNAKE_CASE     ( self      :    int ):






           """simple docstring"""




           self.encoder.init_highway_pooler(self.pooler )







     def   _SCREAMING_SNAKE_CASE     ( self      :    str ):






           """simple docstring"""




           return self.embeddings.word_embeddings







     def   _SCREAMING_SNAKE_CASE     ( self      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    Dict ):






           """simple docstring"""




           UpperCamelCase					       =  value







     def   _SCREAMING_SNAKE_CASE     ( self      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    Tuple ):






           """simple docstring"""




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







     @add_start_docstrings_to_model_forward(_SCREAMING_SNAKE_CASE )
     def   _SCREAMING_SNAKE_CASE     ( self      :    Union[str, Any]   ,						_SCREAMING_SNAKE_CASE      :    List[Any]=None   ,						_SCREAMING_SNAKE_CASE      :    List[str]=None   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]=None   ,						_SCREAMING_SNAKE_CASE      :    Tuple=None   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]=None   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any]=None   ,						_SCREAMING_SNAKE_CASE      :    List[Any]=None   ,						_SCREAMING_SNAKE_CASE      :    Dict=None   ,						):






           """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(_SCREAMING_SNAKE_CASE   ,						device=_SCREAMING_SNAKE_CASE )
           if encoder_attention_mask is None:
                 UpperCamelCase					       =  torch.ones(_SCREAMING_SNAKE_CASE   ,						device=_SCREAMING_SNAKE_CASE )
           if token_type_ids is None:
                 UpperCamelCase					       =  torch.zeros(_SCREAMING_SNAKE_CASE   ,						dtype=torch.long   ,						device=_SCREAMING_SNAKE_CASE )

           # 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(_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE   ,						_SCREAMING_SNAKE_CASE )

           # 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) * -1_0_0_0_0.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(_SCREAMING_SNAKE_CASE   ,						self.config.num_hidden_layers )

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

           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 					A__     (     __snake_case ):

     '''simple docstring'''






     def __init__( self      :    Optional[Any]   ,						_SCREAMING_SNAKE_CASE      :    List[Any]   ,						_SCREAMING_SNAKE_CASE      :    int ):






           """simple docstring"""




           UpperCamelCase					       =  message
           UpperCamelCase					       =  exit_layer  # start from 1!





class 					A__     (     nn.Module ):

     '''simple docstring'''






     def __init__( self      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    str ):






           """simple docstring"""




           super().__init__()
           UpperCamelCase					       =  BertPooler(_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  nn.Dropout(config.hidden_dropout_prob )
           UpperCamelCase					       =  nn.Linear(config.hidden_size   ,						config.num_labels )







     def   _SCREAMING_SNAKE_CASE     ( self      :    Tuple   ,						_SCREAMING_SNAKE_CASE      :    Optional[Any] ):






           """simple docstring"""




           UpperCamelCase					       =  encoder_outputs[0]
           UpperCamelCase					       =  self.pooler(_SCREAMING_SNAKE_CASE )
           # "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(_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  self.classifier(_SCREAMING_SNAKE_CASE )

           return logits, pooled_output





@add_start_docstrings(
    """Bert Model (with early exiting - DeeBERT) with a classifier on top,
    also takes care of multi-layer training. """   ,    __snake_case   ,    )
class 					A__     (     __snake_case ):

     '''simple docstring'''






     def __init__( self      :    List[str]   ,						_SCREAMING_SNAKE_CASE      :    Dict ):






           """simple docstring"""




           super().__init__(_SCREAMING_SNAKE_CASE )
           UpperCamelCase					       =  config.num_labels
           UpperCamelCase					       =  config.num_hidden_layers

           UpperCamelCase					       =  DeeBertModel(_SCREAMING_SNAKE_CASE )
           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(_SCREAMING_SNAKE_CASE )
     def   _SCREAMING_SNAKE_CASE     ( self      :    str   ,						_SCREAMING_SNAKE_CASE      :    Any=None   ,						_SCREAMING_SNAKE_CASE      :    int=None   ,						_SCREAMING_SNAKE_CASE      :    Tuple=None   ,						_SCREAMING_SNAKE_CASE      :    List[Any]=None   ,						_SCREAMING_SNAKE_CASE      :    List[str]=None   ,						_SCREAMING_SNAKE_CASE      :    Dict=None   ,						_SCREAMING_SNAKE_CASE      :    List[str]=None   ,						_SCREAMING_SNAKE_CASE      :    str=-1   ,						_SCREAMING_SNAKE_CASE      :    str=False   ,						):






           """simple docstring"""




           UpperCamelCase					       =  self.num_layers
           try:
                 UpperCamelCase					       =  self.bert(
                     _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   ,						)
                 # sequence_output, pooled_output, (hidden_states), (attentions), highway exits

                 UpperCamelCase					       =  outputs[1]

                 UpperCamelCase					       =  self.dropout(_SCREAMING_SNAKE_CASE )
                 UpperCamelCase					       =  self.classifier(_SCREAMING_SNAKE_CASE )
                 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(_SCREAMING_SNAKE_CASE )
                 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(_SCREAMING_SNAKE_CASE )
                             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(_SCREAMING_SNAKE_CASE )

                 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)



