

import argparse
import os

import gluonnlp as nlp
import mxnet as mx
import numpy as np
import torch
from gluonnlp.base import get_home_dir
from gluonnlp.model.bert import BERTEncoder
from gluonnlp.model.utils import _load_vocab
from gluonnlp.vocab import Vocab
from packaging import version
from torch import nn

from transformers import BertConfig, BertForMaskedLM, BertModel, RobertaTokenizer
from transformers.models.bert.modeling_bert import (
    BertIntermediate,
    BertLayer,
    BertOutput,
    BertSelfAttention,
    BertSelfOutput,
)
from transformers.utils import logging


if version.parse(nlp.__version__) != version.parse("""0.8.3"""):
	raise Exception("""requires gluonnlp == 0.8.3""")

if version.parse(mx.__version__) != version.parse("""1.5.0"""):
	raise Exception("""requires mxnet == 1.5.0""")

logging.set_verbosity_info()
__snake_case				:      Optional[int]   		=			logging.get_logger(__name__)

__snake_case				:      Any   		=			"""The Nymphenburg Palace is a beautiful palace in Munich!"""


def 					_UpperCAmelCase				(						a__  ,	a__):

			'''simple docstring'''



			a_						:						List[Any] 							= {
			    """attention_cell""": """multi_head""",
			    """num_layers""": 4,
			    """units""": 1_0_2_4,
			    """hidden_size""": 7_6_8,
			    """max_length""": 5_1_2,
			    """num_heads""": 8,
			    """scaled""": True,
			    """dropout""": 0.1,
			    """use_residual""": True,
			    """embed_size""": 1_0_2_4,
			    """embed_dropout""": 0.1,
			    """word_embed""": None,
			    """layer_norm_eps""": 1e-5,
			    """token_type_vocab_size""": 2,
			}

			a_						:						Optional[Any] 							= bort_4_8_768_1024_hparams

			# Let's construct the original Bort model here
			# Taken from official BERT implementation, see:
			# https://github.com/alexa/bort/blob/master/bort/bort.py
			a_						:						Optional[int] 							= BERTEncoder(
			    attention_cell=predefined_args["""attention_cell"""]  ,	num_layers=predefined_args["""num_layers"""]  ,	units=predefined_args["""units"""]  ,	hidden_size=predefined_args["""hidden_size"""]  ,	max_length=predefined_args["""max_length"""]  ,	num_heads=predefined_args["""num_heads"""]  ,	scaled=predefined_args["""scaled"""]  ,	dropout=predefined_args["""dropout"""]  ,	output_attention=a__  ,	output_all_encodings=a__  ,	use_residual=predefined_args["""use_residual"""]  ,	activation=predefined_args.get("""activation"""  ,	"""gelu""")  ,	layer_norm_eps=predefined_args.get("""layer_norm_eps"""  ,	a__)  ,	)

			# Vocab information needs to be fetched first
			# It's the same as RoBERTa, so RobertaTokenizer can be used later
			a_						:						int 							= """openwebtext_ccnews_stories_books_cased"""

			# Specify download folder to Gluonnlp's vocab
			a_						:						Optional[Any] 							= os.path.join(get_home_dir()  ,	"""models""")
			a_						:						Tuple 							= _load_vocab(a__  ,	a__  ,	a__  ,	cls=a__)

			a_						:						Union[str, Any] 							= nlp.model.BERTModel(
			    a__  ,	len(a__)  ,	units=predefined_args["""units"""]  ,	embed_size=predefined_args["""embed_size"""]  ,	embed_dropout=predefined_args["""embed_dropout"""]  ,	word_embed=predefined_args["""word_embed"""]  ,	use_pooler=a__  ,	use_token_type_embed=a__  ,	token_type_vocab_size=predefined_args["""token_type_vocab_size"""]  ,	use_classifier=a__  ,	use_decoder=a__  ,	)

			original_bort.load_parameters(a__  ,	cast_dtype=a__  ,	ignore_extra=a__)
			a_						:						Tuple 							= original_bort._collect_params_with_prefix()

			# Build our config 🤗
			a_						:						List[Any] 							= {
			    """architectures""": ["""BertForMaskedLM"""],
			    """attention_probs_dropout_prob""": predefined_args["""dropout"""],
			    """hidden_act""": """gelu""",
			    """hidden_dropout_prob""": predefined_args["""dropout"""],
			    """hidden_size""": predefined_args["""embed_size"""],
			    """initializer_range""": 0.02,
			    """intermediate_size""": predefined_args["""hidden_size"""],
			    """layer_norm_eps""": predefined_args["""layer_norm_eps"""],
			    """max_position_embeddings""": predefined_args["""max_length"""],
			    """model_type""": """bort""",
			    """num_attention_heads""": predefined_args["""num_heads"""],
			    """num_hidden_layers""": predefined_args["""num_layers"""],
			    """pad_token_id""": 1,  # 2 = BERT, 1 = RoBERTa
			    """type_vocab_size""": 1,  # 2 = BERT, 1 = RoBERTa
			    """vocab_size""": len(a__),
			}

			a_						:						Any 							= BertConfig.from_dict(a__)
			a_						:						Dict 							= BertForMaskedLM(a__)
			hf_bort_model.eval()

			# Parameter mapping table (Gluonnlp to Transformers)
			# * denotes layer index
			#
			# | Gluon Parameter                                                | Transformers Parameter
			# | -------------------------------------------------------------- | ----------------------
			# | `encoder.layer_norm.beta`                                      | `bert.embeddings.LayerNorm.bias`
			# | `encoder.layer_norm.gamma`                                     | `bert.embeddings.LayerNorm.weight`
			# | `encoder.position_weight`                                      | `bert.embeddings.position_embeddings.weight`
			# | `word_embed.0.weight`                                          | `bert.embeddings.word_embeddings.weight`
			# | `encoder.transformer_cells.*.attention_cell.proj_key.bias`     | `bert.encoder.layer.*.attention.self.key.bias`
			# | `encoder.transformer_cells.*.attention_cell.proj_key.weight`   | `bert.encoder.layer.*.attention.self.key.weight`
			# | `encoder.transformer_cells.*.attention_cell.proj_query.bias`   | `bert.encoder.layer.*.attention.self.query.bias`
			# | `encoder.transformer_cells.*.attention_cell.proj_query.weight` | `bert.encoder.layer.*.attention.self.query.weight`
			# | `encoder.transformer_cells.*.attention_cell.proj_value.bias`   | `bert.encoder.layer.*.attention.self.value.bias`
			# | `encoder.transformer_cells.*.attention_cell.proj_value.weight` | `bert.encoder.layer.*.attention.self.value.weight`
			# | `encoder.transformer_cells.*.ffn.ffn_2.bias`                   | `bert.encoder.layer.*.attention.output.dense.bias`
			# | `encoder.transformer_cells.*.ffn.ffn_2.weight`                 | `bert.encoder.layer.*.attention.output.dense.weight`
			# | `encoder.transformer_cells.*.layer_norm.beta`                  | `bert.encoder.layer.*.attention.output.LayerNorm.bias`
			# | `encoder.transformer_cells.*.layer_norm.gamma`                 | `bert.encoder.layer.*.attention.output.LayerNorm.weight`
			# | `encoder.transformer_cells.*.ffn.ffn_1.bias`                   | `bert.encoder.layer.*.intermediate.dense.bias`
			# | `encoder.transformer_cells.*.ffn.ffn_1.weight`                 | `bert.encoder.layer.*.intermediate.dense.weight`
			# | `encoder.transformer_cells.*.ffn.layer_norm.beta`              | `bert.encoder.layer.*.output.LayerNorm.bias`
			# | `encoder.transformer_cells.*.ffn.layer_norm.gamma`             | `bert.encoder.layer.*.output.LayerNorm.weight`
			# | `encoder.transformer_cells.*.proj.bias`                        | `bert.encoder.layer.*.output.dense.bias`
			# | `encoder.transformer_cells.*.proj.weight`                      | `bert.encoder.layer.*.output.dense.weight`

			# Helper function to convert MXNET Arrays to PyTorch
			def to_torch(a__) -> nn.Parameter:
						return nn.Parameter(torch.FloatTensor(mx_array.data().asnumpy()))

			# Check param shapes and map new HF param back
			def check_and_map_params(a__  ,	a__):
						a_						:						Optional[int] 							= hf_param.shape

						a_						:						str 							= to_torch(params[gluon_param])
						a_						:						Optional[int] 							= gluon_param.shape

						assert (
						    shape_hf == shape_gluon
						), f'''The gluon parameter {gluon_param} has shape {shape_gluon}, but expects shape {shape_hf} for Transformers'''

						return gluon_param

			a_						:						List[Any] 							= check_and_map_params(
			    hf_bort_model.bert.embeddings.word_embeddings.weight  ,	"""word_embed.0.weight""")
			a_						:						Optional[int] 							= check_and_map_params(
			    hf_bort_model.bert.embeddings.position_embeddings.weight  ,	"""encoder.position_weight""")
			a_						:						int 							= check_and_map_params(
			    hf_bort_model.bert.embeddings.LayerNorm.bias  ,	"""encoder.layer_norm.beta""")
			a_						:						Union[str, Any] 							= check_and_map_params(
			    hf_bort_model.bert.embeddings.LayerNorm.weight  ,	"""encoder.layer_norm.gamma""")

			# Inspired by RoBERTa conversion script, we just zero them out (Bort does not use them)
			a_						:						str 							= torch.zeros_like(
			    hf_bort_model.bert.embeddings.token_type_embeddings.weight.data)

			for i in range(hf_bort_config.num_hidden_layers):
						a_						:						BertLayer 							= hf_bort_model.bert.encoder.layer[i]

						# self attention
						a_						:						BertSelfAttention 							= layer.attention.self

						a_						:						Tuple 							= check_and_map_params(
						    self_attn.key.bias.data  ,	f'''encoder.transformer_cells.{i}.attention_cell.proj_key.bias''')

						a_						:						int 							= check_and_map_params(
						    self_attn.key.weight.data  ,	f'''encoder.transformer_cells.{i}.attention_cell.proj_key.weight''')
						a_						:						Any 							= check_and_map_params(
						    self_attn.query.bias.data  ,	f'''encoder.transformer_cells.{i}.attention_cell.proj_query.bias''')
						a_						:						Union[str, Any] 							= check_and_map_params(
						    self_attn.query.weight.data  ,	f'''encoder.transformer_cells.{i}.attention_cell.proj_query.weight''')
						a_						:						Any 							= check_and_map_params(
						    self_attn.value.bias.data  ,	f'''encoder.transformer_cells.{i}.attention_cell.proj_value.bias''')
						a_						:						Optional[int] 							= check_and_map_params(
						    self_attn.value.weight.data  ,	f'''encoder.transformer_cells.{i}.attention_cell.proj_value.weight''')

						# self attention output
						a_						:						BertSelfOutput 							= layer.attention.output

						a_						:						Union[str, Any] 							= check_and_map_params(
						    self_output.dense.bias  ,	f'''encoder.transformer_cells.{i}.proj.bias''')
						a_						:						List[str] 							= check_and_map_params(
						    self_output.dense.weight  ,	f'''encoder.transformer_cells.{i}.proj.weight''')
						a_						:						Tuple 							= check_and_map_params(
						    self_output.LayerNorm.bias  ,	f'''encoder.transformer_cells.{i}.layer_norm.beta''')
						a_						:						Dict 							= check_and_map_params(
						    self_output.LayerNorm.weight  ,	f'''encoder.transformer_cells.{i}.layer_norm.gamma''')

						# intermediate
						a_						:						BertIntermediate 							= layer.intermediate

						a_						:						Dict 							= check_and_map_params(
						    intermediate.dense.bias  ,	f'''encoder.transformer_cells.{i}.ffn.ffn_1.bias''')
						a_						:						Optional[Any] 							= check_and_map_params(
						    intermediate.dense.weight  ,	f'''encoder.transformer_cells.{i}.ffn.ffn_1.weight''')

						# output
						a_						:						BertOutput 							= layer.output

						a_						:						Optional[int] 							= check_and_map_params(
						    bert_output.dense.bias  ,	f'''encoder.transformer_cells.{i}.ffn.ffn_2.bias''')
						a_						:						Union[str, Any] 							= check_and_map_params(
						    bert_output.dense.weight  ,	f'''encoder.transformer_cells.{i}.ffn.ffn_2.weight''')
						a_						:						List[Any] 							= check_and_map_params(
						    bert_output.LayerNorm.bias  ,	f'''encoder.transformer_cells.{i}.ffn.layer_norm.beta''')
						a_						:						Optional[int] 							= check_and_map_params(
						    bert_output.LayerNorm.weight  ,	f'''encoder.transformer_cells.{i}.ffn.layer_norm.gamma''')

			# Save space and energy 🎄
			hf_bort_model.half()

			# Compare output of both models
			a_						:						Optional[int] 							= RobertaTokenizer.from_pretrained("""roberta-base""")

			a_						:						Dict 							= tokenizer.encode_plus(a__)["""input_ids"""]

			# Get gluon output
			a_						:						List[str] 							= mx.nd.array([input_ids])
			a_						:						Tuple 							= original_bort(inputs=a__  ,	token_types=[])

			# Get Transformer output (save and reload model again)
			hf_bort_model.save_pretrained(a__)
			a_						:						Optional[Any] 							= BertModel.from_pretrained(a__)
			hf_bort_model.eval()

			a_						:						List[str] 							= tokenizer.encode_plus(a__  ,	return_tensors="""pt""")
			a_						:						Optional[Any] 							= hf_bort_model(**a__)[0]

			a_						:						Union[str, Any] 							= output_gluon[0].asnumpy()
			a_						:						Optional[int] 							= output_hf[0].detach().numpy()

			a_						:						Any 							= np.max(np.abs(hf_layer - gluon_layer)).item()
			a_						:						List[str] 							= np.allclose(a__  ,	a__  ,	atol=1e-3)

			if success:
						print("""✔️ Both model do output the same tensors""")
			else:
						print("""❌ Both model do **NOT** output the same tensors""")
						print("""Absolute difference is:"""  ,	a__)


