





import os
from shutil import copyfile
from typing import Any, Dict, List, Optional, Tuple

import sentencepiece as spm

from ...tokenization_utils import AddedToken, BatchEncoding, PreTrainedTokenizer
from ...utils import logging


UpperCAmelCase_					=					logging.get_logger(__name__)

UpperCAmelCase_					=					'▁'

UpperCAmelCase_					=					{'vocab_file': 'sentencepiece.bpe.model'}

UpperCAmelCase_					=					{
    'vocab_file': {
        'facebook/mbart-large-50-one-to-many-mmt': (
            'https://huggingface.co/facebook/mbart-large-50-one-to-many-mmt/resolve/main/sentencepiece.bpe.model'
        ),
    }
}

UpperCAmelCase_					=					{
    'facebook/mbart-large-50-one-to-many-mmt': 1024,
}

# fmt: off
UpperCAmelCase_					=					['ar_AR', 'cs_CZ', 'de_DE', 'en_XX', 'es_XX', 'et_EE', 'fi_FI', 'fr_XX', 'gu_IN', 'hi_IN', 'it_IT', 'ja_XX', 'kk_KZ', 'ko_KR', 'lt_LT', 'lv_LV', 'my_MM', 'ne_NP', 'nl_XX', 'ro_RO', 'ru_RU', 'si_LK', 'tr_TR', 'vi_VN', 'zh_CN', 'af_ZA', 'az_AZ', 'bn_IN', 'fa_IR', 'he_IL', 'hr_HR', 'id_ID', 'ka_GE', 'km_KH', 'mk_MK', 'ml_IN', 'mn_MN', 'mr_IN', 'pl_PL', 'ps_AF', 'pt_XX', 'sv_SE', 'sw_KE', 'ta_IN', 'te_IN', 'th_TH', 'tl_XX', 'uk_UA', 'ur_PK', 'xh_ZA', 'gl_ES', 'sl_SI']






