
import argparse
import os
import sys
from unittest.mock import patch

import pytorch_lightning as pl
import timeout_decorator
import torch
from distillation import SummarizationDistiller, distill_main
from finetune import SummarizationModule, main

from transformers import MarianMTModel
from transformers.file_utils import cached_path
from transformers.testing_utils import TestCasePlus, require_torch_gpu, slow
from utils import load_json


UpperCamelCase_						=						"sshleifer/mar_enro_6_3_student"

class      _SCREAMING_SNAKE_CASE				(		snake_case__				):





				def 							_UpperCAmelCase    (		self			:		List[Any]				):


								"""simple docstring"""

								super().setUp()

								A  :	int				=					cached_path(
								    '''https://cdn-datasets.huggingface.co/translation/wmt_en_ro-tr40k-va0.5k-te0.5k.tar.gz'''       ,   extract_compressed_file=__lowercase       ,   )
								A  :	Dict				=					f"""{data_cached}/wmt_en_ro-tr40k-va0.5k-te0.5k"""





				@slow
				@require_torch_gpu
				def 							_UpperCAmelCase    (		self			:		int				):


								"""simple docstring"""

								MarianMTModel.from_pretrained(__lowercase				)






				@slow
				@require_torch_gpu
				def 							_UpperCAmelCase    (		self			:		Union[str, Any]				):


								"""simple docstring"""

								A  :	Tuple				=					{
								    '''$MAX_LEN''': 64,
								    '''$BS''': 64,
								    '''$GAS''': 1,
								    '''$ENRO_DIR''': self.data_dir,
								    '''facebook/mbart-large-cc25''': MARIAN_MODEL,
								    # "val_check_interval=0.25": "val_check_interval=1.0",
								    '''--learning_rate=3e-5''': '''--learning_rate 3e-4''',
								    '''--num_train_epochs 6''': '''--num_train_epochs 1''',
								}

								# Clean up bash script
								A  :	str				=					(self.test_file_dir / '''train_mbart_cc25_enro.sh''').open().read().split('''finetune.py'''				)[1].strip()
								A  :	Union[str, Any]				=					bash_script.replace('''\\\n'''       ,   ''''''				).strip().replace('''"$@"'''       ,   ''''''				)
								for k, v in env_vars_to_replace.items():
												A  :	Tuple				=					bash_script.replace(__lowercase       ,   str(__lowercase				)				)
								A  :	Union[str, Any]				=					self.get_auto_remove_tmp_dir()

								# bash_script = bash_script.replace("--fp16 ", "")
								A  :	List[Any]				=					f"""\n            --output_dir {output_dir}\n            --tokenizer_name Helsinki-NLP/opus-mt-en-ro\n            --sortish_sampler\n            --do_predict\n            --gpus 1\n            --freeze_encoder\n            --n_train 40000\n            --n_val 500\n            --n_test 500\n            --fp16_opt_level O1\n            --num_sanity_val_steps 0\n            --eval_beams 2\n        """.split()
								# XXX: args.gpus > 1 : handle multi_gpu in the future

								A  :	str				=					['''finetune.py'''] + bash_script.split() + args
								with patch.object(__lowercase       ,   '''argv'''       ,   __lowercase				):
												A  :	str				=					argparse.ArgumentParser()
												A  :	Dict				=					pl.Trainer.add_argparse_args(__lowercase				)
												A  :	List[Any]				=					SummarizationModule.add_model_specific_args(__lowercase       ,   os.getcwd()				)
												A  :	int				=					parser.parse_args()
												A  :	Tuple				=					main(__lowercase				)

								# Check metrics
								A  :	Union[str, Any]				=					load_json(model.metrics_save_path				)
								A  :	Any				=					metrics['''val'''][0]
								A  :	int				=					metrics['''val'''][-1]
								self.assertEqual(len(metrics['''val''']				)       ,   (args.max_epochs / args.val_check_interval)				)
								assert isinstance(last_step_stats[f"""val_avg_{model.val_metric}"""]       ,   __lowercase				)

								self.assertGreater(last_step_stats['''val_avg_gen_time''']       ,   0.01				)
								# model hanging on generate. Maybe bad config was saved. (XXX: old comment/assert?)
								self.assertLessEqual(last_step_stats['''val_avg_gen_time''']       ,   1.0				)

								# test learning requirements:

								# 1. BLEU improves over the course of training by more than 2 pts
								self.assertGreater(last_step_stats['''val_avg_bleu'''] - first_step_stats['''val_avg_bleu''']       ,   2				)

								# 2. BLEU finishes above 17
								self.assertGreater(last_step_stats['''val_avg_bleu''']       ,   17				)

								# 3. test BLEU and val BLEU within ~1.1 pt.
								self.assertLess(abs(metrics['''val'''][-1]['''val_avg_bleu'''] - metrics['''test'''][-1]['''test_avg_bleu''']				)       ,   1.1				)

								# check lightning ckpt can be loaded and has a reasonable statedict
								A  :	Any				=					os.listdir(__lowercase				)
								A  :	Tuple				=					[x for x in contents if x.endswith('''.ckpt'''				)][0]
								A  :	int				=					os.path.join(args.output_dir       ,   __lowercase				)
								A  :	List[str]				=					torch.load(__lowercase       ,   map_location='''cpu'''				)
								A  :	str				=					'''model.model.decoder.layers.0.encoder_attn_layer_norm.weight'''
								assert expected_key in ckpt["state_dict"]
								assert ckpt["state_dict"]["model.model.decoder.layers.0.encoder_attn_layer_norm.weight"].dtype == torch.floataa

								# TODO: turn on args.do_predict when PL bug fixed.
								if args.do_predict:
												A  :	str				=					{os.path.basename(__lowercase				) for p in contents}
												assert "test_generations.txt" in contents
												assert "test_results.txt" in contents
												# assert len(metrics["val"]) ==  desired_n_evals
												assert len(metrics['''test''']				) == 1



