'''simple docstring'''





import tempfile
import unittest

from diffusers import (
    DDIMScheduler,
    DDPMScheduler,
    DPMSolverMultistepScheduler,
    EulerAncestralDiscreteScheduler,
    EulerDiscreteScheduler,
    PNDMScheduler,
    logging,
)
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.utils.testing_utils import CaptureLogger



class    UpperCamelCase (    lowerCAmelCase__    ):







    """simple docstring"""





    A			:  Tuple            =       "config.json"





    @register_to_config
    def __init__( self			:		int				,				UpperCAmelCase_			:		Tuple=2				,				UpperCAmelCase_			:		Optional[int]=5				,				UpperCAmelCase_			:		List[Any]=(2, 5)				,				UpperCAmelCase_			:		Optional[Any]="for diffusion"				,				UpperCAmelCase_			:		Tuple=[1, 3]				,				):


         """simple docstring"""







         pass






class    UpperCamelCase (    lowerCAmelCase__    ):







    """simple docstring"""





    A			:  Tuple            =       "config.json"





    @register_to_config
    def __init__( self			:		List[str]				,				UpperCAmelCase_			:		List[str]=2				,				UpperCAmelCase_			:		str=5				,				UpperCAmelCase_			:		Union[str, Any]=(2, 5)				,				UpperCAmelCase_			:		List[str]="for diffusion"				,				UpperCAmelCase_			:		Optional[int]=[1, 3]				,				):


         """simple docstring"""







         pass






class    UpperCamelCase (    lowerCAmelCase__    ):







    """simple docstring"""





    A			:  Optional[int]            =       "config.json"





    @register_to_config
    def __init__( self			:		Union[str, Any]				,				UpperCAmelCase_			:		str=2				,				UpperCAmelCase_			:		int=5				,				UpperCAmelCase_			:		Optional[int]=(2, 5)				,				UpperCAmelCase_			:		Any="for diffusion"				,				UpperCAmelCase_			:		Any=[1, 3]				,				UpperCAmelCase_			:		int=[1, 3]				,				):


         """simple docstring"""







         pass






class    UpperCamelCase (    lowerCAmelCase__    ):







    """simple docstring"""





    A			:  Optional[int]            =       "config.json"





    @register_to_config
    def __init__( self			:		int				,				UpperCAmelCase_			:		Any=2				,				UpperCAmelCase_			:		str=5				,				UpperCAmelCase_			:		Optional[int]=(2, 5)				,				UpperCAmelCase_			:		Any="for diffusion"				,				UpperCAmelCase_			:		Optional[int]=[1, 5]				,				UpperCAmelCase_			:		Union[str, Any]=[5, 4]				,				):


         """simple docstring"""







         pass






