



"""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()
_lowercase				:				str   		=logging.get_logger(__name__)

_lowercase				:				Optional[Any]   		="The Nymphenburg Palace is a beautiful palace in Munich!"






def   snake_case__     (	__lowerCamelCase       :					str							,	__lowerCamelCase       :					str			):
			"""simple docstring"""
			lowerCamelCase__	:	List[Any]				={
			    '''attention_cell''': '''multi_head''',
			    '''num_layers''': 4,
			    '''units''': 1024,
			    '''hidden_size''': 768,
			    '''max_length''': 512,
			    '''num_heads''': 8,
			    '''scaled''': True,
			    '''dropout''': 0.1,
			    '''use_residual''': True,
			    '''embed_size''': 1024,
			    '''embed_dropout''': 0.1,
			    '''word_embed''': None,
			    '''layer_norm_eps''': 1e-5,
			    '''token_type_vocab_size''': 2,
			}

			lowerCamelCase__	:	List[str]				=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
			lowerCamelCase__	:	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=__lowerCamelCase							,	output_all_encodings=__lowerCamelCase							,	use_residual=predefined_args['''use_residual''']							,	activation=predefined_args.get('''activation'''							,	'''gelu'''			)							,	layer_norm_eps=predefined_args.get('''layer_norm_eps'''							,	__lowerCamelCase			)							,	)

			# Vocab information needs to be fetched first
			# It's the same as RoBERTa, so RobertaTokenizer can be used later
			lowerCamelCase__	:	Optional[Any]				='''openwebtext_ccnews_stories_books_cased'''

			# Specify download folder to Gluonnlp's vocab
			lowerCamelCase__	:	Union[str, Any]				=os.path.join(get_home_dir()							,	'''models'''			)
			lowerCamelCase__	:	Any				=_load_vocab(__lowerCamelCase							,	__lowerCamelCase							,	__lowerCamelCase							,	cls=__lowerCamelCase			)

			lowerCamelCase__	:	Any				=nlp.model.BERTModel(
			    __lowerCamelCase							,	len(__lowerCamelCase			)							,	units=predefined_args['''units''']							,	embed_size=predefined_args['''embed_size''']							,	embed_dropout=predefined_args['''embed_dropout''']							,	word_embed=predefined_args['''word_embed''']							,	use_pooler=__lowerCamelCase							,	use_token_type_embed=__lowerCamelCase							,	token_type_vocab_size=predefined_args['''token_type_vocab_size''']							,	use_classifier=__lowerCamelCase							,	use_decoder=__lowerCamelCase							,	)

			original_bort.load_parameters(__lowerCamelCase							,	cast_dtype=__lowerCamelCase							,	ignore_extra=__lowerCamelCase			)
			lowerCamelCase__	:	Optional[int]				=original_bort._collect_params_with_prefix()

			# Build our config 🤗
			lowerCamelCase__	:	Union[str, 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(__lowerCamelCase			),
			}

			lowerCamelCase__	:	int				=BertConfig.from_dict(__lowerCamelCase			)
			lowerCamelCase__	:	Dict				=BertForMaskedLM(__lowerCamelCase			)
			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(__lowerCamelCase       :					Dict			) -> 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(__lowerCamelCase       :					Union[str, Any]							,	__lowerCamelCase       :					Optional[int]			):
						lowerCamelCase__	:	Union[str, Any]				=hf_param.shape

						lowerCamelCase__	:	Optional[int]				=to_torch(params[gluon_param]			)
						lowerCamelCase__	:	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

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

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

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

						lowerCamelCase__	:	Optional[Any]				=check_and_map_params(
						    self_attn.key.bias.data							,	f'''encoder.transformer_cells.{i}.attention_cell.proj_key.bias'''			)

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

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

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

						# intermediate
						lowerCamelCase__	:	BertIntermediate				=layer.intermediate

						lowerCamelCase__	:	Union[str, Any]				=check_and_map_params(
						    intermediate.dense.bias							,	f'''encoder.transformer_cells.{i}.ffn.ffn_1.bias'''			)
						lowerCamelCase__	:	Any				=check_and_map_params(
						    intermediate.dense.weight							,	f'''encoder.transformer_cells.{i}.ffn.ffn_1.weight'''			)

						# output
						lowerCamelCase__	:	BertOutput				=layer.output

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

			lowerCamelCase__	:	Tuple				=tokenizer.encode_plus(__lowerCamelCase			)['''input_ids''']

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

			# Get Transformer output (save and reload model again)
			hf_bort_model.save_pretrained(__lowerCamelCase			)
			lowerCamelCase__	:	int				=BertModel.from_pretrained(__lowerCamelCase			)
			hf_bort_model.eval()

			lowerCamelCase__	:	Union[str, Any]				=tokenizer.encode_plus(__lowerCamelCase							,	return_tensors='''pt'''			)
			lowerCamelCase__	:	Any				=hf_bort_model(**__lowerCamelCase			)[0]

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

			lowerCamelCase__	:	List[Any]				=np.max(np.abs(hf_layer - gluon_layer			)			).item()
			lowerCamelCase__	:	List[Any]				=np.allclose(__lowerCamelCase							,	__lowerCamelCase							,	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:'''							,	__lowerCamelCase			)


if __name__ == "__main__":
		_lowercase				:				List[str]   		=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."
		)
		_lowercase				:				List[str]   		=parser.parse_args()
		convert_bort_checkpoint_to_pytorch(args.bort_checkpoint_path, args.pytorch_dump_folder_path)









"""simple docstring"""






_lowercase				:				Optional[Any]   		=		{
    "Pillow": "Pillow<10.0.0",
    "accelerate": "accelerate>=0.20.3",
    "av": "av==9.2.0",
    "beautifulsoup4": "beautifulsoup4",
    "black": "black~=23.1",
    "codecarbon": "codecarbon==1.2.0",
    "cookiecutter": "cookiecutter==1.7.3",
    "dataclasses": "dataclasses",
    "datasets": "datasets!=2.5.0",
    "decord": "decord==0.6.0",
    "deepspeed": "deepspeed>=0.9.3",
    "diffusers": "diffusers",
    "dill": "dill<0.3.5",
    "evaluate": "evaluate>=0.2.0",
    "fairscale": "fairscale>0.3",
    "faiss-cpu": "faiss-cpu",
    "fastapi": "fastapi",
    "filelock": "filelock",
    "flax": "flax>=0.4.1,<=0.7.0",
    "ftfy": "ftfy",
    "fugashi": "fugashi>=1.0",
    "GitPython": "GitPython<3.1.19",
    "hf-doc-builder": "hf-doc-builder>=0.3.0",
    "huggingface-hub": "huggingface-hub>=0.14.1,<1.0",
    "importlib_metadata": "importlib_metadata",
    "ipadic": "ipadic>=1.0.0,<2.0",
    "isort": "isort>=5.5.4",
    "jax": "jax>=0.2.8,!=0.3.2,<=0.4.13",
    "jaxlib": "jaxlib>=0.1.65,<=0.4.13",
    "jieba": "jieba",
    "kenlm": "kenlm",
    "keras-nlp": "keras-nlp>=0.3.1",
    "librosa": "librosa",
    "nltk": "nltk",
    "natten": "natten>=0.14.6",
    "numpy": "numpy>=1.17",
    "onnxconverter-common": "onnxconverter-common",
    "onnxruntime-tools": "onnxruntime-tools>=1.4.2",
    "onnxruntime": "onnxruntime>=1.4.0",
    "opencv-python": "opencv-python",
    "optuna": "optuna",
    "optax": "optax>=0.0.8,<=0.1.4",
    "packaging": "packaging>=20.0",
    "parameterized": "parameterized",
    "phonemizer": "phonemizer",
    "protobuf": "protobuf",
    "psutil": "psutil",
    "pyyaml": "pyyaml>=5.1",
    "pydantic": "pydantic<2",
    "pytest": "pytest>=7.2.0",
    "pytest-timeout": "pytest-timeout",
    "pytest-xdist": "pytest-xdist",
    "python": "python>=3.8.0",
    "ray[tune]": "ray[tune]",
    "regex": "regex!=2019.12.17",
    "requests": "requests",
    "rhoknp": "rhoknp>=1.1.0,<1.3.1",
    "rjieba": "rjieba",
    "rouge-score": "rouge-score!=0.0.7,!=0.0.8,!=0.1,!=0.1.1",
    "ruff": "ruff>=0.0.241,<=0.0.259",
    "sacrebleu": "sacrebleu>=1.4.12,<2.0.0",
    "sacremoses": "sacremoses",
    "safetensors": "safetensors>=0.3.1",
    "sagemaker": "sagemaker>=2.31.0",
    "scikit-learn": "scikit-learn",
    "sentencepiece": "sentencepiece>=0.1.91,!=0.1.92",
    "sigopt": "sigopt",
    "starlette": "starlette",
    "sudachipy": "sudachipy>=0.6.6",
    "sudachidict_core": "sudachidict_core>=20220729",
    "tensorflow-cpu": "tensorflow-cpu>=2.6,<2.14",
    "tensorflow": "tensorflow>=2.6,<2.14",
    "tensorflow-text": "tensorflow-text<2.14",
    "tf2onnx": "tf2onnx",
    "timeout-decorator": "timeout-decorator",
    "timm": "timm",
    "tokenizers": "tokenizers>=0.11.1,!=0.11.3,<0.14",
    "torch": "torch>=1.9,!=1.12.0",
    "torchaudio": "torchaudio",
    "torchvision": "torchvision",
    "pyctcdecode": "pyctcdecode>=0.4.0",
    "tqdm": "tqdm>=4.27",
    "unidic": "unidic>=1.0.2",
    "unidic_lite": "unidic_lite>=1.0.7",
    "urllib3": "urllib3<2.0.0",
    "uvicorn": "uvicorn",
}




