






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			:					Optional[Any]         =	logging.get_logger(__name__)

lowerCamelCase			:					List[Any]         =	'''The Nymphenburg Palace is a beautiful palace in Munich!'''





def       __lowerCAmelCase      (						__snake_case					, __snake_case       ):
     __lowerCAmelCase								=					{
         "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								=					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								=					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=__snake_case					, output_all_encodings=__snake_case					, use_residual=predefined_args["use_residual"]					, activation=predefined_args.get("activation"					, "gelu"       )					, layer_norm_eps=predefined_args.get("layer_norm_eps"					, __snake_case       )					, )

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

     # Specify download folder to Gluonnlp's vocab
     __lowerCAmelCase								=					os.path.join(get_home_dir()					, "models"       )
     __lowerCAmelCase								=					_load_vocab(__snake_case					, __snake_case					, __snake_case					, cls=__snake_case       )

     __lowerCAmelCase								=					nlp.model.BERTModel(
         __snake_case					, len(__snake_case       )					, units=predefined_args["units"]					, embed_size=predefined_args["embed_size"]					, embed_dropout=predefined_args["embed_dropout"]					, word_embed=predefined_args["word_embed"]					, use_pooler=__snake_case					, use_token_type_embed=__snake_case					, token_type_vocab_size=predefined_args["token_type_vocab_size"]					, use_classifier=__snake_case					, use_decoder=__snake_case					, )

     original_bort.load_parameters(__snake_case					, cast_dtype=__snake_case					, ignore_extra=__snake_case       )
     __lowerCAmelCase								=					original_bort._collect_params_with_prefix()

     # Build our config 🤗
     __lowerCAmelCase								=					{
         "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(__snake_case       ),
     }

     __lowerCAmelCase								=					BertConfig.from_dict(__snake_case       )
     __lowerCAmelCase								=					BertForMaskedLM(__snake_case       )
     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(__snake_case       ) -> 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(__snake_case					, __snake_case       ):
          __lowerCAmelCase								=					hf_param.shape

          __lowerCAmelCase								=					to_torch(params[gluon_param]       )
          __lowerCAmelCase								=					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								=					check_and_map_params(
         hf_bort_model.bert.embeddings.word_embeddings.weight					, "word_embed.0.weight"       )
     __lowerCAmelCase								=					check_and_map_params(
         hf_bort_model.bert.embeddings.position_embeddings.weight					, "encoder.position_weight"       )
     __lowerCAmelCase								=					check_and_map_params(
         hf_bort_model.bert.embeddings.LayerNorm.bias					, "encoder.layer_norm.beta"       )
     __lowerCAmelCase								=					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								=					torch.zeros_like(
         hf_bort_model.bert.embeddings.token_type_embeddings.weight.data       )

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

          # self attention
          __lowerCAmelCase								=					layer.attention.self

          __lowerCAmelCase								=					check_and_map_params(
              self_attn.key.bias.data					, F"""encoder.transformer_cells.{i}.attention_cell.proj_key.bias"""       )

          __lowerCAmelCase								=					check_and_map_params(
              self_attn.key.weight.data					, F"""encoder.transformer_cells.{i}.attention_cell.proj_key.weight"""       )
          __lowerCAmelCase								=					check_and_map_params(
              self_attn.query.bias.data					, F"""encoder.transformer_cells.{i}.attention_cell.proj_query.bias"""       )
          __lowerCAmelCase								=					check_and_map_params(
              self_attn.query.weight.data					, F"""encoder.transformer_cells.{i}.attention_cell.proj_query.weight"""       )
          __lowerCAmelCase								=					check_and_map_params(
              self_attn.value.bias.data					, F"""encoder.transformer_cells.{i}.attention_cell.proj_value.bias"""       )
          __lowerCAmelCase								=					check_and_map_params(
              self_attn.value.weight.data					, F"""encoder.transformer_cells.{i}.attention_cell.proj_value.weight"""       )

          # self attention output
          __lowerCAmelCase								=					layer.attention.output

          __lowerCAmelCase								=					check_and_map_params(
              self_output.dense.bias					, F"""encoder.transformer_cells.{i}.proj.bias"""       )
          __lowerCAmelCase								=					check_and_map_params(
              self_output.dense.weight					, F"""encoder.transformer_cells.{i}.proj.weight"""       )
          __lowerCAmelCase								=					check_and_map_params(
              self_output.LayerNorm.bias					, F"""encoder.transformer_cells.{i}.layer_norm.beta"""       )
          __lowerCAmelCase								=					check_and_map_params(
              self_output.LayerNorm.weight					, F"""encoder.transformer_cells.{i}.layer_norm.gamma"""       )

          # intermediate
          __lowerCAmelCase								=					layer.intermediate

          __lowerCAmelCase								=					check_and_map_params(
              intermediate.dense.bias					, F"""encoder.transformer_cells.{i}.ffn.ffn_1.bias"""       )
          __lowerCAmelCase								=					check_and_map_params(
              intermediate.dense.weight					, F"""encoder.transformer_cells.{i}.ffn.ffn_1.weight"""       )

          # output
          __lowerCAmelCase								=					layer.output

          __lowerCAmelCase								=					check_and_map_params(
              bert_output.dense.bias					, F"""encoder.transformer_cells.{i}.ffn.ffn_2.bias"""       )
          __lowerCAmelCase								=					check_and_map_params(
              bert_output.dense.weight					, F"""encoder.transformer_cells.{i}.ffn.ffn_2.weight"""       )
          __lowerCAmelCase								=					check_and_map_params(
              bert_output.LayerNorm.bias					, F"""encoder.transformer_cells.{i}.ffn.layer_norm.beta"""       )
          __lowerCAmelCase								=					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								=					RobertaTokenizer.from_pretrained("roberta-base"       )

     __lowerCAmelCase								=					tokenizer.encode_plus(__snake_case       )["input_ids"]

     # Get gluon output
     __lowerCAmelCase								=					mx.nd.array([input_ids]       )
     __lowerCAmelCase								=					original_bort(inputs=__snake_case					, token_types=[]       )

     # Get Transformer output (save and reload model again)
     hf_bort_model.save_pretrained(__snake_case       )
     __lowerCAmelCase								=					BertModel.from_pretrained(__snake_case       )
     hf_bort_model.eval()

     __lowerCAmelCase								=					tokenizer.encode_plus(__snake_case					, return_tensors="pt"       )
     __lowerCAmelCase								=					hf_bort_model(**__snake_case       )[0]

     __lowerCAmelCase								=					output_gluon[0].asnumpy()
     __lowerCAmelCase								=					output_hf[0].detach().numpy()

     __lowerCAmelCase								=					np.max(np.abs(hf_layer - gluon_layer       )       ).item()
     __lowerCAmelCase								=					np.allclose(__snake_case					, __snake_case					, 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:"					, __snake_case       )


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














# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path

import torch

from ...utils import is_npu_available, is_xpu_available
from .config_args import ClusterConfig, default_json_config_file
from .config_utils import SubcommandHelpFormatter


lowerCamelCase			:					List[str]         =	'''Create a default config file for Accelerate with only a few flags set.'''





def       __lowerCAmelCase      (						__snake_case="no"					, __snake_case = default_json_config_file					, __snake_case = False       ):
     __lowerCAmelCase								=					Path(__snake_case       )
     path.parent.mkdir(parents=__snake_case					, exist_ok=__snake_case       )
     if path.exists():
          print(
              F"""Configuration already exists at {save_location}, will not override. Run `accelerate config` manually or pass a different `save_location`."""       )
          return False
     __lowerCAmelCase								=					mixed_precision.lower()
     if mixed_precision not in ["no", "fp16", "bf16", "fp8"]:
          raise ValueError(
              F"""`mixed_precision` should be one of 'no', 'fp16', 'bf16', or 'fp8'. Received {mixed_precision}"""       )
     __lowerCAmelCase								=					{
         "compute_environment": "LOCAL_MACHINE",
         "mixed_precision": mixed_precision,
     }
     if torch.cuda.is_available():
          __lowerCAmelCase								=					torch.cuda.device_count()
          __lowerCAmelCase								=					num_gpus
          __lowerCAmelCase								=					False
          if num_gpus > 1:
               __lowerCAmelCase								=					"MULTI_GPU"
          else:
               __lowerCAmelCase								=					"NO"
     elif is_xpu_available() and use_xpu:
          __lowerCAmelCase								=					torch.xpu.device_count()
          __lowerCAmelCase								=					num_xpus
          __lowerCAmelCase								=					False
          if num_xpus > 1:
               __lowerCAmelCase								=					"MULTI_XPU"
          else:
               __lowerCAmelCase								=					"NO"
     elif is_npu_available():
          __lowerCAmelCase								=					torch.npu.device_count()
          __lowerCAmelCase								=					num_npus
          __lowerCAmelCase								=					False
          if num_npus > 1:
               __lowerCAmelCase								=					"MULTI_NPU"
          else:
               __lowerCAmelCase								=					"NO"
     else:
          __lowerCAmelCase								=					0
          __lowerCAmelCase								=					True
          __lowerCAmelCase								=					1
          __lowerCAmelCase								=					"NO"
     __lowerCAmelCase								=					ClusterConfig(**__snake_case       )
     config.to_json_file(__snake_case       )
     return path





def       __lowerCAmelCase      (						__snake_case					, __snake_case       ):
     __lowerCAmelCase								=					parser.add_parser("default"					, parents=__snake_case					, help=__snake_case					, formatter_class=__snake_case       )
     parser.add_argument(
         "--config_file"					, default=__snake_case					, help=(
             "The path to use to store the config file. Will default to a file named default_config.yaml in the cache "
             "location, which is the content of the environment `HF_HOME` suffixed with 'accelerate', or if you don't have "
             "such an environment variable, your cache directory ('~/.cache' or the content of `XDG_CACHE_HOME`) suffixed "
             "with 'huggingface'."
         )					, dest="save_location"					, )

     parser.add_argument(
         "--mixed_precision"					, choices=["no", "fp16", "bf16"]					, type=__snake_case					, help="Whether or not to use mixed precision training. "
         "Choose between FP16 and BF16 (bfloat16) training. "
         "BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later."					, default="no"					, )
     parser.set_defaults(func=__snake_case       )
     return parser





def       __lowerCAmelCase      (						__snake_case       ):
     __lowerCAmelCase								=					write_basic_config(args.mixed_precision					, args.save_location       )
     if config_file:
          print(F"""accelerate configuration saved at {config_file}"""       )






