





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()
_SCREAMING_SNAKE_CASE 		=						logging.get_logger(__name__)

_SCREAMING_SNAKE_CASE 		=						"""The Nymphenburg Palace is a beautiful palace in Munich!"""


def 							SCREAMING_SNAKE_CASE__  (      __a				,	__a				):
	snake_case_						:     Optional[int]     =		{
	    'attention_cell': 'multi_head',
	    'num_layers': 4,
	    'units': 10_24,
	    'hidden_size': 7_68,
	    'max_length': 5_12,
	    'num_heads': 8,
	    'scaled': True,
	    'dropout': 0.1,
	    'use_residual': True,
	    'embed_size': 10_24,
	    'embed_dropout': 0.1,
	    'word_embed': None,
	    'layer_norm_eps': 1E-5,
	    'token_type_vocab_size': 2,
	}

	snake_case_						:     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
	snake_case_						:     Any     =		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
	snake_case_						:     str     =		'openwebtext_ccnews_stories_books_cased'

	# Specify download folder to Gluonnlp's vocab
	snake_case_						:     List[str]     =		os.path.join(get_home_dir()				,	'models'				)
	snake_case_						:     str     =		_load_vocab(__a				,	__a				,	__a				,	cls=__a				)

	snake_case_						:     Tuple     =		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				)
	snake_case_						:     Any     =		original_bort._collect_params_with_prefix()

	# Build our config 🤗
	snake_case_						:     List[str]     =		{
	    '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				),
	}

	snake_case_						:     Optional[int]     =		BertConfig.from_dict(__a				)
	snake_case_						:     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				):
		snake_case_						:     List[Any]     =		hf_param.shape

		snake_case_						:     Dict     =		to_torch(params[gluon_param]				)
		snake_case_						:     Any     =		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

	snake_case_						:     Any     =		check_and_map_params(
	    hf_bort_model.bert.embeddings.word_embeddings.weight				,	'word_embed.0.weight'				)
	snake_case_						:     Dict     =		check_and_map_params(
	    hf_bort_model.bert.embeddings.position_embeddings.weight				,	'encoder.position_weight'				)
	snake_case_						:     Tuple     =		check_and_map_params(
	    hf_bort_model.bert.embeddings.LayerNorm.bias				,	'encoder.layer_norm.beta'				)
	snake_case_						:     Dict     =		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)
	snake_case_						:     List[Any]     =		torch.zeros_like(
	    hf_bort_model.bert.embeddings.token_type_embeddings.weight.data				)

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

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

		snake_case_						:     Optional[int]     =		check_and_map_params(
		    self_attn.key.bias.data				,	f"""encoder.transformer_cells.{i}.attention_cell.proj_key.bias"""				)

		snake_case_						:     str     =		check_and_map_params(
		    self_attn.key.weight.data				,	f"""encoder.transformer_cells.{i}.attention_cell.proj_key.weight"""				)
		snake_case_						:     Any     =		check_and_map_params(
		    self_attn.query.bias.data				,	f"""encoder.transformer_cells.{i}.attention_cell.proj_query.bias"""				)
		snake_case_						:     str     =		check_and_map_params(
		    self_attn.query.weight.data				,	f"""encoder.transformer_cells.{i}.attention_cell.proj_query.weight"""				)
		snake_case_						:     List[Any]     =		check_and_map_params(
		    self_attn.value.bias.data				,	f"""encoder.transformer_cells.{i}.attention_cell.proj_value.bias"""				)
		snake_case_						:     List[str]     =		check_and_map_params(
		    self_attn.value.weight.data				,	f"""encoder.transformer_cells.{i}.attention_cell.proj_value.weight"""				)

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

		snake_case_						:     List[str]     =		check_and_map_params(
		    self_output.dense.bias				,	f"""encoder.transformer_cells.{i}.proj.bias"""				)
		snake_case_						:     Any     =		check_and_map_params(
		    self_output.dense.weight				,	f"""encoder.transformer_cells.{i}.proj.weight"""				)
		snake_case_						:     int     =		check_and_map_params(
		    self_output.LayerNorm.bias				,	f"""encoder.transformer_cells.{i}.layer_norm.beta"""				)
		snake_case_						:     Any     =		check_and_map_params(
		    self_output.LayerNorm.weight				,	f"""encoder.transformer_cells.{i}.layer_norm.gamma"""				)

		# intermediate
		snake_case_						:     BertIntermediate     =		layer.intermediate

		snake_case_						:     Optional[int]     =		check_and_map_params(
		    intermediate.dense.bias				,	f"""encoder.transformer_cells.{i}.ffn.ffn_1.bias"""				)
		snake_case_						:     Tuple     =		check_and_map_params(
		    intermediate.dense.weight				,	f"""encoder.transformer_cells.{i}.ffn.ffn_1.weight"""				)

		# output
		snake_case_						:     BertOutput     =		layer.output

		snake_case_						:     Optional[Any]     =		check_and_map_params(
		    bert_output.dense.bias				,	f"""encoder.transformer_cells.{i}.ffn.ffn_2.bias"""				)
		snake_case_						:     Tuple     =		check_and_map_params(
		    bert_output.dense.weight				,	f"""encoder.transformer_cells.{i}.ffn.ffn_2.weight"""				)
		snake_case_						:     Dict     =		check_and_map_params(
		    bert_output.LayerNorm.bias				,	f"""encoder.transformer_cells.{i}.ffn.layer_norm.beta"""				)
		snake_case_						:     List[str]     =		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
	snake_case_						:     List[str]     =		RobertaTokenizer.from_pretrained('roberta-base'				)

	snake_case_						:     str     =		tokenizer.encode_plus(__a				)['input_ids']

	# Get gluon output
	snake_case_						:     Tuple     =		mx.nd.array([input_ids]				)
	snake_case_						:     Optional[Any]     =		original_bort(inputs=__a				,	token_types=[]				)

	# Get Transformer output (save and reload model again)
	hf_bort_model.save_pretrained(__a				)
	snake_case_						:     List[Any]     =		BertModel.from_pretrained(__a				)
	hf_bort_model.eval()

	snake_case_						:     Any     =		tokenizer.encode_plus(__a				,	return_tensors='pt'				)
	snake_case_						:     str     =		hf_bort_model(**__a				)[0]

	snake_case_						:     List[str]     =		output_gluon[0].asnumpy()
	snake_case_						:     Union[str, Any]     =		output_hf[0].detach().numpy()

	snake_case_						:     Union[str, Any]     =		np.max(np.abs(hf_layer - gluon_layer				)				).item()
	snake_case_						:     List[Any]     =		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__":
						_SCREAMING_SNAKE_CASE 		=						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."""
						)
						_SCREAMING_SNAKE_CASE 		=						parser.parse_args()
						convert_bort_checkpoint_to_pytorch(args.bort_checkpoint_path, args.pytorch_dump_folder_path)








from __future__ import annotations

import unittest

from transformers import AutoTokenizer, MBartConfig, is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
						import tensorflow as tf

						from transformers import TFAutoModelForSeqaSeqLM, TFMBartForConditionalGeneration, TFMBartModel
@require_tf
class SCREAMING_SNAKE_CASE_  :
						__magic_name__:     int	 =    MBartConfig
						__magic_name__:     str	 =    {}
						__magic_name__:     Union[str, Any]	 =    "gelu"






						def __init__( self : List[str]						,				_A : Optional[int]						,				_A : List[Any]=13						,				_A : List[Any]=7						,				_A : Dict=True						,				_A : Tuple=False						,				_A : Optional[Any]=99						,				_A : Dict=32						,				_A : str=2						,				_A : str=4						,				_A : Tuple=37						,				_A : Tuple=0.1						,				_A : Union[str, Any]=0.1						,				_A : Optional[int]=20						,				_A : Dict=2						,				_A : List[str]=1						,				_A : Union[str, Any]=0						,				)							->      List[Any]:







							"""simple docstring"""


							snake_case_						:     str     =		parent
							snake_case_						:     List[str]     =		batch_size
							snake_case_						:     List[str]     =		seq_length
							snake_case_						:     Union[str, Any]     =		is_training
							snake_case_						:     Optional[int]     =		use_labels
							snake_case_						:     Dict     =		vocab_size
							snake_case_						:     Union[str, Any]     =		hidden_size
							snake_case_						:     str     =		num_hidden_layers
							snake_case_						:     Optional[Any]     =		num_attention_heads
							snake_case_						:     Union[str, Any]     =		intermediate_size
							snake_case_						:     Any     =		hidden_dropout_prob
							snake_case_						:     Any     =		attention_probs_dropout_prob
							snake_case_						:     Optional[Any]     =		max_position_embeddings
							snake_case_						:     Optional[Any]     =		eos_token_id
							snake_case_						:     Tuple     =		pad_token_id
							snake_case_						:     int     =		bos_token_id






						def UpperCAmelCase_			( self : List[str]		)							->      Tuple:







							"""simple docstring"""


							snake_case_						:     List[Any]     =		ids_tensor([self.batch_size, self.seq_length - 1]						,				self.vocab_size		)
							snake_case_						:     int     =		tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size		)						,				1		)
							snake_case_						:     Dict     =		tf.concat([input_ids, eos_tensor]						,				axis=1		)

							snake_case_						:     Any     =		ids_tensor([self.batch_size, self.seq_length]						,				self.vocab_size		)

							snake_case_						:     Optional[Any]     =		self.config_cls(
							    vocab_size=self.vocab_size						,				d_model=self.hidden_size						,				encoder_layers=self.num_hidden_layers						,				decoder_layers=self.num_hidden_layers						,				encoder_attention_heads=self.num_attention_heads						,				decoder_attention_heads=self.num_attention_heads						,				encoder_ffn_dim=self.intermediate_size						,				decoder_ffn_dim=self.intermediate_size						,				dropout=self.hidden_dropout_prob						,				attention_dropout=self.attention_probs_dropout_prob						,				max_position_embeddings=self.max_position_embeddings						,				eos_token_ids=[2]						,				bos_token_id=self.bos_token_id						,				pad_token_id=self.pad_token_id						,				decoder_start_token_id=self.pad_token_id						,				**self.config_updates						,				)
							snake_case_						:     Union[str, Any]     =		prepare_mbart_inputs_dict(_A						,				_A						,				_A		)
							return config, inputs_dict






						def UpperCAmelCase_			( self : Optional[Any]						,				_A : Optional[Any]						,				_A : int		)							->      str:







							"""simple docstring"""


							snake_case_						:     Dict     =		TFMBartModel(config=_A		).get_decoder()
							snake_case_						:     Any     =		inputs_dict['input_ids']

							snake_case_						:     List[Any]     =		input_ids[:1, :]
							snake_case_						:     Dict     =		inputs_dict['attention_mask'][:1, :]
							snake_case_						:     Tuple     =		inputs_dict['head_mask']
							snake_case_						:     List[Any]     =		1

							# first forward pass
							snake_case_						:     Any     =		model(_A						,				attention_mask=_A						,				head_mask=_A						,				use_cache=_A		)

							snake_case_  ,snake_case_						:     str     =		outputs.to_tuple()
							snake_case_						:     int     =		past_key_values[1]


def 							SCREAMING_SNAKE_CASE__  (      __a				,	__a				,	__a				,	__a=None				,	__a=None				,	__a=None				,	__a=None				,	__a=None				,	):
	if attention_mask is None:
		snake_case_						:     Optional[int]     =		tf.cast(tf.math.not_equal(__a				,	config.pad_token_id				)				,	tf.inta				)
	if decoder_attention_mask is None:
		snake_case_						:     str     =		tf.concat(
		    [
		        tf.ones(decoder_input_ids[:, :1].shape				,	dtype=tf.inta				),
		        tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:]				,	config.pad_token_id				)				,	tf.inta				),
		    ]				,	axis=-1				,	)
	if head_mask is None:
		snake_case_						:     str     =		tf.ones((config.encoder_layers, config.encoder_attention_heads)				)
	if decoder_head_mask is None:
		snake_case_						:     int     =		tf.ones((config.decoder_layers, config.decoder_attention_heads)				)
	if cross_attn_head_mask is None:
		snake_case_						:     Dict     =		tf.ones((config.decoder_layers, config.decoder_attention_heads)				)
	return {
	    "input_ids": input_ids,
	    "decoder_input_ids": decoder_input_ids,
	    "attention_mask": attention_mask,
	    "decoder_attention_mask": decoder_attention_mask,
	    "head_mask": head_mask,
	    "decoder_head_mask": decoder_head_mask,
	    "cross_attn_head_mask": cross_attn_head_mask,
	}
@require_tf
class SCREAMING_SNAKE_CASE_  (   snake_case_    ,			snake_case_    ,			unittest.TestCase   ):
						__magic_name__:     Tuple	 =    (TFMBartForConditionalGeneration, TFMBartModel) if is_tf_available() else ()
						__magic_name__:     int	 =    (TFMBartForConditionalGeneration,) if is_tf_available() else ()
						__magic_name__:     Union[str, Any]	 =    (
						    {
						        "conversational": TFMBartForConditionalGeneration,
						        "feature-extraction": TFMBartModel,
						        "summarization": TFMBartForConditionalGeneration,
						        "text2text-generation": TFMBartForConditionalGeneration,
						        "translation": TFMBartForConditionalGeneration,
						    }
						    if is_tf_available()
						    else {}
						)
						__magic_name__:     Tuple	 =    True
						__magic_name__:     Tuple	 =    False
						__magic_name__:     Any	 =    False






						def UpperCAmelCase_			( self : Any						,				_A : Union[str, Any]						,				_A : List[Any]						,				_A : str						,				_A : int						,				_A : Dict		)							->      Union[str, Any]:







							"""simple docstring"""


							if pipeline_test_casse_name != "FeatureExtractionPipelineTests":
								# Exception encountered when calling layer '...'
								return True

							return False






						def UpperCAmelCase_			( self : Dict		)							->      Tuple:







							"""simple docstring"""


							snake_case_						:     Optional[Any]     =		TFMBartModelTester(self		)
							snake_case_						:     List[Any]     =		ConfigTester(self						,				config_class=_A		)






						def UpperCAmelCase_			( self : Optional[Any]		)							->      str:







							"""simple docstring"""


							self.config_tester.run_common_tests()






						def UpperCAmelCase_			( self : Union[str, Any]		)							->      Tuple:







							"""simple docstring"""


							snake_case_						:     Dict     =		self.model_tester.prepare_config_and_inputs_for_common()
							self.model_tester.check_decoder_model_past_large_inputs(*_A		)
@require_sentencepiece
@require_tokenizers
@require_tf
class SCREAMING_SNAKE_CASE_  (   unittest.TestCase   ):
						__magic_name__:     Optional[int]	 =    [
						    " UN Chief Says There Is No Military Solution in Syria",
						]
						__magic_name__:     Union[str, Any]	 =    [
						    "Şeful ONU declară că nu există o soluţie militară în Siria",
						]
						__magic_name__:     List[Any]	 =    "facebook/mbart-large-en-ro"






						@cached_property
						def UpperCAmelCase_			( self : str		)							->      List[Any]:







							"""simple docstring"""


							return AutoTokenizer.from_pretrained(self.model_name		)






						@cached_property
						def UpperCAmelCase_			( self : List[Any]		)							->      Any:







							"""simple docstring"""


							snake_case_						:     List[str]     =		TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name		)
							return model






						def UpperCAmelCase_			( self : Optional[int]						,				**_A : str		)							->      int:







							"""simple docstring"""


							snake_case_						:     List[str]     =		self.translate_src_text(**_A		)
							self.assertListEqual(self.expected_text						,				_A		)






						def UpperCAmelCase_			( self : Union[str, Any]						,				**_A : Dict		)							->      int:







							"""simple docstring"""


							snake_case_						:     Optional[Any]     =		self.tokenizer(self.src_text						,				**_A						,				return_tensors='tf'		)
							snake_case_						:     int     =		self.model.generate(
							    model_inputs.input_ids						,				attention_mask=model_inputs.attention_mask						,				num_beams=2		)
							snake_case_						:     Any     =		self.tokenizer.batch_decode(_A						,				skip_special_tokens=_A		)
							return generated_words






						@slow
						def UpperCAmelCase_			( self : str		)							->      List[str]:







							"""simple docstring"""


							self._assert_generated_batch_equal_expected()

