






'''simple docstring'''







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

__lowerCAmelCase       		=      'The Nymphenburg Palace is a beautiful palace in Munich!'





def 		_UpperCAmelCase   (							__A							:   str						,		__A							:   str      ):
				a_		:  Dict											=      {
				    '''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,
				}

				a_		:  Tuple											=      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[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
				a_		:  Optional[Any]											=      '''openwebtext_ccnews_stories_books_cased'''

				# Specify download folder to Gluonnlp's vocab
				a_		:  Dict											=      os.path.join(get_home_dir()						,		'''models'''      )
				a_		:  List[Any]											=      _load_vocab(__A						,		__A						,		__A						,		cls=__A      )

				a_		:  str											=      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_		:  Union[str, Any]											=      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_		:  List[str]											=      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							:   List[Any]      ) -> 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							:   str						,		__A							:   Tuple      ):
								a_		:  Tuple											=      hf_param.shape

								a_		:  int											=      to_torch(params[gluon_param]      )
								a_		:  Union[str, 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

				a_		:  Any											=      check_and_map_params(
				    hf_bort_model.bert.embeddings.word_embeddings.weight						,		'''word_embed.0.weight'''      )
				a_		:  List[str]											=      check_and_map_params(
				    hf_bort_model.bert.embeddings.position_embeddings.weight						,		'''encoder.position_weight'''      )
				a_		:  Optional[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_		:  Union[str, Any]											=      check_and_map_params(
								    self_attn.key.bias.data						,		f'encoder.transformer_cells.{i}.attention_cell.proj_key.bias'      )

								a_		:  Union[str, Any]											=      check_and_map_params(
								    self_attn.key.weight.data						,		f'encoder.transformer_cells.{i}.attention_cell.proj_key.weight'      )
								a_		:  Optional[int]											=      check_and_map_params(
								    self_attn.query.bias.data						,		f'encoder.transformer_cells.{i}.attention_cell.proj_query.bias'      )
								a_		:  Tuple											=      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_		:  Union[str, Any]											=      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_		:  List[Any]											=      check_and_map_params(
								    self_output.dense.bias						,		f'encoder.transformer_cells.{i}.proj.bias'      )
								a_		:  List[Any]											=      check_and_map_params(
								    self_output.dense.weight						,		f'encoder.transformer_cells.{i}.proj.weight'      )
								a_		:  Optional[int]											=      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_		:  str											=      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_		:  Union[str, Any]											=      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_		:  str											=      check_and_map_params(
								    bert_output.LayerNorm.bias						,		f'encoder.transformer_cells.{i}.ffn.layer_norm.beta'      )
								a_		:  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
				a_		:  int											=      RobertaTokenizer.from_pretrained('''roberta-base'''      )

				a_		:  Union[str, Any]											=      tokenizer.encode_plus(__A      )['''input_ids''']

				# Get gluon output
				a_		:  int											=      mx.nd.array([input_ids]      )
				a_		:  List[str]											=      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_		:  int											=      tokenizer.encode_plus(__A						,		return_tensors='''pt'''      )
				a_		:  Union[str, Any]											=      hf_bort_model(**__A      )[0]

				a_		:  Tuple											=      output_gluon[0].asnumpy()
				a_		:  Optional[Any]											=      output_hf[0].detach().numpy()

				a_		:  Optional[Any]											=      np.max(np.abs(hf_layer - gluon_layer      )      ).item()
				a_		:  Dict											=      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__":
		__lowerCAmelCase       		=      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.'
		)
		__lowerCAmelCase       		=      parser.parse_args()
		convert_bort_checkpoint_to_pytorch(args.bort_checkpoint_path, args.pytorch_dump_folder_path)












'''simple docstring'''







import warnings


warnings.warn(
    'memory_utils has been reorganized to utils.memory. Import `find_executable_batchsize` from the main `__init__`: '
    '`from accelerate import find_executable_batch_size` to avoid this warning.',
    FutureWarning,
)




