





import gc
import os
import tempfile
import unittest

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import CLIPTextConfig, CLIPTextModel, CLIPTokenizer

from diffusers import AutoencoderKL, DDIMScheduler, StableDiffusionPipeline, UNetaDConditionModel
from diffusers.loaders import AttnProcsLayers, LoraLoaderMixin
from diffusers.models.attention_processor import (
    Attention,
    AttnProcessor,
    AttnProcessor2_0,
    LoRAAttnProcessor,
    LoRAAttnProcessor2_0,
    LoRAXFormersAttnProcessor,
    XFormersAttnProcessor,
)
from diffusers.utils import TEXT_ENCODER_ATTN_MODULE, floats_tensor, torch_device




def   _SCREAMING_SNAKE_CASE      (	snake_case_		:						nn.Module						):
     __magic_name__				=       {}
     for name in unet.attn_processors.keys():
          __magic_name__				=       None if name.endswith('''attn1.processor'''						) else unet.config.cross_attention_dim
          if name.startswith('''mid_block'''						):
               __magic_name__				=       unet.config.block_out_channels[-1]
          elif name.startswith('''up_blocks'''						):
               __magic_name__				=       int(name[len('''up_blocks.'''						)]						)
               __magic_name__				=       list(reversed(unet.config.block_out_channels						)						)[block_id]
          elif name.startswith('''down_blocks'''						):
               __magic_name__				=       int(name[len('''down_blocks.'''						)]						)
               __magic_name__				=       unet.config.block_out_channels[block_id]
          __magic_name__				=       (
              LoRAAttnProcessor2_0 if hasattr(snake_case_ ,     '''scaled_dot_product_attention'''						) else LoRAAttnProcessor
          )
          __magic_name__				=       lora_attn_processor_class(
              hidden_size=snake_case_ ,     cross_attention_dim=snake_case_						)
     __magic_name__				=       AttnProcsLayers(snake_case_						)
     return lora_attn_procs, unet_lora_layers




def   _SCREAMING_SNAKE_CASE      (	snake_case_		:						nn.Module						):
     __magic_name__				=       {}
     __magic_name__				=       (
         LoRAAttnProcessor2_0 if hasattr(snake_case_ ,     '''scaled_dot_product_attention'''						) else LoRAAttnProcessor
     )
     for name, module in text_encoder.named_modules():
          if name.endswith(snake_case_						):
               __magic_name__				=       lora_attn_processor_class(
                   hidden_size=module.out_proj.out_features ,     cross_attention_dim=snake_case_						)
     return text_lora_attn_procs




def   _SCREAMING_SNAKE_CASE      (	snake_case_		:						nn.Module						):
     __magic_name__				=       create_text_encoder_lora_attn_procs(snake_case_						)
     __magic_name__				=       AttnProcsLayers(snake_case_						)
     return text_encoder_lora_layers




def   _SCREAMING_SNAKE_CASE      (	snake_case_		:						List[Any] ,     snake_case_		:						Union[str, Any]=False						):
     for _, attn_proc in text_lora_attn_procs.items():
          # set up.weights
          for layer_name, layer_module in attn_proc.named_modules():
               if layer_name.endswith('''_lora'''						):
                    __magic_name__				=       (
                        torch.randn_like(layer_module.up.weight						)
                        if randn_weight
                        else torch.zeros_like(layer_module.up.weight						)
                    )
                    __magic_name__				=       torch.nn.Parameter(snake_case_						)



