




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()
_UpperCAmelCase	:						int			=     logging.get_logger(__name__)

_UpperCAmelCase	:						Tuple			=     "The Nymphenburg Palace is a beautiful palace in Munich!"



def       UpperCAmelCase__  (     lowerCamelCase,			lowerCamelCase   ):
  lowercase    :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,
  }

  lowercase    :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
  lowercase    :List[str]      	=						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
  lowercase    :List[str]      	=						"openwebtext_ccnews_stories_books_cased"

  # Specify download folder to Gluonnlp's vocab
  lowercase    :Dict      	=						os.path.join(get_home_dir(),			"models"   )
  lowercase    :Any      	=						_load_vocab(lowerCamelCase,			lowerCamelCase,			lowerCamelCase,			cls=lowerCamelCase   )

  lowercase    :Union[str, 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   )
  lowercase    :Union[str, Any]      	=						original_bort._collect_params_with_prefix()

  # Build our config 🤗
  lowercase    :Optional[int]      	=						{
      "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   ),
  }

  lowercase    :Any      	=						BertConfig.from_dict(lowerCamelCase   )
  lowercase    :Tuple      	=						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   ) -> 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,			lowerCamelCase   ):
    lowercase    :List[Any]      	=						hf_param.shape

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

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

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

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

    lowercase    :Optional[Any]      	=						check_and_map_params(
        self_attn.key.bias.data,			F"encoder.transformer_cells.{i}.attention_cell.proj_key.bias"   )

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

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

    lowercase    :Union[str, Any]      	=						check_and_map_params(
        self_output.dense.bias,			F"encoder.transformer_cells.{i}.proj.bias"   )
    lowercase    :Any      	=						check_and_map_params(
        self_output.dense.weight,			F"encoder.transformer_cells.{i}.proj.weight"   )
    lowercase    :Union[str, Any]      	=						check_and_map_params(
        self_output.LayerNorm.bias,			F"encoder.transformer_cells.{i}.layer_norm.beta"   )
    lowercase    :Dict      	=						check_and_map_params(
        self_output.LayerNorm.weight,			F"encoder.transformer_cells.{i}.layer_norm.gamma"   )

    # intermediate
    lowercase    :BertIntermediate      	=						layer.intermediate

    lowercase    :str      	=						check_and_map_params(
        intermediate.dense.bias,			F"encoder.transformer_cells.{i}.ffn.ffn_1.bias"   )
    lowercase    :Tuple      	=						check_and_map_params(
        intermediate.dense.weight,			F"encoder.transformer_cells.{i}.ffn.ffn_1.weight"   )

    # output
    lowercase    :BertOutput      	=						layer.output

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

  lowercase    :List[Any]      	=						tokenizer.encode_plus(lowerCamelCase   )["input_ids"]

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

  # Get Transformer output (save and reload model again)
  hf_bort_model.save_pretrained(lowerCamelCase   )
  lowercase    :Any      	=						BertModel.from_pretrained(lowerCamelCase   )
  hf_bort_model.eval()

  lowercase    :str      	=						tokenizer.encode_plus(lowerCamelCase,			return_tensors="pt"   )
  lowercase    :Union[str, Any]      	=						hf_bort_model(**lowerCamelCase   )[0]

  lowercase    :Dict      	=						output_gluon[0].asnumpy()
  lowercase    :List[str]      	=						output_hf[0].detach().numpy()

  lowercase    :Dict      	=						np.max(np.abs(hf_layer - gluon_layer   )   ).item()
  lowercase    :Optional[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__":
  _UpperCAmelCase	:						Union[str, 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."
  )
  _UpperCAmelCase	:						Dict			=     parser.parse_args()
  convert_bort_checkpoint_to_pytorch(args.bort_checkpoint_path, args.pytorch_dump_folder_path)









from collections.abc import Generator
from math import sin



def       UpperCAmelCase__  (     lowerCamelCase   ):
  if len(lowerCamelCase   ) != 32:
    raise ValueError("Input must be of length 32"   )

  lowercase    :Any      	=						B""
  for i in [3, 2, 1, 0]:
    little_endian += string_aa[8 * i : 8 * i + 8]
  return little_endian



def       UpperCAmelCase__  (     lowerCamelCase   ):
  if i < 0:
    raise ValueError("Input must be non-negative"   )

  lowercase    :List[Any]      	=						format(lowerCamelCase,			"08x"   )[-8:]
  lowercase    :Tuple      	=						B""
  for i in [3, 2, 1, 0]:
    little_endian_hex += hex_rep[2 * i : 2 * i + 2].encode("utf-8"   )
  return little_endian_hex



def       UpperCAmelCase__  (     lowerCamelCase   ):
  lowercase    :Union[str, Any]      	=						B""
  for char in message:
    bit_string += format(lowerCamelCase,			"08b"   ).encode("utf-8"   )
  lowercase    :str      	=						format(len(lowerCamelCase   ),			"064b"   ).encode("utf-8"   )

  # Pad bit_string to a multiple of 512 chars
  bit_string += b"1"
  while len(lowerCamelCase   ) % 512 != 448:
    bit_string += b"0"
  bit_string += to_little_endian(start_len[32:]   ) + to_little_endian(start_len[:32]   )

  return bit_string



def       UpperCAmelCase__  (     lowerCamelCase   ):
  if len(lowerCamelCase   ) % 512 != 0:
    raise ValueError("Input must have length that's a multiple of 512"   )

  for pos in range(0,			len(lowerCamelCase   ),			512   ):
    lowercase    :List[str]      	=						bit_string[pos : pos + 512]
    lowercase    :int      	=						[]
    for i in range(0,			512,			32   ):
      block_words.append(int(to_little_endian(block[i : i + 32]   ),			2   )   )
    yield block_words



def       UpperCAmelCase__  (     lowerCamelCase   ):
  if i < 0:
    raise ValueError("Input must be non-negative"   )

  lowercase    :List[str]      	=						format(lowerCamelCase,			"032b"   )
  lowercase    :Optional[Any]      	=						""
  for c in i_str:
    new_str += "1" if c == "0" else "0"
  return int(lowerCamelCase,			2   )



def       UpperCAmelCase__  (     lowerCamelCase,			lowerCamelCase   ):
  return (a + b) % 2**32



def       UpperCAmelCase__  (     lowerCamelCase,			lowerCamelCase   ):
  if i < 0:
    raise ValueError("Input must be non-negative"   )
  if shift < 0:
    raise ValueError("Shift must be non-negative"   )
  return ((i << shift) ^ (i >> (32 - shift))) % 2**32



def       UpperCAmelCase__  (     lowerCamelCase   ):
  lowercase    :Tuple      	=						preprocess(lowerCamelCase   )

  lowercase    :Tuple      	=						[int(2**32 * abs(sin(i + 1   )   )   ) for i in range(64   )]

  # Starting states
  lowercase    :Optional[Any]      	=						0X6745_2301
  lowercase    :Union[str, Any]      	=						0XEFCD_AB89
  lowercase    :Optional[int]      	=						0X98BA_DCFE
  lowercase    :str      	=						0X1032_5476

  lowercase    :str      	=						[
      7,
      12,
      17,
      22,
      7,
      12,
      17,
      22,
      7,
      12,
      17,
      22,
      7,
      12,
      17,
      22,
      5,
      9,
      14,
      20,
      5,
      9,
      14,
      20,
      5,
      9,
      14,
      20,
      5,
      9,
      14,
      20,
      4,
      11,
      16,
      23,
      4,
      11,
      16,
      23,
      4,
      11,
      16,
      23,
      4,
      11,
      16,
      23,
      6,
      10,
      15,
      21,
      6,
      10,
      15,
      21,
      6,
      10,
      15,
      21,
      6,
      10,
      15,
      21,
  ]

  # Process bit string in chunks, each with 16 32-char words
  for block_words in get_block_words(lowerCamelCase   ):
    lowercase    :str      	=						aa
    lowercase    :int      	=						ba
    lowercase    :int      	=						ca
    lowercase    :List[str]      	=						da

    # Hash current chunk
    for i in range(64   ):
      if i <= 15:
        # f = (b & c) | (not_32(b) & d)     # Alternate definition for f
        lowercase    :Optional[int]      	=						d ^ (b & (c ^ d))
        lowercase    :Optional[Any]      	=						i
      elif i <= 31:
        # f = (d & b) | (not_32(d) & c)     # Alternate definition for f
        lowercase    :Any      	=						c ^ (d & (b ^ c))
        lowercase    :Union[str, Any]      	=						(5 * i + 1) % 16
      elif i <= 47:
        lowercase    :Any      	=						b ^ c ^ d
        lowercase    :List[Any]      	=						(3 * i + 5) % 16
      else:
        lowercase    :List[Any]      	=						c ^ (b | not_aa(lowerCamelCase   ))
        lowercase    :Any      	=						(7 * i) % 16
      lowercase    :Dict      	=						(f + a + added_consts[i] + block_words[g]) % 2**32
      lowercase    :Any      	=						d
      lowercase    :str      	=						c
      lowercase    :Union[str, Any]      	=						b
      lowercase    :Union[str, Any]      	=						sum_aa(lowerCamelCase,			left_rotate_aa(lowerCamelCase,			shift_amounts[i]   )   )

    # Add hashed chunk to running total
    lowercase    :Optional[Any]      	=						sum_aa(lowerCamelCase,			lowerCamelCase   )
    lowercase    :List[Any]      	=						sum_aa(lowerCamelCase,			lowerCamelCase   )
    lowercase    :List[str]      	=						sum_aa(lowerCamelCase,			lowerCamelCase   )
    lowercase    :Union[str, Any]      	=						sum_aa(lowerCamelCase,			lowerCamelCase   )

  lowercase    :int      	=						reformat_hex(lowerCamelCase   ) + reformat_hex(lowerCamelCase   ) + reformat_hex(lowerCamelCase   ) + reformat_hex(lowerCamelCase   )
  return digest


if __name__ == "__main__":
     import doctest

     doctest.testmod()



