





"""simple docstring"""





from typing import Dict, List

from nltk.translate import gleu_score

import datasets
from datasets import MetricInfo


_UpperCamelCase					      =   '\\n@misc{wu2016googles,\n      title={Google\'s Neural Machine Translation System: Bridging the Gap between Human and Machine Translation},\n      author={Yonghui Wu and Mike Schuster and Zhifeng Chen and Quoc V. Le and Mohammad Norouzi and Wolfgang Macherey\n              and Maxim Krikun and Yuan Cao and Qin Gao and Klaus Macherey and Jeff Klingner and Apurva Shah and Melvin\n              Johnson and Xiaobing Liu and Łukasz Kaiser and Stephan Gouws and Yoshikiyo Kato and Taku Kudo and Hideto\n              Kazawa and Keith Stevens and George Kurian and Nishant Patil and Wei Wang and Cliff Young and\n              Jason Smith and Jason Riesa and Alex Rudnick and Oriol Vinyals and Greg Corrado and Macduff Hughes\n              and Jeffrey Dean},\n      year={2016},\n      eprint={1609.08144},\n      archivePrefix={arXiv},\n      primaryClass={cs.CL}\n}\n'

_UpperCamelCase					      =   '\\nThe BLEU score has some undesirable properties when used for single\nsentences, as it was designed to be a corpus measure. We therefore\nuse a slightly different score for our RL experiments which we call\nthe \'GLEU score\'. For the GLEU score, we record all sub-sequences of\n1, 2, 3 or 4 tokens in output and target sequence (n-grams). We then\ncompute a recall, which is the ratio of the number of matching n-grams\nto the number of total n-grams in the target (ground truth) sequence,\nand a precision, which is the ratio of the number of matching n-grams\nto the number of total n-grams in the generated output sequence. Then\nGLEU score is simply the minimum of recall and precision. This GLEU\nscore\'s range is always between 0 (no matches) and 1 (all match) and\nit is symmetrical when switching output and target. According to\nour experiments, GLEU score correlates quite well with the BLEU\nmetric on a corpus level but does not have its drawbacks for our per\nsentence reward objective.\n'