if __name__ == "__main__":
	__snake_case				:      Any   		=			argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    """--bort_checkpoint_path""", default=None, type=str, required=True, help="""Path the official Bort params file."""
	)
	parser.add_argument(
	    """--pytorch_dump_folder_path""", default=None, type=str, required=True, help="""Path to the output PyTorch model."""
	)
	__snake_case				:      str   		=			parser.parse_args()
	convert_bort_checkpoint_to_pytorch(args.bort_checkpoint_path, args.pytorch_dump_folder_path)





import logging

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,
    BertEncoder,
    BertModel,
    BertPreTrainedModel,
)


__snake_case				:      Optional[int]   		=			logging.getLogger(__name__)



class        A__(a_							):
						"""simple docstring"""






						def 			UpperCamelCase__					(    self  ,      _lowercase  ,      _lowercase  ,      _lowercase=None  ,      _lowercase=None    )     -> Tuple:
									a_						:						str 							= self.layer[current_layer](_lowercase  ,      _lowercase  ,      head_mask[current_layer]    )

									a_						:						str 							= layer_outputs[0]

									return hidden_states







@add_start_docstrings(
    '''The bare Bert Model transformer with PABEE outputting raw hidden-states without any specific head on top.''',  a_,  )
class        A__(a_							):
						"""simple docstring"""






						def __init__(    self  ,      _lowercase    )     -> List[str]:
									super().__init__(_lowercase    )

									a_						:						Tuple 							= BertEncoderWithPabee(_lowercase    )

									self.init_weights()
									a_						:						int 							= 0
									a_						:						Any 							= 0
									a_						:						Tuple 							= 0

									a_						:						Optional[int] 							= 0





						def 			UpperCamelCase__					(    self  ,      _lowercase    )     -> Tuple:
									a_						:						Dict 							= threshold





						def 			UpperCamelCase__					(    self  ,      _lowercase    )     -> List[Any]:
									a_						:						Optional[int] 							= patience





						def 			UpperCamelCase__					(    self    )     -> Dict:
									a_						:						str 							= 0
									a_						:						Optional[int] 							= 0





						def 			UpperCamelCase__					(    self    )     -> List[Any]:
									a_						:						Union[str, Any] 							= self.inference_layers_num / self.inference_instances_num
									a_						:						Optional[int] 							= (
									    F'''*** Patience = {self.patience} Avg. Inference Layers = {avg_inf_layers:.2f} Speed Up ='''
									    F''' {1 - avg_inf_layers / self.config.num_hidden_layers:.2f} ***'''
									)
									print(_lowercase    )





						@add_start_docstrings_to_model_forward(_lowercase    )
						def 			UpperCamelCase__					(    self  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=False  ,      )     -> str:

									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:
												a_						:						Dict 							= input_ids.size()
									elif inputs_embeds is not None:
												a_						:						Dict 							= inputs_embeds.size()[:-1]
									else:
												raise ValueError("""You have to specify either input_ids or inputs_embeds"""    )

									a_						:						Optional[Any] 							= input_ids.device if input_ids is not None else inputs_embeds.device

									if attention_mask is None:
												a_						:						Tuple 							= torch.ones(_lowercase  ,      device=_lowercase    )
									if token_type_ids is None:
												a_						:						List[str] 							= 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.
									a_						:						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 self.config.is_decoder and encoder_hidden_states is not None:
												a_							,					a_							,					a_						:						int 							= encoder_hidden_states.size()
												a_						:						List[str] 							= (encoder_batch_size, encoder_sequence_length)
												if encoder_attention_mask is None:
															a_						:						Tuple 							= torch.ones(_lowercase  ,      device=_lowercase    )
												a_						:						List[Any] 							= self.invert_attention_mask(_lowercase    )
									else:
												a_						:						Optional[Any] 							= None

									# 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]
									a_						:						List[Any] 							= self.get_head_mask(_lowercase  ,      self.config.num_hidden_layers    )

									a_						:						List[str] 							= self.embeddings(
									    input_ids=_lowercase  ,      position_ids=_lowercase  ,      token_type_ids=_lowercase  ,      inputs_embeds=_lowercase    )
									a_						:						List[Any] 							= embedding_output

									if self.training:
												a_						:						Any 							= []
												for i in range(self.config.num_hidden_layers    ):
															a_						:						int 							= self.encoder.adaptive_forward(
															    _lowercase  ,      current_layer=_lowercase  ,      attention_mask=_lowercase  ,      head_mask=_lowercase    )

															a_						:						List[Any] 							= self.pooler(_lowercase    )
															a_						:						Optional[int] 							= output_layers[i](output_dropout(_lowercase    )    )
															res.append(_lowercase    )
									elif self.patience == 0:  # Use all layers for inference
												a_						:						Union[str, Any] 							= self.encoder(
												    _lowercase  ,      attention_mask=_lowercase  ,      head_mask=_lowercase  ,      encoder_hidden_states=_lowercase  ,      encoder_attention_mask=_lowercase  ,      )
												a_						:						Union[str, Any] 							= self.pooler(encoder_outputs[0]    )
												a_						:						List[str] 							= [output_layers[self.config.num_hidden_layers - 1](_lowercase    )]
									else:
												a_						:						Any 							= 0
												a_						:						Dict 							= None
												a_						:						Tuple 							= 0
												for i in range(self.config.num_hidden_layers    ):
															calculated_layer_num += 1
															a_						:						Optional[Any] 							= self.encoder.adaptive_forward(
															    _lowercase  ,      current_layer=_lowercase  ,      attention_mask=_lowercase  ,      head_mask=_lowercase    )

															a_						:						Optional[int] 							= self.pooler(_lowercase    )
															a_						:						int 							= output_layers[i](_lowercase    )
															if regression:
																		a_						:						Dict 							= logits.detach()
																		if patient_result is not None:
																					a_						:						Optional[Any] 							= patient_result.detach()
																		if (patient_result is not None) and torch.abs(patient_result - labels    ) < self.regression_threshold:
																					patient_counter += 1
																		else:
																					a_						:						Optional[Any] 							= 0
															else:
																		a_						:						str 							= logits.detach().argmax(dim=1    )
																		if patient_result is not None:
																					a_						:						str 							= patient_result.detach().argmax(dim=1    )
																		if (patient_result is not None) and torch.all(labels.eq(_lowercase    )    ):
																					patient_counter += 1
																		else:
																					a_						:						Tuple 							= 0

															a_						:						Union[str, Any] 							= logits
															if patient_counter == self.patience:
																		break
												a_						:						str 							= [patient_result]
												self.inference_layers_num += calculated_layer_num
												self.inference_instances_num += 1

									return res







