





import gc
import random
import tempfile
import unittest

import numpy as np
import torch
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, LMSDiscreteScheduler, PNDMScheduler, UNetaDConditionModel
from diffusers.pipelines.stable_diffusion_safe import StableDiffusionPipelineSafe as StableDiffusionPipeline
from diffusers.utils import floats_tensor, nightly, torch_device
from diffusers.utils.testing_utils import require_torch_gpu

class     snake_case_		(   unittest.TestCase							):






 '''simple docstring'''





 def  __UpperCAmelCase						(  self		)	->       Union[str, Any]:
    # clean up the VRAM after each test
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()

 @property
 def  __UpperCAmelCase						(  self		)	->       Optional[Any]:
    UpperCAmelCase__						=1
    UpperCAmelCase__						=3
    UpperCAmelCase__						=(32, 32)

    UpperCAmelCase__						=floats_tensor((batch_size, num_channels) + sizes,		rng=random.Random(0		)		).to(A_		)
    return image

 @property
 def  __UpperCAmelCase						(  self		)	->       Optional[Any]:
    torch.manual_seed(0		)
    UpperCAmelCase__						=UNetaDConditionModel(
        block_out_channels=(32, 64),		layers_per_block=2,		sample_size=32,		in_channels=4,		out_channels=4,		down_block_types=("DownBlock2D", "CrossAttnDownBlock2D"),		up_block_types=("CrossAttnUpBlock2D", "UpBlock2D"),		cross_attention_dim=32,		)
    return model

 @property
 def  __UpperCAmelCase						(  self		)	->       Dict:
    torch.manual_seed(0		)
    UpperCAmelCase__						=AutoencoderKL(
        block_out_channels=[32, 64],		in_channels=3,		out_channels=3,		down_block_types=["DownEncoderBlock2D", "DownEncoderBlock2D"],		up_block_types=["UpDecoderBlock2D", "UpDecoderBlock2D"],		latent_channels=4,		)
    return model

 @property
 def  __UpperCAmelCase						(  self		)	->       Optional[Any]:
    torch.manual_seed(0		)
    UpperCAmelCase__						=CLIPTextConfig(
        bos_token_id=0,		eos_token_id=2,		hidden_size=32,		intermediate_size=37,		layer_norm_eps=1E-05,		num_attention_heads=4,		num_hidden_layers=5,		pad_token_id=1,		vocab_size=1000,		)
    return CLIPTextModel(A_		)

 @property
 def  __UpperCAmelCase						(  self		)	->       str:
    def extract(*A_,		**A_		):

       class     snake_case_		:






        '''simple docstring'''





        def __init__(  self		)	->       Dict:
           UpperCAmelCase__						=torch.ones([0]		)






        def  __UpperCAmelCase						(  self,		A_		)	->       str:
           self.pixel_values.to(A_		)
           return self

       return Out()

    return extract

 def  __UpperCAmelCase						(  self		)	->       Any:
    UpperCAmelCase__						="cpu"  # ensure determinism for the device-dependent torch.Generator
    UpperCAmelCase__						=self.dummy_cond_unet
    UpperCAmelCase__						=DDIMScheduler(
        beta_start=0.0_00_85,		beta_end=0.0_12,		beta_schedule="scaled_linear",		clip_sample=A_,		set_alpha_to_one=A_,		)

    UpperCAmelCase__						=self.dummy_vae
    UpperCAmelCase__						=self.dummy_text_encoder
    UpperCAmelCase__						=CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)

    # make sure here that pndm scheduler skips prk
    UpperCAmelCase__						=StableDiffusionPipeline(
        unet=A_,		scheduler=A_,		vae=A_,		text_encoder=A_,		tokenizer=A_,		safety_checker=A_,		feature_extractor=self.dummy_extractor,		)
    UpperCAmelCase__						=sd_pipe.to(A_		)
    sd_pipe.set_progress_bar_config(disable=A_		)

    UpperCAmelCase__						="A painting of a squirrel eating a burger"

    UpperCAmelCase__						=torch.Generator(device=A_		).manual_seed(0		)
    UpperCAmelCase__						=sd_pipe([prompt],		generator=A_,		guidance_scale=6.0,		num_inference_steps=2,		output_type="np"		)
    UpperCAmelCase__						=output.images

    UpperCAmelCase__						=torch.Generator(device=A_		).manual_seed(0		)
    UpperCAmelCase__						=sd_pipe(
        [prompt],		generator=A_,		guidance_scale=6.0,		num_inference_steps=2,		output_type="np",		return_dict=A_,		)[0]

    UpperCAmelCase__						=image[0, -3:, -3:, -1]
    UpperCAmelCase__						=image_from_tuple[0, -3:, -3:, -1]

    assert image.shape == (1, 64, 64, 3)
    UpperCAmelCase__						=np.array([0.57_56, 0.61_18, 0.50_05, 0.50_41, 0.54_71, 0.47_26, 0.49_76, 0.48_65, 0.48_64]		)

    assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2
    assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() < 1E-2

 def  __UpperCAmelCase						(  self		)	->       int:
    UpperCAmelCase__						="cpu"  # ensure determinism for the device-dependent torch.Generator
    UpperCAmelCase__						=self.dummy_cond_unet
    UpperCAmelCase__						=PNDMScheduler(skip_prk_steps=A_		)
    UpperCAmelCase__						=self.dummy_vae
    UpperCAmelCase__						=self.dummy_text_encoder
    UpperCAmelCase__						=CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)

    # make sure here that pndm scheduler skips prk
    UpperCAmelCase__						=StableDiffusionPipeline(
        unet=A_,		scheduler=A_,		vae=A_,		text_encoder=A_,		tokenizer=A_,		safety_checker=A_,		feature_extractor=self.dummy_extractor,		)
    UpperCAmelCase__						=sd_pipe.to(A_		)
    sd_pipe.set_progress_bar_config(disable=A_		)

    UpperCAmelCase__						="A painting of a squirrel eating a burger"
    UpperCAmelCase__						=torch.Generator(device=A_		).manual_seed(0		)
    UpperCAmelCase__						=sd_pipe([prompt],		generator=A_,		guidance_scale=6.0,		num_inference_steps=2,		output_type="np"		)

    UpperCAmelCase__						=output.images

    UpperCAmelCase__						=torch.Generator(device=A_		).manual_seed(0		)
    UpperCAmelCase__						=sd_pipe(
        [prompt],		generator=A_,		guidance_scale=6.0,		num_inference_steps=2,		output_type="np",		return_dict=A_,		)[0]

    UpperCAmelCase__						=image[0, -3:, -3:, -1]
    UpperCAmelCase__						=image_from_tuple[0, -3:, -3:, -1]

    assert image.shape == (1, 64, 64, 3)
    UpperCAmelCase__						=np.array([0.51_25, 0.57_16, 0.48_28, 0.50_60, 0.56_50, 0.47_68, 0.51_85, 0.48_95, 0.49_93]		)

    assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2
    assert np.abs(image_from_tuple_slice.flatten() - expected_slice		).max() < 1E-2

 def  __UpperCAmelCase						(  self		)	->       str:
    UpperCAmelCase__						=StableDiffusionPipeline.from_pretrained(
        "hf-internal-testing/tiny-stable-diffusion-lms-pipe",		safety_checker=A_		)
    assert isinstance(A_,		A_		)
    assert isinstance(pipe.scheduler,		A_		)
    assert pipe.safety_checker is None

    UpperCAmelCase__						=pipe("example prompt",		num_inference_steps=2		).images[0]
    assert image is not None

    # check that there's no error when saving a pipeline with one of the models being None
    with tempfile.TemporaryDirectory() as tmpdirname:
       pipe.save_pretrained(A_		)
       UpperCAmelCase__						=StableDiffusionPipeline.from_pretrained(A_		)

    # sanity check that the pipeline still works
    assert pipe.safety_checker is None
    UpperCAmelCase__						=pipe("example prompt",		num_inference_steps=2		).images[0]
    assert image is not None






 @unittest.skipIf(torch_device != "cuda",		"This test requires a GPU"		)
 def  __UpperCAmelCase						(  self		)	->       Optional[int]:
    UpperCAmelCase__						=self.dummy_cond_unet
    UpperCAmelCase__						=PNDMScheduler(skip_prk_steps=A_		)
    UpperCAmelCase__						=self.dummy_vae
    UpperCAmelCase__						=self.dummy_text_encoder
    UpperCAmelCase__						=CLIPTokenizer.from_pretrained("hf-internal-testing/tiny-random-clip"		)

    # put models in fp16
    UpperCAmelCase__						=unet.half()
    UpperCAmelCase__						=vae.half()
    UpperCAmelCase__						=bert.half()

    # make sure here that pndm scheduler skips prk
    UpperCAmelCase__						=StableDiffusionPipeline(
        unet=A_,		scheduler=A_,		vae=A_,		text_encoder=A_,		tokenizer=A_,		safety_checker=A_,		feature_extractor=self.dummy_extractor,		)
    UpperCAmelCase__						=sd_pipe.to(A_		)
    sd_pipe.set_progress_bar_config(disable=A_		)

    UpperCAmelCase__						="A painting of a squirrel eating a burger"
    UpperCAmelCase__						=sd_pipe([prompt],		num_inference_steps=2,		output_type="np"		).images

    assert image.shape == (1, 64, 64, 3)