class  SCREAMING_SNAKE_CASE_  ( unittest.TestCase			):


  """simple docstring"""





  def 			__A      (  self						)		->   List[str]:





       '''simple docstring'''







       torch.manual_seed(0						)
       __magic_name__				=       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	,  )
       __magic_name__				=       DDIMScheduler(
           beta_start=0.0_00_85	,  beta_end=0.0_12	,  beta_schedule='''scaled_linear'''	,  clip_sample=A	,  set_alpha_to_one=A	,  steps_offset=1	,  )
       torch.manual_seed(0						)
       __magic_name__				=       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	,  )
       __magic_name__				=       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=10_00	,  )
       __magic_name__				=       CLIPTextModel(A						)
       __magic_name__				=       CLIPTokenizer.from_pretrained('''hf-internal-testing/tiny-random-clip'''						)

       __magic_name__		,  __magic_name__				=       create_unet_lora_layers(A						)
       __magic_name__				=       create_text_encoder_lora_layers(A						)

       __magic_name__				=       {
           '''unet''': unet,
           '''scheduler''': scheduler,
           '''vae''': vae,
           '''text_encoder''': text_encoder,
           '''tokenizer''': tokenizer,
           '''safety_checker''': None,
           '''feature_extractor''': None,
       }
       __magic_name__				=       {
           '''unet_lora_layers''': unet_lora_layers,
           '''text_encoder_lora_layers''': text_encoder_lora_layers,
           '''unet_lora_attn_procs''': unet_lora_attn_procs,
       }
       return pipeline_components, lora_components





  def 			__A      (  self						)		->   Tuple:





       '''simple docstring'''







       __magic_name__				=       1
       __magic_name__				=       10
       __magic_name__				=       4
       __magic_name__				=       (32, 32)

       __magic_name__				=       torch.manual_seed(0						)
       __magic_name__				=       floats_tensor((batch_size, num_channels) + sizes						)
       __magic_name__				=       torch.randint(1	,  A	,  size=(batch_size, sequence_length)	,  generator=A						)

       __magic_name__				=       {
           '''prompt''': '''A painting of a squirrel eating a burger''',
           '''generator''': generator,
           '''num_inference_steps''': 2,
           '''guidance_scale''': 6.0,
           '''output_type''': '''numpy''',
       }

       return noise, input_ids, pipeline_inputs

       # copied from: https://colab.research.google.com/gist/sayakpaul/df2ef6e1ae6d8c10a49d859883b10860/scratchpad.ipynb





  def 			__A      (  self						)		->   Dict:





       '''simple docstring'''







       __magic_name__				=       77

       __magic_name__				=       torch.randint(2	,  56	,  size=(1, max_seq_length)	,  generator=torch.manual_seed(0						)						)

       __magic_name__				=       {}
       __magic_name__				=       inputs
       return prepared_inputs





  def 			__A      (  self	,  A						)		->   List[Any]:





       '''simple docstring'''







       __magic_name__		,  __magic_name__				=       self.get_dummy_components()
       LoraLoaderMixin.save_lora_weights(
           save_directory=A	,  unet_lora_layers=lora_components['''unet_lora_layers''']	,  text_encoder_lora_layers=lora_components['''text_encoder_lora_layers''']	,  )
       self.assertTrue(os.path.isfile(os.path.join(A	,  '''pytorch_lora_weights.bin'''						)						)						)





  def 			__A      (  self						)		->   Union[str, Any]:





       '''simple docstring'''







       __magic_name__		,  __magic_name__				=       self.get_dummy_components()
       __magic_name__				=       StableDiffusionPipeline(**A						)
       __magic_name__				=       sd_pipe.to(A						)
       sd_pipe.set_progress_bar_config(disable=A						)

       __magic_name__		,  __magic_name__		,  __magic_name__				=       self.get_dummy_inputs()

       __magic_name__				=       sd_pipe(**A						).images
       __magic_name__				=       original_images[0, -3:, -3:, -1]

       with tempfile.TemporaryDirectory() as tmpdirname:
            LoraLoaderMixin.save_lora_weights(
                save_directory=A	,  unet_lora_layers=lora_components['''unet_lora_layers''']	,  text_encoder_lora_layers=lora_components['''text_encoder_lora_layers''']	,  )
            self.assertTrue(os.path.isfile(os.path.join(A	,  '''pytorch_lora_weights.bin'''						)						)						)
            sd_pipe.load_lora_weights(A						)

       __magic_name__				=       sd_pipe(**A						).images
       __magic_name__				=       lora_images[0, -3:, -3:, -1]

       # Outputs shouldn't match.
       self.assertFalse(torch.allclose(torch.from_numpy(A						)	,  torch.from_numpy(A						)						)						)





  def 			__A      (  self						)		->   Union[str, Any]:





       '''simple docstring'''







       __magic_name__		,  __magic_name__				=       self.get_dummy_components()
       __magic_name__				=       StableDiffusionPipeline(**A						)
       __magic_name__				=       sd_pipe.to(A						)
       sd_pipe.set_progress_bar_config(disable=A						)

       __magic_name__		,  __magic_name__		,  __magic_name__				=       self.get_dummy_inputs()

       __magic_name__				=       sd_pipe(**A						).images
       __magic_name__				=       original_images[0, -3:, -3:, -1]

       with tempfile.TemporaryDirectory() as tmpdirname:
            LoraLoaderMixin.save_lora_weights(
                save_directory=A	,  unet_lora_layers=lora_components['''unet_lora_layers''']	,  text_encoder_lora_layers=lora_components['''text_encoder_lora_layers''']	,  safe_serialization=A	,  )
            self.assertTrue(os.path.isfile(os.path.join(A	,  '''pytorch_lora_weights.safetensors'''						)						)						)
            sd_pipe.load_lora_weights(A						)

       __magic_name__				=       sd_pipe(**A						).images
       __magic_name__				=       lora_images[0, -3:, -3:, -1]

       # Outputs shouldn't match.
       self.assertFalse(torch.allclose(torch.from_numpy(A						)	,  torch.from_numpy(A						)						)						)





  def 			__A      (  self						)		->   int:





       '''simple docstring'''







       __magic_name__		,  __magic_name__				=       self.get_dummy_components()
       __magic_name__				=       lora_components['''unet_lora_attn_procs''']
       __magic_name__				=       StableDiffusionPipeline(**A						)
       __magic_name__				=       sd_pipe.to(A						)
       sd_pipe.set_progress_bar_config(disable=A						)

       __magic_name__		,  __magic_name__		,  __magic_name__				=       self.get_dummy_inputs()

       __magic_name__				=       sd_pipe(**A						).images
       __magic_name__				=       original_images[0, -3:, -3:, -1]

       with tempfile.TemporaryDirectory() as tmpdirname:
            __magic_name__				=       sd_pipe.unet
            unet.set_attn_processor(A						)
            unet.save_attn_procs(A						)
            self.assertTrue(os.path.isfile(os.path.join(A	,  '''pytorch_lora_weights.bin'''						)						)						)
            sd_pipe.load_lora_weights(A						)

       __magic_name__				=       sd_pipe(**A						).images
       __magic_name__				=       lora_images[0, -3:, -3:, -1]

       # Outputs shouldn't match.
       self.assertFalse(torch.allclose(torch.from_numpy(A						)	,  torch.from_numpy(A						)						)						)





  def 			__A      (  self						)		->   Dict:





       '''simple docstring'''







       __magic_name__		,  __magic_name__				=       self.get_dummy_components()
       __magic_name__				=       StableDiffusionPipeline(**A						)

       __magic_name__				=       self.get_dummy_tokens()

       # inference without lora
       __magic_name__				=       pipe.text_encoder(**A						)[0]
       assert outputs_without_lora.shape == (1, 77, 32)

       # create lora_attn_procs with zeroed out up.weights
       __magic_name__				=       create_text_encoder_lora_attn_procs(pipe.text_encoder						)
       set_lora_up_weights(A	,  randn_weight=A						)

       # monkey patch
       pipe._modify_text_encoder(A						)

       # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
       del text_attn_procs
       gc.collect()

       # inference with lora
       __magic_name__				=       pipe.text_encoder(**A						)[0]
       assert outputs_with_lora.shape == (1, 77, 32)

       assert torch.allclose(
           A	,  A						), "lora_up_weight are all zero, so the lora outputs should be the same to without lora outputs"

       # create lora_attn_procs with randn up.weights
       __magic_name__				=       create_text_encoder_lora_attn_procs(pipe.text_encoder						)
       set_lora_up_weights(A	,  randn_weight=A						)

       # monkey patch
       pipe._modify_text_encoder(A						)

       # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
       del text_attn_procs
       gc.collect()

       # inference with lora
       __magic_name__				=       pipe.text_encoder(**A						)[0]
       assert outputs_with_lora.shape == (1, 77, 32)

       assert not torch.allclose(
           A	,  A						), "lora_up_weight are not zero, so the lora outputs should be different to without lora outputs"





  def 			__A      (  self						)		->   Union[str, Any]:





       '''simple docstring'''







       __magic_name__		,  __magic_name__				=       self.get_dummy_components()
       __magic_name__				=       StableDiffusionPipeline(**A						)

       __magic_name__				=       self.get_dummy_tokens()

       # inference without lora
       __magic_name__				=       pipe.text_encoder(**A						)[0]
       assert outputs_without_lora.shape == (1, 77, 32)

       # create lora_attn_procs with randn up.weights
       __magic_name__				=       create_text_encoder_lora_attn_procs(pipe.text_encoder						)
       set_lora_up_weights(A	,  randn_weight=A						)

       # monkey patch
       pipe._modify_text_encoder(A						)

       # verify that it's okay to release the text_attn_procs which holds the LoRAAttnProcessor.
       del text_attn_procs
       gc.collect()

       # inference with lora
       __magic_name__				=       pipe.text_encoder(**A						)[0]
       assert outputs_with_lora.shape == (1, 77, 32)

       assert not torch.allclose(
           A	,  A						), "lora outputs should be different to without lora outputs"

       # remove monkey patch
       pipe._remove_text_encoder_monkey_patch()

       # inference with removed lora
       __magic_name__				=       pipe.text_encoder(**A						)[0]
       assert outputs_without_lora_removed.shape == (1, 77, 32)

       assert torch.allclose(
           A	,  A						), "remove lora monkey patch should restore the original outputs"





  def 			__A      (  self						)		->   List[Any]:





       '''simple docstring'''







       __magic_name__		,  __magic_name__				=       self.get_dummy_components()
       __magic_name__				=       StableDiffusionPipeline(**A						)
       __magic_name__				=       sd_pipe.to(A						)
       sd_pipe.set_progress_bar_config(disable=A						)

       __magic_name__		,  __magic_name__		,  __magic_name__				=       self.get_dummy_inputs()

       with tempfile.TemporaryDirectory() as tmpdirname:
            LoraLoaderMixin.save_lora_weights(
                save_directory=A	,  unet_lora_layers=lora_components['''unet_lora_layers''']	,  text_encoder_lora_layers=lora_components['''text_encoder_lora_layers''']	,  )
            self.assertTrue(os.path.isfile(os.path.join(A	,  '''pytorch_lora_weights.bin'''						)						)						)
            sd_pipe.load_lora_weights(A						)

       __magic_name__				=       sd_pipe(**A						).images
       __magic_name__				=       lora_images[0, -3:, -3:, -1]

       __magic_name__				=       sd_pipe(**A	,  cross_attention_kwargs={'''scale''': 0.5}						).images
       __magic_name__				=       lora_images_with_scale[0, -3:, -3:, -1]

       # Outputs shouldn't match.
       self.assertFalse(
           torch.allclose(torch.from_numpy(A						)	,  torch.from_numpy(A						)						)						)





  def 			__A      (  self						)		->   Optional[int]:





       '''simple docstring'''







       with tempfile.TemporaryDirectory() as tmpdirname:
            self.create_lora_weight_file(A						)

            __magic_name__		,  __magic_name__				=       self.get_dummy_components()
            __magic_name__				=       StableDiffusionPipeline(**A						)
            __magic_name__				=       sd_pipe.to(A						)
            sd_pipe.set_progress_bar_config(disable=A						)

            # check if vanilla attention processors are used
            for _, module in sd_pipe.unet.named_modules():
                 if isinstance(A	,  A						):
                      self.assertIsInstance(module.processor	,  (AttnProcessor, AttnProcessor2_0)						)

            # load LoRA weight file
            sd_pipe.load_lora_weights(A						)

            # check if lora attention processors are used
            for _, module in sd_pipe.unet.named_modules():
                 if isinstance(A	,  A						):
                      __magic_name__				=       (
                          LoRAAttnProcessor2_0 if hasattr(A	,  '''scaled_dot_product_attention'''						) else LoRAAttnProcessor
                      )
                      self.assertIsInstance(module.processor	,  A						)





  @unittest.skipIf(torch_device != '''cuda'''	,  '''This test is supposed to run on GPU'''						)
  def 			__A      (  self						)		->   str:





       '''simple docstring'''







       with tempfile.TemporaryDirectory() as tmpdirname:
            self.create_lora_weight_file(A						)

            __magic_name__		,  __magic_name__				=       self.get_dummy_components()
            __magic_name__				=       StableDiffusionPipeline(**A						)
            __magic_name__				=       sd_pipe.to(A						)
            sd_pipe.set_progress_bar_config(disable=A						)

            # enable XFormers
            sd_pipe.enable_xformers_memory_efficient_attention()

            # check if xFormers attention processors are used
            for _, module in sd_pipe.unet.named_modules():
                 if isinstance(A	,  A						):
                      self.assertIsInstance(module.processor	,  A						)

            # load LoRA weight file
            sd_pipe.load_lora_weights(A						)

            # check if lora attention processors are used
            for _, module in sd_pipe.unet.named_modules():
                 if isinstance(A	,  A						):
                      self.assertIsInstance(module.processor	,  A						)







  @unittest.skipIf(torch_device != '''cuda'''	,  '''This test is supposed to run on GPU'''						)
  def 			__A      (  self						)		->   Optional[Any]:





       '''simple docstring'''







       __magic_name__		,  __magic_name__				=       self.get_dummy_components()
       __magic_name__				=       StableDiffusionPipeline(**A						)
       __magic_name__				=       sd_pipe.to(A						)
       sd_pipe.set_progress_bar_config(disable=A						)

       __magic_name__		,  __magic_name__		,  __magic_name__				=       self.get_dummy_inputs()

       # enable XFormers
       sd_pipe.enable_xformers_memory_efficient_attention()

       __magic_name__				=       sd_pipe(**A						).images
       __magic_name__				=       original_images[0, -3:, -3:, -1]

       with tempfile.TemporaryDirectory() as tmpdirname:
            LoraLoaderMixin.save_lora_weights(
                save_directory=A	,  unet_lora_layers=lora_components['''unet_lora_layers''']	,  text_encoder_lora_layers=lora_components['''text_encoder_lora_layers''']	,  )
            self.assertTrue(os.path.isfile(os.path.join(A	,  '''pytorch_lora_weights.bin'''						)						)						)
            sd_pipe.load_lora_weights(A						)

       __magic_name__				=       sd_pipe(**A						).images
       __magic_name__				=       lora_images[0, -3:, -3:, -1]

       # Outputs shouldn't match.
       self.assertFalse(torch.allclose(torch.from_numpy(A						)	,  torch.from_numpy(A						)						)						)