class        lowercase__			(      A__							):

   '''simple docstring'''


   a		:				int  						=						VOCAB_FILES_NAMES
   a		:				Optional[int]  						=						PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
   a		:				List[str]  						=						PRETRAINED_VOCAB_FILES_MAP
   a		:				Union[str, Any]  						=						["input_ids", "attention_mask"]

   a		:				Optional[Any]  						=						[]
   a		:				Tuple  						=						[]




   def __init__(       self,						__magic_name__,						__magic_name__=None,						__magic_name__=None,						__magic_name__="</s>",						__magic_name__="</s>",						__magic_name__="<s>",						__magic_name__="<unk>",						__magic_name__="<pad>",						__magic_name__="<mask>",						__magic_name__ = None,						**__magic_name__,						)					-> None:






       """simple docstring"""
       # Mask token behave like a normal word, i.e. include the space before it
       UpperCamelCase__				:							Tuple			  = AddedToken(snake_case_,						lstrip=snake_case_,						rstrip=snake_case_				) if isinstance(snake_case_,						snake_case_				) else mask_token

       UpperCamelCase__				:							List[str]			  = {} if sp_model_kwargs is None else sp_model_kwargs

       UpperCamelCase__				:							int			  = kwargs.get('''additional_special_tokens''',						[]				)
       kwargs["additional_special_tokens"] += [
           code for code in FAIRSEQ_LANGUAGE_CODES if code not in kwargs["additional_special_tokens"]
       ]

       super().__init__(
           src_lang=snake_case_,						tgt_lang=snake_case_,						eos_token=snake_case_,						unk_token=snake_case_,						sep_token=snake_case_,						cls_token=snake_case_,						pad_token=snake_case_,						mask_token=snake_case_,						sp_model_kwargs=self.sp_model_kwargs,						**snake_case_,						)

       UpperCamelCase__				:							int			  = spm.SentencePieceProcessor(**self.sp_model_kwargs				)
       self.sp_model.Load(str(snake_case_				)				)
       UpperCamelCase__				:							Dict			  = vocab_file

       # Original fairseq vocab and spm vocab must be "aligned":
       # Vocab    |    0    |    1    |   2    |    3    |  4  |  5  |  6  |   7   |   8   |  9
       # -------- | ------- | ------- | ------ | ------- | --- | --- | --- | ----- | ----- | ----
       # fairseq  | '<s>'   | '<pad>' | '</s>' | '<unk>' | ',' | '.' | '▁' | 's'   | '▁de' | '-'
       # spm      | '<unk>' | '<s>'   | '</s>' | ','     | '.' | '▁' | 's' | '▁de' | '-'   | '▁a'

       # Mimic fairseq token-to-id alignment for the first 4 token
       UpperCamelCase__				:							Any			  = {'''<s>''': 0, '''<pad>''': 1, '''</s>''': 2, '''<unk>''': 3}

       # The first "real" token "," has position 4 in the original fairseq vocab and position 3 in the spm vocab
       UpperCamelCase__				:							Any			  = 1

       UpperCamelCase__				:							Any			  = len(self.sp_model				)
       UpperCamelCase__				:							Union[str, Any]			  = {
           code: self.sp_model_size + i + self.fairseq_offset for i, code in enumerate(snake_case_				)
       }
       UpperCamelCase__				:							Optional[int]			  = {v: k for k, v in self.lang_code_to_id.items()}
       UpperCamelCase__				:							str			  = len(self.sp_model				) + len(self.lang_code_to_id				) + self.fairseq_offset

       self.fairseq_tokens_to_ids.update(self.lang_code_to_id				)
       UpperCamelCase__				:							int			  = {v: k for k, v in self.fairseq_tokens_to_ids.items()}

       UpperCamelCase__				:							Any			  = src_lang if src_lang is not None else '''en_XX'''
       UpperCamelCase__				:							Optional[int]			  = self.lang_code_to_id[self._src_lang]
       UpperCamelCase__				:							Optional[int]			  = tgt_lang
       self.set_src_lang_special_tokens(self._src_lang				)




   @property
   def  UpperCamelCase__    (       self				)					-> int:






       """simple docstring"""
       return len(self.sp_model				) + len(self.lang_code_to_id				) + self.fairseq_offset + 1  # Plus 1 for the mask token




   @property
   def  UpperCamelCase__    (       self				)					-> str:






       """simple docstring"""
       return self._src_lang




   @src_lang.setter
   def  UpperCamelCase__    (       self,						__magic_name__				)					-> None:






       """simple docstring"""
       UpperCamelCase__				:							Tuple			  = new_src_lang
       self.set_src_lang_special_tokens(self._src_lang				)




   def __getstate__(       self				)					-> Dict:






       """simple docstring"""
       UpperCamelCase__				:							int			  = self.__dict__.copy()
       UpperCamelCase__				:							int			  = None
       return state




   def __setstate__(       self,						__magic_name__				)					-> None:






       """simple docstring"""
       UpperCamelCase__				:							int			  = d

       # for backward compatibility
       if not hasattr(self,						'''sp_model_kwargs'''				):
           UpperCamelCase__				:							Tuple			  = {}

       UpperCamelCase__				:							int			  = spm.SentencePieceProcessor(**self.sp_model_kwargs				)
       self.sp_model.Load(self.vocab_file				)




   def  UpperCamelCase__    (       self				)					-> Dict:






       """simple docstring"""
       UpperCamelCase__				:							Tuple			  = {self.convert_ids_to_tokens(snake_case_				): i for i in range(self.vocab_size				)}
       vocab.update(self.added_tokens_encoder				)
       return vocab




   def  UpperCamelCase__    (       self,						__magic_name__				)					-> List[str]:






       """simple docstring"""
       return self.sp_model.encode(snake_case_,						out_type=snake_case_				)




   def  UpperCamelCase__    (       self,						__magic_name__				)					-> int:






       """simple docstring"""
       if token in self.fairseq_tokens_to_ids:
           return self.fairseq_tokens_to_ids[token]
       UpperCamelCase__				:							Tuple			  = self.sp_model.PieceToId(snake_case_				)

       # Need to return unknown token if the SP model returned 0
       return spm_id + self.fairseq_offset if spm_id else self.unk_token_id




   def  UpperCamelCase__    (       self,						__magic_name__				)					-> str:






       """simple docstring"""
       if index in self.fairseq_ids_to_tokens:
           return self.fairseq_ids_to_tokens[index]
       return self.sp_model.IdToPiece(index - self.fairseq_offset				)




   def  UpperCamelCase__    (       self,						__magic_name__				)					-> Dict:






       """simple docstring"""
       UpperCamelCase__				:							Any			  = []
       UpperCamelCase__				:							List[str]			  = ''''''
       UpperCamelCase__				:							Union[str, Any]			  = False
       for token in tokens:
           # make sure that special tokens are not decoded using sentencepiece model
           if token in self.all_special_tokens:
               if not prev_is_special:
                   out_string += " "
               out_string += self.sp_model.decode(snake_case_				) + token
               UpperCamelCase__				:							Tuple			  = True
               UpperCamelCase__				:							Tuple			  = []
           else:
               current_sub_tokens.append(snake_case_				)
               UpperCamelCase__				:							List[Any]			  = False
       out_string += self.sp_model.decode(snake_case_				)
       return out_string.strip()




   def  UpperCamelCase__    (       self,						__magic_name__,						__magic_name__ = None				)					-> Tuple[str]:






       """simple docstring"""
       if not os.path.isdir(snake_case_				):
           logger.error(f"Vocabulary path ({save_directory}) should be a directory"				)
           return
       UpperCamelCase__				:							Any			  = os.path.join(
           snake_case_,						(filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']				)

       if os.path.abspath(self.vocab_file				) != os.path.abspath(snake_case_				) and os.path.isfile(self.vocab_file				):
           copyfile(self.vocab_file,						snake_case_				)
       elif not os.path.isfile(self.vocab_file				):
           with open(snake_case_,						'''wb'''				) as fi:
               UpperCamelCase__				:							Optional[Any]			  = self.sp_model.serialized_model_proto()
               fi.write(snake_case_				)

       return (out_vocab_file,)




   def  UpperCamelCase__    (       self,						__magic_name__,						__magic_name__ = None,						__magic_name__ = False				)					-> List[int]:






       """simple docstring"""

       if already_has_special_tokens:
           return super().get_special_tokens_mask(
               token_ids_a=snake_case_,						token_ids_a=snake_case_,						already_has_special_tokens=snake_case_				)

       UpperCamelCase__				:							Optional[int]			  = [1] * len(self.prefix_tokens				)
       UpperCamelCase__				:							Optional[int]			  = [1] * len(self.suffix_tokens				)
       if token_ids_a is None:
           return prefix_ones + ([0] * len(snake_case_				)) + suffix_ones
       return prefix_ones + ([0] * len(snake_case_				)) + ([0] * len(snake_case_				)) + suffix_ones




   def  UpperCamelCase__    (       self,						__magic_name__,						__magic_name__ = None				)					-> List[int]:






       """simple docstring"""
       if token_ids_a is None:
           return self.prefix_tokens + token_ids_a + self.suffix_tokens
       # We don't expect to process pairs, but leave the pair logic for API consistency
       return self.prefix_tokens + token_ids_a + token_ids_a + self.suffix_tokens




   def  UpperCamelCase__    (       self,						__magic_name__,						__magic_name__,						__magic_name__,						__magic_name__,						**__magic_name__				)					-> List[str]:






       """simple docstring"""
       if src_lang is None or tgt_lang is None:
           raise ValueError('''Translation requires a `src_lang` and a `tgt_lang` for this model'''				)
       UpperCamelCase__				:							Optional[Any]			  = src_lang
       UpperCamelCase__				:							Optional[Any]			  = self(snake_case_,						add_special_tokens=snake_case_,						return_tensors=snake_case_,						**snake_case_				)
       UpperCamelCase__				:							List[Any]			  = self.convert_tokens_to_ids(snake_case_				)
       UpperCamelCase__				:							Union[str, Any]			  = tgt_lang_id
       return inputs




   def  UpperCamelCase__    (       self,						__magic_name__,						__magic_name__ = "en_XX",						__magic_name__ = None,						__magic_name__ = "ro_RO",						**__magic_name__,						)					-> BatchEncoding:






       """simple docstring"""
       UpperCamelCase__				:							Optional[int]			  = src_lang
       UpperCamelCase__				:							str			  = tgt_lang
       return super().prepare_seqaseq_batch(snake_case_,						snake_case_,						**snake_case_				)




   def  UpperCamelCase__    (       self				)					-> int:






       """simple docstring"""
       return self.set_src_lang_special_tokens(self.src_lang				)




   def  UpperCamelCase__    (       self				)					-> int:






       """simple docstring"""
       return self.set_tgt_lang_special_tokens(self.tgt_lang				)




   def  UpperCamelCase__    (       self,						__magic_name__				)					-> None:






       """simple docstring"""
       UpperCamelCase__				:							Tuple			  = self.lang_code_to_id[src_lang]
       UpperCamelCase__				:							Optional[Any]			  = [self.cur_lang_code_id]
       UpperCamelCase__				:							Optional[Any]			  = [self.eos_token_id]




   def  UpperCamelCase__    (       self,						__magic_name__				)					-> None:






       """simple docstring"""
       UpperCamelCase__				:							List[Any]			  = self.lang_code_to_id[tgt_lang]
       UpperCamelCase__				:							List[Any]			  = [self.cur_lang_code_id]
       UpperCamelCase__				:							int			  = [self.eos_token_id]



"""simple docstring"""
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     (_lowerCAmelCase						):
						__lowerCAmelCase     =  torch.exp(_lowerCAmelCase						)
						__lowerCAmelCase     =  torch.sum(_lowerCAmelCase ,    dim=1						)  # sum of exp(x_i)
						__lowerCAmelCase     =  torch.sum(x * exp_x ,    dim=1						)  # sum of x_i * exp(x_i)
						return torch.log(_lowerCAmelCase						) - B / A
class       lowerCAmelCase_		(				nn.Module ):



							'''simple docstring'''





							def __init__( self    ,				snake_case_				)       ->  Optional[Any]:
													super().__init__()
													__lowerCAmelCase     =  config.output_attentions
													__lowerCAmelCase     =  config.output_hidden_states
													__lowerCAmelCase     =  nn.ModuleList([BertLayer(snake_case_				) for _ in range(config.num_hidden_layers				)]				)
													__lowerCAmelCase     =  nn.ModuleList([BertHighway(snake_case_				) for _ in range(config.num_hidden_layers				)]				)

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




							def    A__						( self    ,				snake_case_				)       ->  List[Any]:
													if (type(snake_case_				) is float) or (type(snake_case_				) is int):
																			for i in range(len(self.early_exit_entropy				)				):
																									__lowerCAmelCase     =  x
													else:
																			__lowerCAmelCase     =  x




							def    A__						( self    ,				snake_case_				)       ->  Dict:
													__lowerCAmelCase     =  pooler.state_dict()
													for highway in self.highway:
																			for name, param in highway.pooler.state_dict().items():
																									param.copy_(loaded_model[name]				)




							def    A__						( self    ,				snake_case_    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				)       ->  List[str]:
													__lowerCAmelCase     =  ()
													__lowerCAmelCase     =  ()
													__lowerCAmelCase     =  ()
													for i, layer_module in enumerate(self.layer				):
																			if self.output_hidden_states:
																									__lowerCAmelCase     =  all_hidden_states + (hidden_states,)

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

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

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

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

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

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

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

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

													__lowerCAmelCase     =  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). '''				,				A__				,				)
class       lowerCAmelCase_		(				A__ ):



							'''simple docstring'''





							def __init__( self    ,				snake_case_				)       ->  Tuple:
													super().__init__(snake_case_				)
													__lowerCAmelCase     =  config

													__lowerCAmelCase     =  BertEmbeddings(snake_case_				)
													__lowerCAmelCase     =  DeeBertEncoder(snake_case_				)
													__lowerCAmelCase     =  BertPooler(snake_case_				)

													self.init_weights()




							def    A__						( self				)       ->  Dict:
													self.encoder.init_highway_pooler(self.pooler				)




							def    A__						( self				)       ->  List[Any]:
													return self.embeddings.word_embeddings




							def    A__						( self    ,				snake_case_				)       ->  Tuple:
													__lowerCAmelCase     =  value




							def    A__						( self    ,				snake_case_				)       ->  List[str]:
													for layer, heads in heads_to_prune.items():
																			self.encoder.layer[layer].attention.prune_heads(snake_case_				)




							@add_start_docstrings_to_model_forward(snake_case_				)
							def    A__						( self    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				)       ->  Optional[Any]:
													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:
																			__lowerCAmelCase     =  input_ids.size()
													elif inputs_embeds is not None:
																			__lowerCAmelCase     =  inputs_embeds.size()[:-1]
													else:
																			raise ValueError("""You have to specify either input_ids or inputs_embeds"""				)

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

													if attention_mask is None:
																			__lowerCAmelCase     =  torch.ones(snake_case_    ,				device=snake_case_				)
													if encoder_attention_mask is None:
																			__lowerCAmelCase     =  torch.ones(snake_case_    ,				device=snake_case_				)
													if token_type_ids is None:
																			__lowerCAmelCase     =  torch.zeros(snake_case_    ,				dtype=torch.long    ,				device=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.
													__lowerCAmelCase     =  self.get_extended_attention_mask(snake_case_    ,				snake_case_    ,				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:
																			__lowerCAmelCase     =  encoder_attention_mask[:, None, :, :]
													if encoder_attention_mask.dim() == 2:
																			__lowerCAmelCase     =  encoder_attention_mask[:, None, None, :]

													__lowerCAmelCase     =  encoder_extended_attention_mask.to(
													    dtype=next(self.parameters()				).dtype				)  # fp16 compatibility
													__lowerCAmelCase     =  (1.0 - encoder_extended_attention_mask) * -10_000.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]
													__lowerCAmelCase     =  self.get_head_mask(snake_case_    ,				self.config.num_hidden_layers				)

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

													__lowerCAmelCase     =  (
													    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       lowerCAmelCase_		(				A__ ):



							'''simple docstring'''





							def __init__( self    ,				snake_case_    ,				snake_case_				)       ->  Optional[int]:
													__lowerCAmelCase     =  message
													__lowerCAmelCase     =  exit_layer  # start from 1!
class       lowerCAmelCase_		(				nn.Module ):



							'''simple docstring'''





							def __init__( self    ,				snake_case_				)       ->  List[Any]:
													super().__init__()
													__lowerCAmelCase     =  BertPooler(snake_case_				)
													__lowerCAmelCase     =  nn.Dropout(config.hidden_dropout_prob				)
													__lowerCAmelCase     =  nn.Linear(config.hidden_size    ,				config.num_labels				)




							def    A__						( self    ,				snake_case_				)       ->  Optional[int]:
													# Pooler
													__lowerCAmelCase     =  encoder_outputs[0]
													__lowerCAmelCase     =  self.pooler(snake_case_				)
													# "return" pooler_output

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

													# Dropout and classification
													__lowerCAmelCase     =  bmodel_output[1]

													__lowerCAmelCase     =  self.dropout(snake_case_				)
													__lowerCAmelCase     =  self.classifier(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. '''				,				A__				,				)