class    UpperCamelCase (    unittest.TestCase    ):







    """simple docstring"""





    def 					SCREAMING_SNAKE_CASE_			( self			:		Optional[int]):


         """simple docstring"""







         with self.assertRaises(a__):
              ConfigMixin.load_config('dummy_path')





    def 					SCREAMING_SNAKE_CASE_			( self			:		Any):


         """simple docstring"""







         a :				List[str]				=      SampleObject()
         a :				str				=      obj.config
         assert config["a"] == 2
         assert config["b"] == 5
         assert config["c"] == (2, 5)
         assert config["d"] == "for diffusion"
         assert config["e"] == [1, 3]

         # init ignore private arguments
         a :				Union[str, Any]				=      SampleObject(_name_or_path='lalala')
         a :				List[Any]				=      obj.config
         assert config["a"] == 2
         assert config["b"] == 5
         assert config["c"] == (2, 5)
         assert config["d"] == "for diffusion"
         assert config["e"] == [1, 3]

         # can override default
         a :				Union[str, Any]				=      SampleObject(c=6)
         a :				Optional[int]				=      obj.config
         assert config["a"] == 2
         assert config["b"] == 5
         assert config["c"] == 6
         assert config["d"] == "for diffusion"
         assert config["e"] == [1, 3]

         # can use positional arguments.
         a :				Dict				=      SampleObject(1				,				c=6)
         a :				Optional[int]				=      obj.config
         assert config["a"] == 1
         assert config["b"] == 5
         assert config["c"] == 6
         assert config["d"] == "for diffusion"
         assert config["e"] == [1, 3]





    def 					SCREAMING_SNAKE_CASE_			( self			:		Optional[Any]):


         """simple docstring"""







         a :				Any				=      SampleObject()
         a :				Optional[Any]				=      obj.config

         assert config["a"] == 2
         assert config["b"] == 5
         assert config["c"] == (2, 5)
         assert config["d"] == "for diffusion"
         assert config["e"] == [1, 3]

         with tempfile.TemporaryDirectory() as tmpdirname:
              obj.save_config(a__)
              a :				str				=      SampleObject.from_config(SampleObject.load_config(a__))
              a :				Dict				=      new_obj.config

         # unfreeze configs
         a :				Dict				=      dict(a__)
         a :				Any				=      dict(a__)

         assert config.pop('c') == (2, 5)  # instantiated as tuple
         assert new_config.pop('c') == [2, 5]  # saved & loaded as list because of json
         config.pop('_use_default_values')
         assert config == new_config





    def 					SCREAMING_SNAKE_CASE_			( self			:		List[Any]):


         """simple docstring"""







         a :				Optional[Any]				=      logging.get_logger('diffusers.configuration_utils')
         # 30 for warning
         logger.setLevel(3_0)

         with CaptureLogger(a__) as cap_logger:
              a :				Optional[int]				=      DDIMScheduler.from_pretrained(
                  'hf-internal-testing/tiny-stable-diffusion-torch'				,				subfolder='scheduler')

         assert ddim.__class__ == DDIMScheduler
         # no warning should be thrown
         assert cap_logger.out == ""





    def 					SCREAMING_SNAKE_CASE_			( self			:		List[str]):


         """simple docstring"""







         a :				int				=      logging.get_logger('diffusers.configuration_utils')
         # 30 for warning
         logger.setLevel(3_0)

         with CaptureLogger(a__) as cap_logger:
              a :				List[Any]				=      EulerDiscreteScheduler.from_pretrained(
                  'hf-internal-testing/tiny-stable-diffusion-torch'				,				subfolder='scheduler')

         assert euler.__class__ == EulerDiscreteScheduler
         # no warning should be thrown
         assert cap_logger.out == ""





    def 					SCREAMING_SNAKE_CASE_			( self			:		Dict):


         """simple docstring"""







         a :				List[Any]				=      logging.get_logger('diffusers.configuration_utils')
         # 30 for warning
         logger.setLevel(3_0)

         with CaptureLogger(a__) as cap_logger:
              a :				List[str]				=      EulerAncestralDiscreteScheduler.from_pretrained(
                  'hf-internal-testing/tiny-stable-diffusion-torch'				,				subfolder='scheduler')

         assert euler.__class__ == EulerAncestralDiscreteScheduler
         # no warning should be thrown
         assert cap_logger.out == ""





    def 					SCREAMING_SNAKE_CASE_			( self			:		Optional[int]):


         """simple docstring"""







         a :				List[str]				=      logging.get_logger('diffusers.configuration_utils')
         # 30 for warning
         logger.setLevel(3_0)

         with CaptureLogger(a__) as cap_logger:
              a :				int				=      PNDMScheduler.from_pretrained(
                  'hf-internal-testing/tiny-stable-diffusion-torch'				,				subfolder='scheduler')

         assert pndm.__class__ == PNDMScheduler
         # no warning should be thrown
         assert cap_logger.out == ""





    def 					SCREAMING_SNAKE_CASE_			( self			:		Tuple):


         """simple docstring"""







         a :				str				=      logging.get_logger('diffusers.configuration_utils')
         # 30 for warning
         logger.setLevel(3_0)

         with CaptureLogger(a__) as cap_logger:
              a :				Optional[Any]				=      DDPMScheduler.from_pretrained(
                  'hf-internal-testing/tiny-stable-diffusion-torch'				,				subfolder='scheduler'				,				prediction_type='sample'				,				beta_end=8				,				)

         with CaptureLogger(a__) as cap_logger_a:
              a :				Union[str, Any]				=      DDPMScheduler.from_pretrained('google/ddpm-celebahq-256'				,				beta_start=8_8)

         assert ddpm.__class__ == DDPMScheduler
         assert ddpm.config.prediction_type == "sample"
         assert ddpm.config.beta_end == 8
         assert ddpm_a.config.beta_start == 8_8

         # no warning should be thrown
         assert cap_logger.out == ""
         assert cap_logger_a.out == ""





    def 					SCREAMING_SNAKE_CASE_			( self			:		Any):


         """simple docstring"""







         a :				List[Any]				=      logging.get_logger('diffusers.configuration_utils')
         # 30 for warning
         logger.setLevel(3_0)

         with CaptureLogger(a__) as cap_logger:
              a :				Any				=      DPMSolverMultistepScheduler.from_pretrained(
                  'hf-internal-testing/tiny-stable-diffusion-torch'				,				subfolder='scheduler')

         assert dpm.__class__ == DPMSolverMultistepScheduler
         # no warning should be thrown
         assert cap_logger.out == ""







    def 					SCREAMING_SNAKE_CASE_			( self			:		str):


         """simple docstring"""







         a :				Any				=      SampleObject()

         a :				int				=      {k: v for k, v in config.config.items() if not k.startswith('_')}

         # make sure that default config has all keys in `_use_default_values`
         assert set(config_dict.keys()) == set(config.config._use_default_values)

         with tempfile.TemporaryDirectory() as tmpdirname:
              config.save_config(a__)

              # now loading it with SampleObject2 should put f into `_use_default_values`
              a :				str				=      SampleObjecta.from_config(a__)

              assert "f" in config._use_default_values
              assert config.f == [1, 3]

         # now loading the config, should **NOT** use [1, 3] for `f`, but the default [1, 4] value
         # **BECAUSE** it is part of `config._use_default_values`
         a :				List[Any]				=      SampleObjecta.from_config(config.config)
         assert new_config.f == [5, 4]

         config.config._use_default_values.pop()
         a :				Optional[int]				=      SampleObjecta.from_config(config.config)
         assert new_config_a.f == [1, 3]

         # Nevertheless "e" should still be correctly loaded to [1, 3] from SampleObject2 instead of defaulting to [1, 5]
         assert new_config_a.e == [1, 3]