class      _SCREAMING_SNAKE_CASE				(		snake_case__				):





				@timeout_decorator.timeout(600				)
				@slow
				@require_torch_gpu
				def 							_UpperCAmelCase    (		self			:		str				):


								"""simple docstring"""

								A  :	int				=					f"""{self.test_file_dir_str}/test_data/wmt_en_ro"""
								A  :	List[Any]				=					{
								    '''--fp16_opt_level=O1''': '''''',
								    '''$MAX_LEN''': 128,
								    '''$BS''': 16,
								    '''$GAS''': 1,
								    '''$ENRO_DIR''': data_dir,
								    '''$m''': '''sshleifer/student_marian_en_ro_6_1''',
								    '''val_check_interval=0.25''': '''val_check_interval=1.0''',
								}

								# Clean up bash script
								A  :	Any				=					(
								    (self.test_file_dir / '''distil_marian_no_teacher.sh''').open().read().split('''distillation.py'''				)[1].strip()
								)
								A  :	Optional[Any]				=					bash_script.replace('''\\\n'''       ,   ''''''				).strip().replace('''"$@"'''       ,   ''''''				)
								A  :	Union[str, Any]				=					bash_script.replace('''--fp16 '''       ,   ''' '''				)

								for k, v in env_vars_to_replace.items():
												A  :	str				=					bash_script.replace(__lowercase       ,   str(__lowercase				)				)
								A  :	Any				=					self.get_auto_remove_tmp_dir()
								A  :	Optional[int]				=					bash_script.replace('''--fp16'''       ,   ''''''				)
								A  :	Tuple				=					6
								A  :	str				=					(
								    ['''distillation.py''']
								    + bash_script.split()
								    + [
								        f"""--output_dir={output_dir}""",
								        '''--gpus=1''',
								        '''--learning_rate=1e-3''',
								        f"""--num_train_epochs={epochs}""",
								        '''--warmup_steps=10''',
								        '''--val_check_interval=1.0''',
								        '''--do_predict''',
								    ]
								)
								with patch.object(__lowercase       ,   '''argv'''       ,   __lowercase				):
												A  :	int				=					argparse.ArgumentParser()
												A  :	int				=					pl.Trainer.add_argparse_args(__lowercase				)
												A  :	Optional[int]				=					SummarizationDistiller.add_model_specific_args(__lowercase       ,   os.getcwd()				)
												A  :	Union[str, Any]				=					parser.parse_args()
												# assert args.gpus == gpus THIS BREAKS for multi_gpu

												A  :	Union[str, Any]				=					distill_main(__lowercase				)

								# Check metrics
								A  :	Dict				=					load_json(model.metrics_save_path				)
								A  :	List[str]				=					metrics['''val'''][0]
								A  :	Tuple				=					metrics['''val'''][-1]
								assert len(metrics['''val''']				) >= (args.max_epochs / args.val_check_interval)  # +1 accounts for val_sanity_check

								assert last_step_stats["val_avg_gen_time"] >= 0.01

								assert first_step_stats["val_avg_bleu"] < last_step_stats["val_avg_bleu"]  # model learned nothing
								assert 1.0 >= last_step_stats["val_avg_gen_time"]  # model hanging on generate. Maybe bad config was saved.
								assert isinstance(last_step_stats[f"""val_avg_{model.val_metric}"""]       ,   __lowercase				)

								# check lightning ckpt can be loaded and has a reasonable statedict
								A  :	int				=					os.listdir(__lowercase				)
								A  :	str				=					[x for x in contents if x.endswith('''.ckpt'''				)][0]
								A  :	Union[str, Any]				=					os.path.join(args.output_dir       ,   __lowercase				)
								A  :	Optional[int]				=					torch.load(__lowercase       ,   map_location='''cpu'''				)
								A  :	int				=					'''model.model.decoder.layers.0.encoder_attn_layer_norm.weight'''
								assert expected_key in ckpt["state_dict"]
								assert ckpt["state_dict"]["model.model.decoder.layers.0.encoder_attn_layer_norm.weight"].dtype == torch.floataa

								# TODO: turn on args.do_predict when PL bug fixed.
								if args.do_predict:
												A  :	int				=					{os.path.basename(__lowercase				) for p in contents}
												assert "test_generations.txt" in contents
												assert "test_results.txt" in contents
												# assert len(metrics["val"]) ==  desired_n_evals
												assert len(metrics['''test''']				) == 1






"""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)