# Copyright (c) 2021-, NVIDIA CORPORATION.  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.

####################################################################################################

#
# Note: If when running this conversion script you're getting an exception:
#     ModuleNotFoundError: No module named 'megatron.model.enums'
# you need to tell python where to find the clone of Megatron-LM, e.g.:
#
# cd /tmp
# git clone https://github.com/NVIDIA/Megatron-LM
# PYTHONPATH=/tmp/Megatron-LM python src/transformers/models/megatron_bert/convert_megatron_bert_checkpoint.py ...
#
# if you already have it cloned elsewhere, simply adjust the path to the existing path
#
# If the training was done using a Megatron-LM fork, e.g.,
# https://github.com/microsoft/Megatron-DeepSpeed/ then chances are that you need to have that one
# in your path, i.e., /path/to/Megatron-DeepSpeed/
#

import argparse
import os
import re
import zipfile

import torch

from transformers import MegatronBertConfig




def   _SCREAMING_SNAKE_CASE      (	snake_case_		:						Optional[int] ,     snake_case_		:						int ,     snake_case_		:						int=0						):
     # Format the message.
     if name is None:
          __magic_name__				=       None
     else:
          __magic_name__				=       '''.''' * max(0 ,     spaces - 2						) + '''# {:''' + str(50 - spaces						) + '''s}'''
          __magic_name__				=       fmt.format(snake_case_						)

     # Print and recurse (if needed).
     if isinstance(snake_case_ ,     snake_case_						):
          if msg is not None:
               print(snake_case_						)
          for k in val.keys():
               recursive_print(snake_case_ ,     val[k] ,     spaces + 2						)
     elif isinstance(snake_case_ ,     torch.Tensor						):
          print(snake_case_ ,     ''':''' ,     val.size()						)
     else:
          print(snake_case_ ,     ''':''' ,     snake_case_						)