"""simple docstring"""


# 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[Any]						,  snake_case   :					List[Any]						,  snake_case   :					Optional[int]=0		)->		Tuple:
       # Format the message.
       if name is None:
              _lowerCamelCase    	=					None
       else:
              _lowerCamelCase    	=					'.' * max(0						,  spaces - 2		) + '# {:' + str(50 - spaces		) + 's}'
              _lowerCamelCase    	=					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   :					Optional[int]						,  snake_case   :					Optional[int]						,  snake_case   :					List[str]						,  snake_case   :					Union[str, Any]						,  snake_case   :					int		)->		List[str]:
       # 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.
       _lowerCamelCase    	=					param.size()
       if checkpoint_version == 1.0:
              # version 1.0 stores [num_heads * hidden_size * num_splits, :]
              _lowerCamelCase    	=					(num_heads, hidden_size, num_splits) + input_shape[1:]
              _lowerCamelCase    	=					param.view(*snake_case		)
              _lowerCamelCase    	=					param.transpose(0						,  2		)
              _lowerCamelCase    	=					param.transpose(1						,  2		).contiguous()
       elif checkpoint_version >= 2.0:
              # other versions store [num_heads * num_splits * hidden_size, :]
              _lowerCamelCase    	=					(num_heads, num_splits, hidden_size) + input_shape[1:]
              _lowerCamelCase    	=					param.view(*snake_case		)
              _lowerCamelCase    	=					param.transpose(0						,  1		).contiguous()
       _lowerCamelCase    	=					param.view(*snake_case		)
       return param





def SCREAMING_SNAKE_CASE_   (       snake_case   :					Union[str, Any]						,  snake_case   :					List[Any]						,  snake_case   :					List[Any]		)->		Optional[Any]:
       # The converted output model.
       _lowerCamelCase    	=					{}

       # old versions did not store training args
       _lowerCamelCase    	=					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))

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

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

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

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

       # The position embeddings.
       _lowerCamelCase    	=					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.
       _lowerCamelCase    	=					pos_embeddings

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

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

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

       # The simple map of names for "automated" rules.
       _lowerCamelCase    	=					{
           '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.
       _lowerCamelCase    	=					None

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

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

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

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

              # For layernorm(s), simply store the layer norm.
              if op_name.endswith('layernorm'		):
                     _lowerCamelCase    	=					'attention.ln' if op_name.startswith('input'		) else 'ln'
                     _lowerCamelCase    	=					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, ""

                     _lowerCamelCase    	=					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.
                     _lowerCamelCase    	=					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.
                     _lowerCamelCase    	=					attention_qkv_weight[0 * config.hidden_size : 1 * config.hidden_size, :]
                     _lowerCamelCase    	=					attention_qkv_weight[1 * config.hidden_size : 2 * config.hidden_size, :]
                     _lowerCamelCase    	=					attention_qkv_weight[2 * config.hidden_size : 3 * config.hidden_size, :]

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

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

                     # Clear the stored tensor.
                     _lowerCamelCase    	=					None

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

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

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

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

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

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

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

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

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

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

       # It should be done!
       return output_state_dict





def SCREAMING_SNAKE_CASE_   (       )->		Tuple:
       # Create the argument parser.
       _lowerCamelCase    	=					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.'						,  )
       _lowerCamelCase    	=					parser.parse_args()

       # Extract the basename.
       _lowerCamelCase    	=					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:
                            _lowerCamelCase    	=					torch.load(snake_case						,  map_location='cpu'		)
       else:
              _lowerCamelCase    	=					torch.load(args.path_to_checkpoint						,  map_location='cpu'		)

       if args.config_file == "":
              # Default config of megatron-bert 345m
              _lowerCamelCase    	=					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
              _lowerCamelCase    	=					input_state_dict['model']['lm_head']['bias'].numel()
       else:
              _lowerCamelCase    	=					MegatronBertConfig.from_json_file(args.config_file		)

       # Convert.
       print('Converting'		)
       _lowerCamelCase    	=					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.
       _lowerCamelCase    	=					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()

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