_UpperCamelCase					      =   '\\nComputes corpus-level Google BLEU (GLEU) score of translated segments against one or more references.\nInstead of averaging the sentence level GLEU scores (i.e. macro-average precision), Wu et al. (2016) sum up the matching\ntokens and the max of hypothesis and reference tokens for each sentence, then compute using the aggregate values.\n\nArgs:\n    predictions (list of str): list of translations to score.\n        Each translation should be tokenized into a list of tokens.\n    references (list of list of str): list of lists of references for each translation.\n        Each reference should be tokenized into a list of tokens.\n    min_len (int): The minimum order of n-gram this function should extract. Defaults to 1.\n    max_len (int): The maximum order of n-gram this function should extract. Defaults to 4.\n\nReturns:\n    \'google_bleu\': google_bleu score\n\nExamples:\n    Example 1:\n        >>> hyp1 = [\'It\', \'is\', \'a\', \'guide\', \'to\', \'action\', \'which\',\n        ...         \'ensures\', \'that\', \'the\', \'rubber\', \'duck\', \'always\',\n        ...         \'disobeys\', \'the\', \'commands\', \'of\', \'the\', \'cat\']\n        >>> ref1a = [\'It\', \'is\', \'the\', \'guiding\', \'principle\', \'which\',\n        ...          \'guarantees\', \'the\', \'rubber\', \'duck\', \'forces\', \'never\',\n        ...          \'being\', \'under\', \'the\', \'command\', \'of\', \'the\', \'cat\']\n\n        >>> hyp2 = [\'he\', \'read\', \'the\', \'book\', \'because\', \'he\', \'was\',\n        ...         \'interested\', \'in\', \'world\', \'history\']\n        >>> ref2a = [\'he\', \'was\', \'interested\', \'in\', \'world\', \'history\',\n        ...          \'because\', \'he\', \'read\', \'the\', \'book\']\n\n        >>> list_of_references = [[ref1a], [ref2a]]\n        >>> hypotheses = [hyp1, hyp2]\n        >>> google_bleu = datasets.load_metric("google_bleu")\n        >>> results = google_bleu.compute(predictions=hypotheses, references=list_of_references)\n        >>> print(round(results["google_bleu"], 2))\n        0.44\n\n    Example 2:\n        >>> hyp1 = [\'It\', \'is\', \'a\', \'guide\', \'to\', \'action\', \'which\',\n        ...         \'ensures\', \'that\', \'the\', \'rubber\', \'duck\', \'always\',\n        ...         \'disobeys\', \'the\', \'commands\', \'of\', \'the\', \'cat\']\n        >>> ref1a = [\'It\', \'is\', \'the\', \'guiding\', \'principle\', \'which\',\n        ...          \'guarantees\', \'the\', \'rubber\', \'duck\', \'forces\', \'never\',\n        ...          \'being\', \'under\', \'the\', \'command\', \'of\', \'the\', \'cat\']\n        >>> ref1b = [\'It\', \'is\', \'a\', \'guide\', \'to\', \'action\', \'that\',\n        ...          \'ensures\', \'that\', \'the\', \'rubber\', \'duck\', \'will\', \'never\',\n        ...          \'heed\', \'the\', \'cat\', \'commands\']\n        >>> ref1c = [\'It\', \'is\', \'the\', \'practical\', \'guide\', \'for\', \'the\',\n        ...          \'rubber\', \'duck\', \'army\', \'never\', \'to\', \'heed\', \'the\', \'directions\',\n        ...          \'of\', \'the\', \'cat\']\n\n        >>> hyp2 = [\'he\', \'read\', \'the\', \'book\', \'because\', \'he\', \'was\',\n        ...         \'interested\', \'in\', \'world\', \'history\']\n        >>> ref2a = [\'he\', \'was\', \'interested\', \'in\', \'world\', \'history\',\n        ...          \'because\', \'he\', \'read\', \'the\', \'book\']\n\n        >>> list_of_references = [[ref1a, ref1b, ref1c], [ref2a]]\n        >>> hypotheses = [hyp1, hyp2]\n        >>> google_bleu = datasets.load_metric("google_bleu")\n        >>> results = google_bleu.compute(predictions=hypotheses, references=list_of_references)\n        >>> print(round(results["google_bleu"], 2))\n        0.61\n\n    Example 3:\n        >>> hyp1 = [\'It\', \'is\', \'a\', \'guide\', \'to\', \'action\', \'which\',\n        ...         \'ensures\', \'that\', \'the\', \'rubber\', \'duck\', \'always\',\n        ...         \'disobeys\', \'the\', \'commands\', \'of\', \'the\', \'cat\']\n        >>> ref1a = [\'It\', \'is\', \'the\', \'guiding\', \'principle\', \'which\',\n        ...          \'guarantees\', \'the\', \'rubber\', \'duck\', \'forces\', \'never\',\n        ...          \'being\', \'under\', \'the\', \'command\', \'of\', \'the\', \'cat\']\n        >>> ref1b = [\'It\', \'is\', \'a\', \'guide\', \'to\', \'action\', \'that\',\n        ...          \'ensures\', \'that\', \'the\', \'rubber\', \'duck\', \'will\', \'never\',\n        ...          \'heed\', \'the\', \'cat\', \'commands\']\n        >>> ref1c = [\'It\', \'is\', \'the\', \'practical\', \'guide\', \'for\', \'the\',\n        ...          \'rubber\', \'duck\', \'army\', \'never\', \'to\', \'heed\', \'the\', \'directions\',\n        ...          \'of\', \'the\', \'cat\']\n\n        >>> hyp2 = [\'he\', \'read\', \'the\', \'book\', \'because\', \'he\', \'was\',\n        ...         \'interested\', \'in\', \'world\', \'history\']\n        >>> ref2a = [\'he\', \'was\', \'interested\', \'in\', \'world\', \'history\',\n        ...          \'because\', \'he\', \'read\', \'the\', \'book\']\n\n        >>> list_of_references = [[ref1a, ref1b, ref1c], [ref2a]]\n        >>> hypotheses = [hyp1, hyp2]\n        >>> google_bleu = datasets.load_metric("google_bleu")\n        >>> results = google_bleu.compute(predictions=hypotheses, references=list_of_references, min_len=2)\n        >>> print(round(results["google_bleu"], 2))\n        0.53\n\n    Example 4:\n        >>> hyp1 = [\'It\', \'is\', \'a\', \'guide\', \'to\', \'action\', \'which\',\n        ...         \'ensures\', \'that\', \'the\', \'rubber\', \'duck\', \'always\',\n        ...         \'disobeys\', \'the\', \'commands\', \'of\', \'the\', \'cat\']\n        >>> ref1a = [\'It\', \'is\', \'the\', \'guiding\', \'principle\', \'which\',\n        ...          \'guarantees\', \'the\', \'rubber\', \'duck\', \'forces\', \'never\',\n        ...          \'being\', \'under\', \'the\', \'command\', \'of\', \'the\', \'cat\']\n        >>> ref1b = [\'It\', \'is\', \'a\', \'guide\', \'to\', \'action\', \'that\',\n        ...          \'ensures\', \'that\', \'the\', \'rubber\', \'duck\', \'will\', \'never\',\n        ...          \'heed\', \'the\', \'cat\', \'commands\']\n        >>> ref1c = [\'It\', \'is\', \'the\', \'practical\', \'guide\', \'for\', \'the\',\n        ...          \'rubber\', \'duck\', \'army\', \'never\', \'to\', \'heed\', \'the\', \'directions\',\n        ...          \'of\', \'the\', \'cat\']\n\n        >>> hyp2 = [\'he\', \'read\', \'the\', \'book\', \'because\', \'he\', \'was\',\n        ...         \'interested\', \'in\', \'world\', \'history\']\n        >>> ref2a = [\'he\', \'was\', \'interested\', \'in\', \'world\', \'history\',\n        ...          \'because\', \'he\', \'read\', \'the\', \'book\']\n\n        >>> list_of_references = [[ref1a, ref1b, ref1c], [ref2a]]\n        >>> hypotheses = [hyp1, hyp2]\n        >>> google_bleu = datasets.load_metric("google_bleu")\n        >>> results = google_bleu.compute(predictions=hypotheses,references=list_of_references, min_len=2, max_len=6)\n        >>> print(round(results["google_bleu"], 2))\n        0.4\n'
@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION							,	_KWARGS_DESCRIPTION )
class 		SCREAMING_SNAKE_CASE_				(     datasets.Metric ):
				"""simple docstring"""




				def        __lowercase	(  self	:List[Any]      ):
											return datasets.MetricInfo(
											    description=_DESCRIPTION    ,     citation=_CITATION    ,     inputs_description=_KWARGS_DESCRIPTION    ,     features=datasets.Features(
											        {
											            '''predictions''': datasets.Sequence(datasets.Value('''string'''    ,     id='''token'''      )    ,     id='''sequence'''      ),
											            '''references''': datasets.Sequence(
											                datasets.Sequence(datasets.Value('''string'''    ,     id='''token'''      )    ,     id='''sequence'''      )    ,     id='''references'''      ),
											        }      )    ,     )




				def        __lowercase	(  self	:Dict    ,     __lowercase	:List[List[List[str]]]    ,     __lowercase	:List[List[str]]    ,     __lowercase	:int = 1    ,     __lowercase	:int = 4    ,     ):
											return {
											    "google_bleu": gleu_score.corpus_gleu(
											        list_of_references=__lowercase    ,     hypotheses=__lowercase    ,     min_len=__lowercase    ,     max_len=__lowercase      )
											}