def   _SCREAMING_SNAKE_CASE      (	snake_case_		:						int ,     snake_case_		:						Optional[int] ,     snake_case_		:						Any ,     snake_case_		:						Tuple ,     snake_case_		:						Tuple						):
     # Permutes layout of param tensor to [num_splits * num_heads * hidden_size, :]
     # for compatibility with later versions of NVIDIA Megatron-LM.
     # The inverse operation is performed inside Megatron-LM to read checkpoints:
     # https://github.com/NVIDIA/Megatron-LM/blob/v2.4/megatron/checkpointing.py#L209
     # If param is the weight tensor of the self-attention block, the returned tensor
     # will have to be transposed one more time to be read by HuggingFace BERT.
     __magic_name__				=       param.size()
     if checkpoint_version == 1.0:
          # version 1.0 stores [num_heads * hidden_size * num_splits, :]
          __magic_name__				=       (num_heads, hidden_size, num_splits) + input_shape[1:]
          __magic_name__				=       param.view(*snake_case_						)
          __magic_name__				=       param.transpose(0 ,     2						)
          __magic_name__				=       param.transpose(1 ,     2						).contiguous()
     elif checkpoint_version >= 2.0:
          # other versions store [num_heads * num_splits * hidden_size, :]
          __magic_name__				=       (num_heads, num_splits, hidden_size) + input_shape[1:]
          __magic_name__				=       param.view(*snake_case_						)
          __magic_name__				=       param.transpose(0 ,     1						).contiguous()
     __magic_name__				=       param.view(*snake_case_						)
     return param




