


"""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_  :      Optional[int]  ):
				SCREAMING_SNAKE_CASE__     							=				torch.exp(lowerCamelCase_  )
				SCREAMING_SNAKE_CASE__     							=				torch.sum(lowerCamelCase_						,					dim=1  )  # sum of exp(x_i)
				SCREAMING_SNAKE_CASE__     							=				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       ,		UpperCAmelCase__     ):
								super().__init__()
								SCREAMING_SNAKE_CASE__     							=				config.output_attentions
								SCREAMING_SNAKE_CASE__     							=				config.output_hidden_states
								SCREAMING_SNAKE_CASE__     							=				nn.ModuleList([BertLayer(UpperCAmelCase__     ) for _ in range(config.num_hidden_layers     )]     )
								SCREAMING_SNAKE_CASE__     							=				nn.ModuleList([BertHighway(UpperCAmelCase__     ) for _ in range(config.num_hidden_layers     )]     )

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



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__     ):
								if (type(UpperCAmelCase__     ) is float) or (type(UpperCAmelCase__     ) is int):
												for i in range(len(self.early_exit_entropy     )     ):
																SCREAMING_SNAKE_CASE__     							=				x
								else:
												SCREAMING_SNAKE_CASE__     							=				x



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__     ):
								SCREAMING_SNAKE_CASE__     							=				pooler.state_dict()
								for highway in self.highway:
												for name, param in highway.pooler.state_dict().items():
																param.copy_(loaded_model[name]     )






				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__=None       ,		UpperCAmelCase__=None       ,		UpperCAmelCase__=None       ,		UpperCAmelCase__=None       ,		):
								SCREAMING_SNAKE_CASE__     							=				()
								SCREAMING_SNAKE_CASE__     							=				()
								SCREAMING_SNAKE_CASE__     							=				()
								for i, layer_module in enumerate(self.layer     ):
												if self.output_hidden_states:
																SCREAMING_SNAKE_CASE__     							=				all_hidden_states + (hidden_states,)

												SCREAMING_SNAKE_CASE__     							=				layer_module(
												    UpperCAmelCase__       ,		UpperCAmelCase__       ,		head_mask[i]       ,		UpperCAmelCase__       ,		UpperCAmelCase__     )
												SCREAMING_SNAKE_CASE__     							=				layer_outputs[0]

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

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

												SCREAMING_SNAKE_CASE__     							=				self.highway[i](UpperCAmelCase__     )
												# logits, pooled_output

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

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

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

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

								SCREAMING_SNAKE_CASE__     							=				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). "       ,  lowercase       ,  )
class 			lowerCamelCase_		(			lowercase   ):





				"""simple docstring"""







				def __init__(						self       ,		UpperCAmelCase__     ):
								super().__init__(UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				config

								SCREAMING_SNAKE_CASE__     							=				BertEmbeddings(UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				DeeBertEncoder(UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				BertPooler(UpperCAmelCase__     )

								self.init_weights()



				def   lowerCAmelCase__		(						self     ):
								self.encoder.init_highway_pooler(self.pooler     )



				def   lowerCAmelCase__		(						self     ):
								return self.embeddings.word_embeddings



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__     ):
								SCREAMING_SNAKE_CASE__     							=				value



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






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

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

								if attention_mask is None:
												SCREAMING_SNAKE_CASE__     							=				torch.ones(UpperCAmelCase__       ,		device=UpperCAmelCase__     )
								if encoder_attention_mask is None:
												SCREAMING_SNAKE_CASE__     							=				torch.ones(UpperCAmelCase__       ,		device=UpperCAmelCase__     )
								if token_type_ids is None:
												SCREAMING_SNAKE_CASE__     							=				torch.zeros(UpperCAmelCase__       ,		dtype=torch.long       ,		device=UpperCAmelCase__     )

								# 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.
								SCREAMING_SNAKE_CASE__     							=				self.get_extended_attention_mask(UpperCAmelCase__       ,		UpperCAmelCase__       ,		UpperCAmelCase__     )

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

								SCREAMING_SNAKE_CASE__     							=				encoder_extended_attention_mask.to(
								    dtype=next(self.parameters()     ).dtype     )  # fp16 compatibility
								SCREAMING_SNAKE_CASE__     							=				(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]
								SCREAMING_SNAKE_CASE__     							=				self.get_head_mask(UpperCAmelCase__       ,		self.config.num_hidden_layers     )

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

								SCREAMING_SNAKE_CASE__     							=				(
								    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_		(			lowercase   ):





				"""simple docstring"""







				def __init__(						self       ,		UpperCAmelCase__       ,		UpperCAmelCase__     ):
								SCREAMING_SNAKE_CASE__     							=				message
								SCREAMING_SNAKE_CASE__     							=				exit_layer  # start from 1!


class 			lowerCamelCase_		(			nn.Module   ):





				"""simple docstring"""







				def __init__(						self       ,		UpperCAmelCase__     ):
								super().__init__()
								SCREAMING_SNAKE_CASE__     							=				BertPooler(UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				nn.Dropout(config.hidden_dropout_prob     )
								SCREAMING_SNAKE_CASE__     							=				nn.Linear(config.hidden_size       ,		config.num_labels     )






				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__     ):
								# Pooler
								SCREAMING_SNAKE_CASE__     							=				encoder_outputs[0]
								SCREAMING_SNAKE_CASE__     							=				self.pooler(UpperCAmelCase__     )
								# "return" pooler_output

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

								# Dropout and classification
								SCREAMING_SNAKE_CASE__     							=				bmodel_output[1]

								SCREAMING_SNAKE_CASE__     							=				self.dropout(UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				self.classifier(UpperCAmelCase__     )

								return logits, pooled_output


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





				"""simple docstring"""







				def __init__(						self       ,		UpperCAmelCase__     ):
								super().__init__(UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				config.num_labels
								SCREAMING_SNAKE_CASE__     							=				config.num_hidden_layers

								SCREAMING_SNAKE_CASE__     							=				DeeBertModel(UpperCAmelCase__     )
								SCREAMING_SNAKE_CASE__     							=				nn.Dropout(config.hidden_dropout_prob     )
								SCREAMING_SNAKE_CASE__     							=				nn.Linear(config.hidden_size       ,		self.config.num_labels     )

								self.init_weights()






				@add_start_docstrings_to_model_forward(UpperCAmelCase__     )
				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__=None       ,		UpperCAmelCase__=None       ,		UpperCAmelCase__=None       ,		UpperCAmelCase__=None       ,		UpperCAmelCase__=None       ,		UpperCAmelCase__=None       ,		UpperCAmelCase__=None       ,		UpperCAmelCase__=-1       ,		UpperCAmelCase__=False       ,		):
								SCREAMING_SNAKE_CASE__     							=				self.num_layers
								try:
												SCREAMING_SNAKE_CASE__     							=				self.bert(
												    UpperCAmelCase__       ,		attention_mask=UpperCAmelCase__       ,		token_type_ids=UpperCAmelCase__       ,		position_ids=UpperCAmelCase__       ,		head_mask=UpperCAmelCase__       ,		inputs_embeds=UpperCAmelCase__       ,		)
												# sequence_output, pooled_output, (hidden_states), (attentions), highway exits

												SCREAMING_SNAKE_CASE__     							=				outputs[1]

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

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

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

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




"""simple docstring"""






import json
import os
import subprocess
import unittest
from ast import literal_eval

import pytest
from parameterized import parameterized_class

from . import is_sagemaker_available


if is_sagemaker_available():
	from sagemaker import Session, TrainingJobAnalytics
	from sagemaker.huggingface import HuggingFace


@pytest.mark.skipif(
    literal_eval(os.getenv("TEST_SAGEMAKER"       ,  "False"   )   ) is not True       ,  reason="Skipping test because should only be run when releasing minor transformers version"       ,  )
@pytest.mark.usefixtures("sm_env"   )
@parameterized_class(
    [
        {
            "framework": "pytorch",
            "script": "run_glue.py",
            "model_name_or_path": "distilbert-base-cased",
            "instance_type": "ml.g4dn.xlarge",
            "results": {"train_runtime": 650, "eval_accuracy": 0.6, "eval_loss": 0.9},
        },
        {
            "framework": "tensorflow",
            "script": "run_tf.py",
            "model_name_or_path": "distilbert-base-cased",
            "instance_type": "ml.g4dn.xlarge",
            "results": {"train_runtime": 600, "eval_accuracy": 0.3, "eval_loss": 0.9},
        },
    ]   )
class 			lowerCamelCase_		(			unittest.TestCase   ):





				"""simple docstring"""







				def   lowerCAmelCase__		(						self     ):
								if self.framework == "pytorch":
												subprocess.run(
												    f'''cp ./examples/pytorch/text-classification/run_glue.py {self.env.test_path}/run_glue.py'''.split()       ,		encoding="utf-8"       ,		check=UpperCAmelCase__       ,		)
								assert hasattr(self       ,		"env"     )



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__=1     ):
								# creates estimator
								return HuggingFace(
								    entry_point=self.script       ,		source_dir=self.env.test_path       ,		role=self.env.role       ,		image_uri=self.env.image_uri       ,		base_job_name=f'''{self.env.base_job_name}-single'''       ,		instance_count=UpperCAmelCase__       ,		instance_type=self.instance_type       ,		debugger_hook_config=UpperCAmelCase__       ,		hyperparameters={**self.env.hyperparameters, "model_name_or_path": self.model_name_or_path}       ,		metric_definitions=self.env.metric_definitions       ,		py_version="py36"       ,		)



				def   lowerCAmelCase__		(						self       ,		UpperCAmelCase__     ):
								TrainingJobAnalytics(UpperCAmelCase__     ).export_csv(f'''{self.env.test_path}/{job_name}_metrics.csv'''     )






				def   lowerCAmelCase__		(						self     ):
								# create estimator
								SCREAMING_SNAKE_CASE__     							=				self.create_estimator()

								# run training
								estimator.fit()

								# result dataframe
								SCREAMING_SNAKE_CASE__     							=				TrainingJobAnalytics(estimator.latest_training_job.name     ).dataframe()

								# extract kpis
								SCREAMING_SNAKE_CASE__     							=				list(result_metrics_df[result_metrics_df.metric_name == "eval_accuracy"]["value"]     )
								SCREAMING_SNAKE_CASE__     							=				list(result_metrics_df[result_metrics_df.metric_name == "eval_loss"]["value"]     )
								# get train time from SageMaker job, this includes starting, preprocessing, stopping
								SCREAMING_SNAKE_CASE__     							=				(
								    Session().describe_training_job(estimator.latest_training_job.name     ).get("TrainingTimeInSeconds"       ,		99_9999     )
								)

								# assert kpis
								assert train_runtime <= self.results["train_runtime"]
								assert all(t >= self.results["eval_accuracy"] for t in eval_accuracy     )
								assert all(t <= self.results["eval_loss"] for t in eval_loss     )

								# dump tests result into json file to share in PR
								with open(f'''{estimator.latest_training_job.name}.json'''       ,		"w"     ) as outfile:
												json.dump({"train_time": train_runtime, "eval_accuracy": eval_accuracy, "eval_loss": eval_loss}       ,		UpperCAmelCase__     )