@nightly
@require_torch_gpu
class     snake_case_		(   unittest.TestCase							):






 '''simple docstring'''





 def  __UpperCAmelCase						(  self		)	->       Union[str, Any]:
    # clean up the VRAM after each test
    super().tearDown()
    gc.collect()
    torch.cuda.empty_cache()

 def  __UpperCAmelCase						(  self		)	->       Optional[int]:
    UpperCAmelCase__						=StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5",		safety_checker=A_		)
    UpperCAmelCase__						=LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config		)
    UpperCAmelCase__						=sd_pipe.to(A_		)
    sd_pipe.set_progress_bar_config(disable=A_		)

    UpperCAmelCase__						=(
        "portrait of girl with smokey eyes makeup in abandoned hotel, grange clothes, redshift, wide high angle"
        " coloured polaroid photograph with flash, kodak film, hyper real, stunning moody cinematography, with"
        " anamorphic lenses, by maripol, fallen angels by wong kar - wai, style of suspiria and neon demon and"
        " children from bahnhof zoo, detailed "
    )
    UpperCAmelCase__						=40_0366_0346
    UpperCAmelCase__						=7

    # without safety guidance (sld_guidance_scale = 0)
    UpperCAmelCase__						=torch.manual_seed(A_		)
    UpperCAmelCase__						=sd_pipe(
        [prompt],		generator=A_,		guidance_scale=A_,		num_inference_steps=50,		output_type="np",		width=512,		height=512,		sld_guidance_scale=0,		)

    UpperCAmelCase__						=output.images
    UpperCAmelCase__						=image[0, -3:, -3:, -1]
    UpperCAmelCase__						=[0.22_78, 0.22_31, 0.22_49, 0.23_33, 0.23_03, 0.18_85, 0.22_73, 0.21_44, 0.21_76]

    assert image.shape == (1, 512, 512, 3)

    assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2

    # without safety guidance (strong configuration)
    UpperCAmelCase__						=torch.manual_seed(A_		)
    UpperCAmelCase__						=sd_pipe(
        [prompt],		generator=A_,		guidance_scale=A_,		num_inference_steps=50,		output_type="np",		width=512,		height=512,		sld_guidance_scale=2000,		sld_warmup_steps=7,		sld_threshold=0.0_25,		sld_momentum_scale=0.5,		sld_mom_beta=0.7,		)

    UpperCAmelCase__						=output.images
    UpperCAmelCase__						=image[0, -3:, -3:, -1]
    UpperCAmelCase__						=[0.23_83, 0.22_76, 0.2_36, 0.21_92, 0.21_86, 0.20_53, 0.19_71, 0.19_01, 0.17_19]

    assert image.shape == (1, 512, 512, 3)

    assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2

 def  __UpperCAmelCase						(  self		)	->       int:
    UpperCAmelCase__						=StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5",		safety_checker=A_		)
    UpperCAmelCase__						=LMSDiscreteScheduler.from_config(sd_pipe.scheduler.config		)
    UpperCAmelCase__						=sd_pipe.to(A_		)
    sd_pipe.set_progress_bar_config(disable=A_		)

    UpperCAmelCase__						="padme amidala taking a bath artwork, safe for work, no nudity"
    UpperCAmelCase__						=27_3497_1755
    UpperCAmelCase__						=7

    UpperCAmelCase__						=torch.manual_seed(A_		)
    UpperCAmelCase__						=sd_pipe(
        [prompt],		generator=A_,		guidance_scale=A_,		num_inference_steps=50,		output_type="np",		width=512,		height=512,		sld_guidance_scale=0,		)

    UpperCAmelCase__						=output.images
    UpperCAmelCase__						=image[0, -3:, -3:, -1]
    UpperCAmelCase__						=[0.35_02, 0.36_22, 0.33_96, 0.36_42, 0.34_78, 0.33_18, 0.35, 0.33_48, 0.32_97]

    assert image.shape == (1, 512, 512, 3)

    assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2

    UpperCAmelCase__						=torch.manual_seed(A_		)
    UpperCAmelCase__						=sd_pipe(
        [prompt],		generator=A_,		guidance_scale=A_,		num_inference_steps=50,		output_type="np",		width=512,		height=512,		sld_guidance_scale=2000,		sld_warmup_steps=7,		sld_threshold=0.0_25,		sld_momentum_scale=0.5,		sld_mom_beta=0.7,		)

    UpperCAmelCase__						=output.images
    UpperCAmelCase__						=image[0, -3:, -3:, -1]
    UpperCAmelCase__						=[0.55_31, 0.52_06, 0.48_95, 0.51_56, 0.51_82, 0.47_51, 0.48_02, 0.48_03, 0.44_43]

    assert image.shape == (1, 512, 512, 3)

    assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2






 def  __UpperCAmelCase						(  self		)	->       Union[str, Any]:
    UpperCAmelCase__						=StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5"		)
    UpperCAmelCase__						=sd_pipe.to(A_		)
    sd_pipe.set_progress_bar_config(disable=A_		)

    UpperCAmelCase__						=(
        "the four horsewomen of the apocalypse, painting by tom of finland, gaston bussiere, craig mullins, j. c."
        " leyendecker"
    )
    UpperCAmelCase__						=10_4435_5234
    UpperCAmelCase__						=12

    UpperCAmelCase__						=torch.manual_seed(A_		)
    UpperCAmelCase__						=sd_pipe(
        [prompt],		generator=A_,		guidance_scale=A_,		num_inference_steps=50,		output_type="np",		width=512,		height=512,		sld_guidance_scale=0,		)

    UpperCAmelCase__						=output.images
    UpperCAmelCase__						=image[0, -3:, -3:, -1]
    UpperCAmelCase__						=np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]		)

    assert image.shape == (1, 512, 512, 3)

    assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-7

    UpperCAmelCase__						=torch.manual_seed(A_		)
    UpperCAmelCase__						=sd_pipe(
        [prompt],		generator=A_,		guidance_scale=A_,		num_inference_steps=50,		output_type="np",		width=512,		height=512,		sld_guidance_scale=2000,		sld_warmup_steps=7,		sld_threshold=0.0_25,		sld_momentum_scale=0.5,		sld_mom_beta=0.7,		)

    UpperCAmelCase__						=output.images
    UpperCAmelCase__						=image[0, -3:, -3:, -1]
    UpperCAmelCase__						=np.array([0.58_18, 0.62_85, 0.68_35, 0.60_19, 0.6_25, 0.67_54, 0.60_96, 0.63_34, 0.65_61]		)
    assert image.shape == (1, 512, 512, 3)

    assert np.abs(image_slice.flatten() - expected_slice		).max() < 1E-2










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

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