def   _SCREAMING_SNAKE_CASE      (	snake_case_		:						int ,     snake_case_		:						Optional[int] ,     snake_case_		:						Union[str, Any]						):
     # The converted output model.
     __magic_name__				=       {}

     # old versions did not store training args
     __magic_name__				=       input_state_dict.get('''args''' ,     snake_case_						)
     if ds_args is not None:
          # do not make the user write a config file when the exact dimensions/sizes are already in the checkpoint
          # from pprint import pprint
          # pprint(vars(ds_args))

          __magic_name__				=       ds_args.tokenizer_type
          __magic_name__				=       ds_args.padded_vocab_size
          __magic_name__				=       ds_args.max_position_embeddings
          __magic_name__				=       ds_args.hidden_size
          __magic_name__				=       ds_args.num_layers
          __magic_name__				=       ds_args.num_attention_heads
          __magic_name__				=       ds_args.ffn_hidden_size if '''ffn_hidden_size''' in ds_args else 4 * ds_args.hidden_size
          # pprint(config)

     # The number of heads.
     __magic_name__				=       config.num_attention_heads
     # The hidden_size per head.
     __magic_name__				=       config.hidden_size // heads
     # Megatron-LM checkpoint version
     if "checkpoint_version" in input_state_dict.keys():
          __magic_name__				=       input_state_dict['''checkpoint_version''']
     else:
          __magic_name__				=       0.0

     # The model.
     __magic_name__				=       input_state_dict['''model''']
     # The language model.
     __magic_name__				=       model['''language_model''']
     # The embeddings.
     __magic_name__				=       lm['''embedding''']

     # The word embeddings.
     __magic_name__				=       embeddings['''word_embeddings''']['''weight''']
     # Truncate the embedding table to vocab_size rows.
     __magic_name__				=       word_embeddings[: config.vocab_size, :]
     # Store the word embeddings.
     __magic_name__				=       word_embeddings

     # The position embeddings.
     __magic_name__				=       embeddings['''position_embeddings''']['''weight''']
     assert pos_embeddings.size(0						) == config.max_position_embeddings and pos_embeddings.size(1						) == config.hidden_size
     # Store the position embeddings.
     __magic_name__				=       pos_embeddings

     # The token-type embeddings.
     __magic_name__				=       embeddings['''tokentype_embeddings''']['''weight''']
     # Store the position embeddings.
     __magic_name__				=       tokentype_embeddings

     # The transformer.
     __magic_name__				=       lm['''transformer'''] if '''transformer''' in lm.keys() else lm['''encoder''']

     # The regex to extract layer names.
     __magic_name__				=       re.compile(r'''layers\.(\d+)\.([a-z0-9_.]+)\.([a-z]+)'''						)

     # The simple map of names for "automated" rules.
     __magic_name__				=       {
         '''attention.dense''': '''.attention.output.dense.''',
         '''self_attention.dense''': '''.attention.output.dense.''',
         '''mlp.dense_h_to_4h''': '''.intermediate.dense.''',
         '''mlp.dense_4h_to_h''': '''.output.dense.''',
     }

     # Keep track of the attention/query/value tensor.
     __magic_name__				=       None

     # Extract the layers.
     for key, val in transformer.items():
          # Match the name.
          __magic_name__				=       layer_re.match(snake_case_						)

          # Stop if that's not a layer
          if m is None:
               break

          # The index of the layer.
          __magic_name__				=       int(m.group(1						)						)
          # The name of the operation.
          __magic_name__				=       m.group(2						)
          # Is it a weight or a bias?
          __magic_name__				=       m.group(3						)

          # The name of the layer.
          __magic_name__				=       f'bert.encoder.layer.{layer_idx}'

          # For layernorm(s), simply store the layer norm.
          if op_name.endswith('''layernorm'''						):
               __magic_name__				=       '''attention.ln''' if op_name.startswith('''input'''						) else '''ln'''
               __magic_name__				=       val

          # Transpose the QKV matrix.
          elif (
              op_name == "attention.query_key_value" or op_name == "self_attention.query_key_value"
          ) and weight_or_bias == "weight":
               # Make sure the QKV pointer is nil.
               assert attention_qkv_weight is None, ""

               __magic_name__				=       fix_query_key_value_ordering(snake_case_ ,     snake_case_ ,     3 ,     snake_case_ ,     snake_case_						)
               # Store the tensor as we need the bias as well to interleave QKV and biases.
               __magic_name__				=       out_val

          # Transpose the bias.
          elif (
              op_name == "attention.query_key_value" or op_name == "self_attention.query_key_value"
          ) and weight_or_bias == "bias":
               # Make sure we read the weight tensor.
               assert attention_qkv_weight is not None, ""

               # Split the QKV matrix into Q, K and V. Megatron stores Q,K,V interleaved.
               __magic_name__				=       attention_qkv_weight[0 * config.hidden_size : 1 * config.hidden_size, :]
               __magic_name__				=       attention_qkv_weight[1 * config.hidden_size : 2 * config.hidden_size, :]
               __magic_name__				=       attention_qkv_weight[2 * config.hidden_size : 3 * config.hidden_size, :]

               __magic_name__				=       fix_query_key_value_ordering(snake_case_ ,     snake_case_ ,     3 ,     snake_case_ ,     snake_case_						)
               # Split the bias.
               __magic_name__				=       out_val[0 * config.hidden_size : 1 * config.hidden_size]
               __magic_name__				=       out_val[1 * config.hidden_size : 2 * config.hidden_size]
               __magic_name__				=       out_val[2 * config.hidden_size : 3 * config.hidden_size]

               # Store.
               __magic_name__				=       q
               __magic_name__				=       q_bias
               __magic_name__				=       k
               __magic_name__				=       k_bias
               __magic_name__				=       v
               __magic_name__				=       v_bias

               # Clear the stored tensor.
               __magic_name__				=       None

          # Copy weights and biases as is.
          elif weight_or_bias in ["weight", "bias"]:
               __magic_name__				=       megatron_to_transformers[op_name]
               __magic_name__				=       val

    # The final layernorm.
     __magic_name__				=       transformer['''final_layernorm.weight''']
     __magic_name__				=       transformer['''final_layernorm.bias''']

     # The pooler.
     __magic_name__				=       lm['''pooler''']

     # Store the matrix and the bias.
     __magic_name__				=       pooler['''dense.weight''']
     __magic_name__				=       pooler['''dense.bias''']

     # The LM head from Megatron (for RACE).
     __magic_name__				=       model['''lm_head''']

     # The transform matrix.
     __magic_name__				=       lm_head['''dense.weight''']
     __magic_name__				=       lm_head['''dense.bias''']

     # The transform LN.
     __magic_name__				=       lm_head['''layernorm.weight''']
     __magic_name__				=       lm_head['''layernorm.bias''']

     # For the decoder, we replicate the weights.
     __magic_name__				=       word_embeddings
     __magic_name__				=       lm_head['''bias''']

     # The classifier from Megatron (for MLNI).
     __magic_name__				=       model['''binary_head''']

     # Store the classifier.
     __magic_name__				=       binary_head['''weight''']
     __magic_name__				=       binary_head['''bias''']

     # It should be done!
     return output_state_dict