class       lowerCAmelCase_		(				A__ ):



							'''simple docstring'''





							def __init__( self    ,				snake_case_				)       ->  Optional[Any]:
													super().__init__(snake_case_				)
													__lowerCAmelCase     =  config.num_labels
													__lowerCAmelCase     =  config.num_hidden_layers

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

													self.init_weights()




							@add_start_docstrings_to_model_forward(snake_case_				)
							def    A__						( self    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=None    ,				snake_case_=-1    ,				snake_case_=False    ,				)       ->  Optional[int]:
													__lowerCAmelCase     =  self.num_layers
													try:
																			__lowerCAmelCase     =  self.bert(
																			    snake_case_    ,				attention_mask=snake_case_    ,				token_type_ids=snake_case_    ,				position_ids=snake_case_    ,				head_mask=snake_case_    ,				inputs_embeds=snake_case_    ,				)
																			# sequence_output, pooled_output, (hidden_states), (attentions), highway exits

																			__lowerCAmelCase     =  outputs[1]

																			__lowerCAmelCase     =  self.dropout(snake_case_				)
																			__lowerCAmelCase     =  self.classifier(snake_case_				)
																			__lowerCAmelCase     =  (logits,) + outputs[2:]  # add hidden states and attention if they are here
													except HighwayException as e:
																			__lowerCAmelCase     =  e.message
																			__lowerCAmelCase     =  e.exit_layer
																			__lowerCAmelCase     =  outputs[0]

													if not self.training:
																			__lowerCAmelCase     =  entropy(snake_case_				)
																			__lowerCAmelCase     =  []
																			__lowerCAmelCase     =  []
													if labels is not None:
																			if self.num_labels == 1:
																									#  We are doing regression
																									__lowerCAmelCase     =  MSELoss()
																									__lowerCAmelCase     =  loss_fct(logits.view(-1				)    ,				labels.view(-1				)				)
																			else:
																									__lowerCAmelCase     =  CrossEntropyLoss()
																									__lowerCAmelCase     =  loss_fct(logits.view(-1    ,				self.num_labels				)    ,				labels.view(-1				)				)

																			# work with highway exits
																			__lowerCAmelCase     =  []
																			for highway_exit in outputs[-1]:
																									__lowerCAmelCase     =  highway_exit[0]
																									if not self.training:
																															highway_logits_all.append(snake_case_				)
																															highway_entropy.append(highway_exit[2]				)
																									if self.num_labels == 1:
																															#  We are doing regression
																															__lowerCAmelCase     =  MSELoss()
																															__lowerCAmelCase     =  loss_fct(highway_logits.view(-1				)    ,				labels.view(-1				)				)
																									else:
																															__lowerCAmelCase     =  CrossEntropyLoss()
																															__lowerCAmelCase     =  loss_fct(highway_logits.view(-1    ,				self.num_labels				)    ,				labels.view(-1				)				)
																									highway_losses.append(snake_case_				)

																			if train_highway:
																									__lowerCAmelCase     =  (sum(highway_losses[:-1]				),) + outputs
																									# exclude the final highway, of course
																			else:
																									__lowerCAmelCase     =  (loss,) + outputs
													if not self.training:
																			__lowerCAmelCase     =  outputs + ((original_entropy, highway_entropy), exit_layer)
																			if output_layer >= 0:
																									__lowerCAmelCase     =  (
																									    (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)