def 				_UpperCAmelCase			(    A				,							A     ):




   '''simple docstring'''

   UpperCAmelCase__						={
       "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,
   }

   UpperCAmelCase__						=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
   UpperCAmelCase__						=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
   UpperCAmelCase__						="openwebtext_ccnews_stories_books_cased"

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

   UpperCAmelCase__						=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     )
   UpperCAmelCase__						=original_bort._collect_params_with_prefix()

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

   UpperCAmelCase__						=BertConfig.from_dict(A     )
   UpperCAmelCase__						=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     ):
      UpperCAmelCase__						=hf_param.shape

      UpperCAmelCase__						=to_torch(params[gluon_param]     )
      UpperCAmelCase__						=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

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

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

      # self attention
      UpperCAmelCase__						=layer.attention.self

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

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

      # self attention output
      UpperCAmelCase__						=layer.attention.output

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

      # intermediate
      UpperCAmelCase__						=layer.intermediate

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

      # output
      UpperCAmelCase__						=layer.output

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

   UpperCAmelCase__						=tokenizer.encode_plus(A     )["input_ids"]

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

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

   UpperCAmelCase__						=tokenizer.encode_plus(A				,							return_tensors="pt"     )
   UpperCAmelCase__						=hf_bort_model(**A     )[0]

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

   UpperCAmelCase__						=np.max(np.abs(hf_layer - gluon_layer     )     ).item()
   UpperCAmelCase__						=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__":
    UpperCamelCase_												=	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_												=	parser.parse_args()
    convert_bort_checkpoint_to_pytorch(args.bort_checkpoint_path, args.pytorch_dump_folder_path)