def   _SCREAMING_SNAKE_CASE      (	):
     # Create the argument parser.
     __magic_name__				=       argparse.ArgumentParser()
     parser.add_argument('''--print-checkpoint-structure''' ,     action='''store_true'''						)
     parser.add_argument('''path_to_checkpoint''' ,     type=snake_case_ ,     help='''Path to the ZIP file containing the checkpoint'''						)
     parser.add_argument(
         '''--config_file''' ,     default='''''' ,     type=snake_case_ ,     help='''An optional config json file describing the pre-trained model.''' ,     )
     __magic_name__				=       parser.parse_args()

     # Extract the basename.
     __magic_name__				=       os.path.dirname(args.path_to_checkpoint						)

     # Load the model.
     # the .zip is very optional, let's keep it for backward compatibility
     print(f'Extracting PyTorch state dictionary from "{args.path_to_checkpoint}"'						)
     if args.path_to_checkpoint.endswith('''.zip'''						):
          with zipfile.ZipFile(args.path_to_checkpoint ,     '''r'''						) as checkpoint:
               with checkpoint.open('''release/mp_rank_00/model_optim_rng.pt'''						) as pytorch_dict:
                    __magic_name__				=       torch.load(snake_case_ ,     map_location='''cpu'''						)
     else:
          __magic_name__				=       torch.load(args.path_to_checkpoint ,     map_location='''cpu'''						)

     if args.config_file == "":
          # Default config of megatron-bert 345m
          __magic_name__				=       MegatronBertConfig()

          # different megatron-bert-*-345m models have different vocab sizes, so override the default
          # config (which is for megatron-bert-cased-345m) with the actual vocab dimension
          __magic_name__				=       input_state_dict['''model''']['''lm_head''']['''bias'''].numel()
     else:
          __magic_name__				=       MegatronBertConfig.from_json_file(args.config_file						)

     # Convert.
     print('''Converting'''						)
     __magic_name__				=       convert_megatron_checkpoint(snake_case_ ,     snake_case_ ,     snake_case_						)

     # Print the structure of converted state dict.
     if args.print_checkpoint_structure:
          recursive_print(snake_case_ ,     snake_case_						)

     # Store the config to file.
     print('''Saving config'''						)
     config.save_pretrained(snake_case_						)

     # Store the state_dict to file.
     __magic_name__				=       os.path.join(snake_case_ ,     '''pytorch_model.bin'''						)
     print(f'Saving checkpoint to "{output_checkpoint_file}"'						)
     torch.save(snake_case_ ,     snake_case_						)


####################################################################################################

if __name__ == "__main__":
    main()

####################################################################################################