"""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 					lowerCAmelCase_      (   SCREAMING_SNAKE_CASE							:				Any ):


							'''simple docstring'''







							__lowerCamelCase   :   Tuple       =torch.exp(SCREAMING_SNAKE_CASE )
							__lowerCamelCase   :   Dict       =torch.sum(SCREAMING_SNAKE_CASE     ,    dim=1 )  # sum of exp(x_i)
							__lowerCamelCase   :   Optional[int]       =torch.sum(x * exp_x     ,    dim=1 )  # sum of x_i * exp(x_i)
							return torch.log(SCREAMING_SNAKE_CASE ) - B / A
class 		SCREAMING_SNAKE_CASE_				(     nn.Module ):
				"""simple docstring"""




				def __init__(  self	:List[str]    ,     __lowercase	:int      ):
											super().__init__()
											__lowerCamelCase   :   str       =config.output_attentions
											__lowerCamelCase   :   List[Any]       =config.output_hidden_states
											__lowerCamelCase   :   Dict       =nn.ModuleList([BertLayer(__lowercase      ) for _ in range(config.num_hidden_layers      )]      )
											__lowerCamelCase   :   str       =nn.ModuleList([BertHighway(__lowercase      ) for _ in range(config.num_hidden_layers      )]      )

											__lowerCamelCase   :   Optional[Any]       =[-1 for _ in range(config.num_hidden_layers      )]




				def        __lowercase	(  self	:Union[str, Any]    ,     __lowercase	:Union[str, Any]      ):
											if (type(__lowercase      ) is float) or (type(__lowercase      ) is int):
																		for i in range(len(self.early_exit_entropy      )      ):
																									__lowerCamelCase   :   Tuple       =x
											else:
																		__lowerCamelCase   :   Any       =x




				def        __lowercase	(  self	:Union[str, Any]    ,     __lowercase	:Tuple      ):
											__lowerCamelCase   :   Union[str, Any]       =pooler.state_dict()
											for highway in self.highway:
																		for name, param in highway.pooler.state_dict().items():
																									param.copy_(loaded_model[name]      )




				def        __lowercase	(  self	:Tuple    ,     __lowercase	:Optional[int]    ,     __lowercase	:Dict=None    ,     __lowercase	:Union[str, Any]=None    ,     __lowercase	:List[str]=None    ,     __lowercase	:str=None    ,     ):
											__lowerCamelCase   :   Any       =()
											__lowerCamelCase   :   List[str]       =()
											__lowerCamelCase   :   Optional[int]       =()
											for i, layer_module in enumerate(self.layer      ):
																		if self.output_hidden_states:
																									__lowerCamelCase   :   int       =all_hidden_states + (hidden_states,)

																		__lowerCamelCase   :   List[Any]       =layer_module(
																		    __lowercase    ,     __lowercase    ,     head_mask[i]    ,     __lowercase    ,     __lowercase      )
																		__lowerCamelCase   :   Optional[int]       =layer_outputs[0]

																		if self.output_attentions:
																									__lowerCamelCase   :   Optional[Any]       =all_attentions + (layer_outputs[1],)

																		__lowerCamelCase   :   Any       =(hidden_states,)
																		if self.output_hidden_states:
																									__lowerCamelCase   :   Optional[Any]       =current_outputs + (all_hidden_states,)
																		if self.output_attentions:
																									__lowerCamelCase   :   Dict       =current_outputs + (all_attentions,)

																		__lowerCamelCase   :   str       =self.highway[i](__lowercase      )
																		# logits, pooled_output

																		if not self.training:
																									__lowerCamelCase   :   Tuple       =highway_exit[0]
																									__lowerCamelCase   :   Tuple       =entropy(__lowercase      )
																									__lowerCamelCase   :   Tuple       =highway_exit + (highway_entropy,)  # logits, hidden_states(?), entropy
																									__lowerCamelCase   :   Optional[int]       =all_highway_exits + (highway_exit,)

																									if highway_entropy < self.early_exit_entropy[i]:
																																__lowerCamelCase   :   Dict       =(highway_logits,) + current_outputs[1:] + (all_highway_exits,)
																																raise HighwayException(__lowercase    ,     i + 1      )
																		else:
																									__lowerCamelCase   :   Union[str, Any]       =all_highway_exits + (highway_exit,)

        # Add last layer
											if self.output_hidden_states:
																		__lowerCamelCase   :   Optional[Any]       =all_hidden_states + (hidden_states,)

											__lowerCamelCase   :   List[Any]       =(hidden_states,)
											if self.output_hidden_states:
																		__lowerCamelCase   :   Tuple       =outputs + (all_hidden_states,)
											if self.output_attentions:
																		__lowerCamelCase   :   Optional[int]       =outputs + (all_attentions,)

											__lowerCamelCase   :   int       =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 		SCREAMING_SNAKE_CASE_				(     snake_case__ ):
				"""simple docstring"""




				def __init__(  self	:Union[str, Any]    ,     __lowercase	:str      ):
											super().__init__(__lowercase      )
											__lowerCamelCase   :   Union[str, Any]       =config

											__lowerCamelCase   :   List[str]       =BertEmbeddings(__lowercase      )
											__lowerCamelCase   :   Dict       =DeeBertEncoder(__lowercase      )
											__lowerCamelCase   :   List[Any]       =BertPooler(__lowercase      )

											self.init_weights()




				def        __lowercase	(  self	:Tuple      ):
											self.encoder.init_highway_pooler(self.pooler      )




				def        __lowercase	(  self	:Dict      ):
											return self.embeddings.word_embeddings




				def        __lowercase	(  self	:List[str]    ,     __lowercase	:int      ):
											__lowerCamelCase   :   Union[str, Any]       =value




				def        __lowercase	(  self	:List[Any]    ,     __lowercase	:Dict      ):
											for layer, heads in heads_to_prune.items():
																		self.encoder.layer[layer].attention.prune_heads(__lowercase      )




				@add_start_docstrings_to_model_forward(__lowercase      )
				def        __lowercase	(  self	:Optional[Any]    ,     __lowercase	:List[str]=None    ,     __lowercase	:List[Any]=None    ,     __lowercase	:Any=None    ,     __lowercase	:Tuple=None    ,     __lowercase	:Union[str, Any]=None    ,     __lowercase	:Optional[Any]=None    ,     __lowercase	:Union[str, Any]=None    ,     __lowercase	:Tuple=None    ,     ):
											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   :   List[str]       =input_ids.size()
											elif inputs_embeds is not None:
																		__lowerCamelCase   :   str       =inputs_embeds.size()[:-1]
											else:
																		raise ValueError('''You have to specify either input_ids or inputs_embeds'''      )

											__lowerCamelCase   :   Optional[int]       =input_ids.device if input_ids is not None else inputs_embeds.device

											if attention_mask is None:
																		__lowerCamelCase   :   str       =torch.ones(__lowercase    ,     device=__lowercase      )
											if encoder_attention_mask is None:
																		__lowerCamelCase   :   Tuple       =torch.ones(__lowercase    ,     device=__lowercase      )
											if token_type_ids is None:
																		__lowerCamelCase   :   List[Any]       =torch.zeros(__lowercase    ,     dtype=torch.long    ,     device=__lowercase      )

											# 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   :   torch.Tensor       =self.get_extended_attention_mask(__lowercase    ,     __lowercase    ,     __lowercase      )

											# 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   :   List[str]       =encoder_attention_mask[:, None, :, :]
											if encoder_attention_mask.dim() == 2:
																		__lowerCamelCase   :   Any       =encoder_attention_mask[:, None, None, :]

											__lowerCamelCase   :   Optional[Any]       =encoder_extended_attention_mask.to(
											    dtype=next(self.parameters()      ).dtype      )  # fp16 compatibility
											__lowerCamelCase   :   List[str]       =(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]
											__lowerCamelCase   :   Union[str, Any]       =self.get_head_mask(__lowercase    ,     self.config.num_hidden_layers      )

											__lowerCamelCase   :   str       =self.embeddings(
											    input_ids=__lowercase    ,     position_ids=__lowercase    ,     token_type_ids=__lowercase    ,     inputs_embeds=__lowercase      )
											__lowerCamelCase   :   Dict       =self.encoder(
											    __lowercase    ,     attention_mask=__lowercase    ,     head_mask=__lowercase    ,     encoder_hidden_states=__lowercase    ,     encoder_attention_mask=__lowercase    ,     )
											__lowerCamelCase   :   int       =encoder_outputs[0]
											__lowerCamelCase   :   Tuple       =self.pooler(__lowercase      )

											__lowerCamelCase   :   int       =(
											    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 		SCREAMING_SNAKE_CASE_				(     snake_case__ ):
				"""simple docstring"""




				def __init__(  self	:List[Any]    ,     __lowercase	:Optional[Any]    ,     __lowercase	:Dict      ):
											__lowerCamelCase   :   List[Any]       =message
											__lowerCamelCase   :   int       =exit_layer  # start from 1!

class 		SCREAMING_SNAKE_CASE_				(     nn.Module ):
				"""simple docstring"""




				def __init__(  self	:Any    ,     __lowercase	:str      ):
											super().__init__()
											__lowerCamelCase   :   str       =BertPooler(__lowercase      )
											__lowerCamelCase   :   Union[str, Any]       =nn.Dropout(config.hidden_dropout_prob      )
											__lowerCamelCase   :   List[str]       =nn.Linear(config.hidden_size    ,     config.num_labels      )




				def        __lowercase	(  self	:Union[str, Any]    ,     __lowercase	:List[str]      ):
											# Pooler
											__lowerCamelCase   :   Optional[Any]       =encoder_outputs[0]
											__lowerCamelCase   :   Any       =self.pooler(__lowercase      )
											# "return" pooler_output

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

											# Dropout and classification
											__lowerCamelCase   :   List[Any]       =bmodel_output[1]

											__lowerCamelCase   :   Optional[Any]       =self.dropout(__lowercase      )
											__lowerCamelCase   :   int       =self.classifier(__lowercase      )

											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 		SCREAMING_SNAKE_CASE_				(     snake_case__ ):
				"""simple docstring"""




				def __init__(  self	:Union[str, Any]    ,     __lowercase	:Dict      ):
											super().__init__(__lowercase      )
											__lowerCamelCase   :   Any       =config.num_labels
											__lowerCamelCase   :   int       =config.num_hidden_layers

											__lowerCamelCase   :   Tuple       =DeeBertModel(__lowercase      )
											__lowerCamelCase   :   Optional[int]       =nn.Dropout(config.hidden_dropout_prob      )
											__lowerCamelCase   :   Optional[int]       =nn.Linear(config.hidden_size    ,     self.config.num_labels      )

											self.init_weights()




				@add_start_docstrings_to_model_forward(__lowercase      )
				def        __lowercase	(  self	:List[str]    ,     __lowercase	:List[str]=None    ,     __lowercase	:str=None    ,     __lowercase	:Optional[Any]=None    ,     __lowercase	:List[Any]=None    ,     __lowercase	:Union[str, Any]=None    ,     __lowercase	:Dict=None    ,     __lowercase	:int=None    ,     __lowercase	:int=-1    ,     __lowercase	:List[str]=False    ,     ):
											__lowerCamelCase   :   Union[str, Any]       =self.num_layers
											try:
																		__lowerCamelCase   :   Union[str, Any]       =self.bert(
																		    __lowercase    ,     attention_mask=__lowercase    ,     token_type_ids=__lowercase    ,     position_ids=__lowercase    ,     head_mask=__lowercase    ,     inputs_embeds=__lowercase    ,     )
																		# sequence_output, pooled_output, (hidden_states), (attentions), highway exits

																		__lowerCamelCase   :   List[Any]       =outputs[1]

																		__lowerCamelCase   :   Optional[Any]       =self.dropout(__lowercase      )
																		__lowerCamelCase   :   Tuple       =self.classifier(__lowercase      )
																		__lowerCamelCase   :   int       =(logits,) + outputs[2:]  # add hidden states and attention if they are here
											except HighwayException as e:
																		__lowerCamelCase   :   Union[str, Any]       =e.message
																		__lowerCamelCase   :   Optional[Any]       =e.exit_layer
																		__lowerCamelCase   :   Any       =outputs[0]

											if not self.training:
																		__lowerCamelCase   :   List[Any]       =entropy(__lowercase      )
																		__lowerCamelCase   :   Union[str, Any]       =[]
																		__lowerCamelCase   :   int       =[]
											if labels is not None:
																		if self.num_labels == 1:
																									#  We are doing regression
																									__lowerCamelCase   :   Union[str, Any]       =MSELoss()
																									__lowerCamelCase   :   List[Any]       =loss_fct(logits.view(-1      )    ,     labels.view(-1      )      )
																		else:
																									__lowerCamelCase   :   Dict       =CrossEntropyLoss()
																									__lowerCamelCase   :   List[Any]       =loss_fct(logits.view(-1    ,     self.num_labels      )    ,     labels.view(-1      )      )

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

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