@add_start_docstrings(
    '''Bert Model transformer with PABEE and a sequence classification/regression head on top (a linear layer on top of
    the pooled output) e.g. for GLUE tasks. ''',  a_,  )
class        A__(a_							):
						"""simple docstring"""






						def __init__(    self  ,      _lowercase    )     -> str:
									super().__init__(_lowercase    )
									a_						:						str 							= config.num_labels

									a_						:						Optional[Any] 							= BertModelWithPabee(_lowercase    )
									a_						:						int 							= nn.Dropout(config.hidden_dropout_prob    )
									a_						:						str 							= nn.ModuleList(
									    [nn.Linear(config.hidden_size  ,      self.config.num_labels    ) for _ in range(config.num_hidden_layers    )]    )

									self.init_weights()





						@add_start_docstrings_to_model_forward(_lowercase    )
						def 			UpperCamelCase__					(    self  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      _lowercase=None  ,      )     -> Tuple:
									a_						:						Optional[Any] 							= self.bert(
									    input_ids=_lowercase  ,      attention_mask=_lowercase  ,      token_type_ids=_lowercase  ,      position_ids=_lowercase  ,      head_mask=_lowercase  ,      inputs_embeds=_lowercase  ,      output_dropout=self.dropout  ,      output_layers=self.classifiers  ,      regression=self.num_labels == 1  ,      )

									a_						:						Optional[Any] 							= (logits[-1],)

									if labels is not None:
												a_						:						int 							= None
												a_						:						Union[str, Any] 							= 0
												for ix, logits_item in enumerate(_lowercase    ):
															if self.num_labels == 1:
																		#  We are doing regression
																		a_						:						Any 							= MSELoss()
																		a_						:						Union[str, Any] 							= loss_fct(logits_item.view(-1    )  ,      labels.view(-1    )    )
															else:
																		a_						:						Any 							= CrossEntropyLoss()
																		a_						:						int 							= loss_fct(logits_item.view(-1  ,      self.num_labels    )  ,      labels.view(-1    )    )
															if total_loss is None:
																		a_						:						str 							= loss
															else:
																		total_loss += loss * (ix + 1)
															total_weights += ix + 1
												a_						:						Any 							= (total_loss / total_weights,) + outputs

									return outputs


