


'''simple docstring'''







from __future__ import annotations

import unittest

from transformers import DistilBertConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
       import tensorflow as tf

       from transformers.models.distilbert.modeling_tf_distilbert import (
           TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
           TFDistilBertForMaskedLM,
           TFDistilBertForMultipleChoice,
           TFDistilBertForQuestionAnswering,
           TFDistilBertForSequenceClassification,
           TFDistilBertForTokenClassification,
           TFDistilBertModel,
       )
class 	UpperCamelCase__					:



     """simple docstring"""






     def __init__( self      ,	snake_case__      ,	):

            '''simple docstring'''

            _lowerCAmelCase    :							Optional[int]   			=       parent
            _lowerCAmelCase    :							Any   			=       13
            _lowerCAmelCase    :							Dict   			=       7
            _lowerCAmelCase    :							Any   			=       True
            _lowerCAmelCase    :							Union[str, Any]   			=       True
            _lowerCAmelCase    :							Union[str, Any]   			=       False
            _lowerCAmelCase    :							List[Any]   			=       True
            _lowerCAmelCase    :							int   			=       99
            _lowerCAmelCase    :							int   			=       32
            _lowerCAmelCase    :							List[Any]   			=       2
            _lowerCAmelCase    :							List[Any]   			=       4
            _lowerCAmelCase    :							Dict   			=       37
            _lowerCAmelCase    :							str   			=       'gelu'
            _lowerCAmelCase    :							str   			=       0.1
            _lowerCAmelCase    :							Any   			=       0.1
            _lowerCAmelCase    :							Union[str, Any]   			=       512
            _lowerCAmelCase    :							int   			=       16
            _lowerCAmelCase    :							List[Any]   			=       2
            _lowerCAmelCase    :							Optional[int]   			=       0.02
            _lowerCAmelCase    :							int   			=       3
            _lowerCAmelCase    :							Optional[int]   			=       4
            _lowerCAmelCase    :							str   			=       None







     def a	( self   ):

            '''simple docstring'''

            _lowerCAmelCase    :							List[Any]   			=       ids_tensor([self.batch_size, self.seq_length]      ,	self.vocab_size   )

            _lowerCAmelCase    :							List[Any]   			=       None
            if self.use_input_mask:
                   _lowerCAmelCase    :							Tuple   			=       random_attention_mask([self.batch_size, self.seq_length]   )

            _lowerCAmelCase    :							List[Any]   			=       None
            _lowerCAmelCase    :							Union[str, Any]   			=       None
            _lowerCAmelCase    :							str   			=       None
            if self.use_labels:
                   _lowerCAmelCase    :							List[str]   			=       ids_tensor([self.batch_size]      ,	self.type_sequence_label_size   )
                   _lowerCAmelCase    :							Optional[Any]   			=       ids_tensor([self.batch_size, self.seq_length]      ,	self.num_labels   )
                   _lowerCAmelCase    :							List[str]   			=       ids_tensor([self.batch_size]      ,	self.num_choices   )

            _lowerCAmelCase    :							str   			=       DistilBertConfig(
                vocab_size=self.vocab_size      ,	dim=self.hidden_size      ,	n_layers=self.num_hidden_layers      ,	n_heads=self.num_attention_heads      ,	hidden_dim=self.intermediate_size      ,	hidden_act=self.hidden_act      ,	dropout=self.hidden_dropout_prob      ,	attention_dropout=self.attention_probs_dropout_prob      ,	max_position_embeddings=self.max_position_embeddings      ,	initializer_range=self.initializer_range      ,	)

            return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels







     def a	( self      ,	snake_case__      ,	snake_case__      ,	snake_case__      ,	snake_case__      ,	snake_case__      ,	snake_case__   ):

            '''simple docstring'''

            _lowerCAmelCase    :							Optional[int]   			=       TFDistilBertModel(config=snake_case__   )
            _lowerCAmelCase    :							List[Any]   			=       {'input_ids': input_ids, 'attention_mask': input_mask}

            _lowerCAmelCase    :							List[str]   			=       model(snake_case__   )

            _lowerCAmelCase    :							List[str]   			=       [input_ids, input_mask]

            _lowerCAmelCase    :							Tuple   			=       model(snake_case__   )

            self.parent.assertEqual(result.last_hidden_state.shape      ,	(self.batch_size, self.seq_length, self.hidden_size)   )







     def a	( self      ,	snake_case__      ,	snake_case__      ,	snake_case__      ,	snake_case__      ,	snake_case__      ,	snake_case__   ):

            '''simple docstring'''

            _lowerCAmelCase    :							Tuple   			=       TFDistilBertForMaskedLM(config=snake_case__   )
            _lowerCAmelCase    :							List[Any]   			=       {'input_ids': input_ids, 'attention_mask': input_mask}
            _lowerCAmelCase    :							List[Any]   			=       model(snake_case__   )
            self.parent.assertEqual(result.logits.shape      ,	(self.batch_size, self.seq_length, self.vocab_size)   )







     def a	( self      ,	snake_case__      ,	snake_case__      ,	snake_case__      ,	snake_case__      ,	snake_case__      ,	snake_case__   ):

            '''simple docstring'''

            _lowerCAmelCase    :							int   			=       TFDistilBertForQuestionAnswering(config=snake_case__   )
            _lowerCAmelCase    :							Tuple   			=       {
                'input_ids': input_ids,
                'attention_mask': input_mask,
            }
            _lowerCAmelCase    :							Union[str, Any]   			=       model(snake_case__   )
            self.parent.assertEqual(result.start_logits.shape      ,	(self.batch_size, self.seq_length)   )
            self.parent.assertEqual(result.end_logits.shape      ,	(self.batch_size, self.seq_length)   )







     def a	( self      ,	snake_case__      ,	snake_case__      ,	snake_case__      ,	snake_case__      ,	snake_case__      ,	snake_case__   ):

            '''simple docstring'''

            _lowerCAmelCase    :							List[Any]   			=       self.num_labels
            _lowerCAmelCase    :							List[str]   			=       TFDistilBertForSequenceClassification(snake_case__   )
            _lowerCAmelCase    :							int   			=       {'input_ids': input_ids, 'attention_mask': input_mask}
            _lowerCAmelCase    :							Any   			=       model(snake_case__   )
            self.parent.assertEqual(result.logits.shape      ,	(self.batch_size, self.num_labels)   )







     def a	( self      ,	snake_case__      ,	snake_case__      ,	snake_case__      ,	snake_case__      ,	snake_case__      ,	snake_case__   ):

            '''simple docstring'''

            _lowerCAmelCase    :							Optional[Any]   			=       self.num_choices
            _lowerCAmelCase    :							Optional[Any]   			=       TFDistilBertForMultipleChoice(snake_case__   )
            _lowerCAmelCase    :							Tuple   			=       tf.tile(tf.expand_dims(snake_case__      ,	1   )      ,	(1, self.num_choices, 1)   )
            _lowerCAmelCase    :							List[Any]   			=       tf.tile(tf.expand_dims(snake_case__      ,	1   )      ,	(1, self.num_choices, 1)   )
            _lowerCAmelCase    :							List[str]   			=       {
                'input_ids': multiple_choice_inputs_ids,
                'attention_mask': multiple_choice_input_mask,
            }
            _lowerCAmelCase    :							int   			=       model(snake_case__   )
            self.parent.assertEqual(result.logits.shape      ,	(self.batch_size, self.num_choices)   )







     def a	( self      ,	snake_case__      ,	snake_case__      ,	snake_case__      ,	snake_case__      ,	snake_case__      ,	snake_case__   ):

            '''simple docstring'''

            _lowerCAmelCase    :							List[Any]   			=       self.num_labels
            _lowerCAmelCase    :							Dict   			=       TFDistilBertForTokenClassification(snake_case__   )
            _lowerCAmelCase    :							Optional[Any]   			=       {'input_ids': input_ids, 'attention_mask': input_mask}
            _lowerCAmelCase    :							List[str]   			=       model(snake_case__   )
            self.parent.assertEqual(result.logits.shape      ,	(self.batch_size, self.seq_length, self.num_labels)   )







     def a	( self   ):

            '''simple docstring'''

            _lowerCAmelCase    :							Any   			=       self.prepare_config_and_inputs()
            ((_lowerCAmelCase)			,     (_lowerCAmelCase)			,     (_lowerCAmelCase)			,     (_lowerCAmelCase)			,     (_lowerCAmelCase)			,     (_lowerCAmelCase))    :							str   			=       config_and_inputs
            _lowerCAmelCase    :							Tuple   			=       {'input_ids': input_ids, 'attention_mask': input_mask}
            return config, inputs_dict
@require_tf
class 	UpperCamelCase__					( SCREAMING_SNAKE_CASE_    ,					SCREAMING_SNAKE_CASE_    ,					unittest.TestCase	):



     """simple docstring"""






     __magic_name__    			=    (
         (
             TFDistilBertModel,
             TFDistilBertForMaskedLM,
             TFDistilBertForQuestionAnswering,
             TFDistilBertForSequenceClassification,
             TFDistilBertForTokenClassification,
             TFDistilBertForMultipleChoice,
         )
         if is_tf_available()
         else None
     )
     __magic_name__    			=    (
         {
             "feature-extraction": TFDistilBertModel,
             "fill-mask": TFDistilBertForMaskedLM,
             "question-answering": TFDistilBertForQuestionAnswering,
             "text-classification": TFDistilBertForSequenceClassification,
             "token-classification": TFDistilBertForTokenClassification,
             "zero-shot": TFDistilBertForSequenceClassification,
         }
         if is_tf_available()
         else {}
     )
     __magic_name__    			=    False
     __magic_name__    			=    False







     def a	( self   ):

            '''simple docstring'''

            _lowerCAmelCase    :							Dict   			=       TFDistilBertModelTester(self   )
            _lowerCAmelCase    :							Any   			=       ConfigTester(self      ,	config_class=snake_case__      ,	dim=37   )







     def a	( self   ):

            '''simple docstring'''

            self.config_tester.run_common_tests()







     def a	( self   ):

            '''simple docstring'''

            _lowerCAmelCase    :							List[str]   			=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_distilbert_model(*snake_case__   )







     def a	( self   ):

            '''simple docstring'''

            _lowerCAmelCase    :							int   			=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_distilbert_for_masked_lm(*snake_case__   )







     def a	( self   ):

            '''simple docstring'''

            _lowerCAmelCase    :							Optional[int]   			=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_distilbert_for_question_answering(*snake_case__   )







     def a	( self   ):

            '''simple docstring'''

            _lowerCAmelCase    :							Union[str, Any]   			=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_distilbert_for_sequence_classification(*snake_case__   )







     def a	( self   ):

            '''simple docstring'''

            _lowerCAmelCase    :							List[Any]   			=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_distilbert_for_multiple_choice(*snake_case__   )







     def a	( self   ):

            '''simple docstring'''

            _lowerCAmelCase    :							Optional[int]   			=       self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_distilbert_for_token_classification(*snake_case__   )







     @slow
     def a	( self   ):

            '''simple docstring'''

            for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]   ):
                   _lowerCAmelCase    :							Any   			=       TFDistilBertModel.from_pretrained(snake_case__   )
                   self.assertIsNotNone(snake_case__   )
@require_tf
class 	UpperCamelCase__					( unittest.TestCase	):



     """simple docstring"""






     @slow
     def a	( self   ):

            '''simple docstring'''

            _lowerCAmelCase    :							Optional[int]   			=       TFDistilBertModel.from_pretrained('distilbert-base-uncased'   )
            _lowerCAmelCase    :							List[str]   			=       tf.constant([[0, 1, 2, 3, 4, 5]]   )
            _lowerCAmelCase    :							Dict   			=       model(snake_case__   )[0]

            _lowerCAmelCase    :							int   			=       [1, 6, 768]
            self.assertEqual(output.shape      ,	snake_case__   )

            _lowerCAmelCase    :							List[Any]   			=       tf.constant(
                [
                    [
                        [0.1926_1885, -0.1373_2955, 0.411_9799],
                        [0.2215_0156, -0.0742_2661, 0.3903_7204],
                        [0.2275_6018, -0.089_6414, 0.370_1467],
                    ]
                ]   )
            tf.debugging.assert_near(output[:, :3, :3]      ,	snake_case__      ,	atol=1E-4   )





'''simple docstring'''







import argparse
import json
import os

import torch
from transformers.file_utils import has_file

from diffusers import UNetaDConditionModel, UNetaDModel


lowerCAmelCase    :     Tuple					=						False
lowerCAmelCase    :     str					=						True
lowerCAmelCase    :     List[Any]					=						False


if __name__ == "__main__":
       lowerCAmelCase    :     Any					=						argparse.ArgumentParser()

       parser.add_argument(
           """--repo_path""",
           default=None,
           type=str,
           required=True,
           help="""The config json file corresponding to the architecture.""",
       )

       parser.add_argument("""--dump_path""", default=None, type=str, required=True, help="""Path to the output model.""")

       lowerCAmelCase    :     Optional[int]					=						parser.parse_args()

       lowerCAmelCase    :     int					=						{
           """image_size""": """sample_size""",
           """num_res_blocks""": """layers_per_block""",
           """block_channels""": """block_out_channels""",
           """down_blocks""": """down_block_types""",
           """up_blocks""": """up_block_types""",
           """downscale_freq_shift""": """freq_shift""",
           """resnet_num_groups""": """norm_num_groups""",
           """resnet_act_fn""": """act_fn""",
           """resnet_eps""": """norm_eps""",
           """num_head_channels""": """attention_head_dim""",
       }

       lowerCAmelCase    :     int					=						{
           """time_steps""": """time_proj""",
           """mid""": """mid_block""",
           """downsample_blocks""": """down_blocks""",
           """upsample_blocks""": """up_blocks""",
       }

       lowerCAmelCase    :     Optional[Any]					=						"""""" if has_file(args.repo_path, """config.json""") else """unet"""

       with open(os.path.join(args.repo_path, subfolder, """config.json"""), """r""", encoding="""utf-8""") as reader:
              lowerCAmelCase    :     int					=						reader.read()
              lowerCAmelCase    :     List[str]					=						json.loads(text)

       if do_only_config:
              for key in config_parameters_to_change.keys():
                     config.pop(key, None)

       if has_file(args.repo_path, """config.json"""):
              lowerCAmelCase    :     str					=						UNetaDModel(**config)
       else:
              lowerCAmelCase    :     Union[str, Any]					=						UNetaDConditionModel if """ldm-text2im-large-256""" in args.repo_path else UNetaDModel
              lowerCAmelCase    :     Dict					=						class_name(**config)

       if do_only_config:
              model.save_config(os.path.join(args.repo_path, subfolder))

       lowerCAmelCase    :     Union[str, Any]					=						dict(model.config)

       if do_only_renaming:
              for key, value in config_parameters_to_change.items():
                     if key in config:
                            lowerCAmelCase    :     str					=						config[key]
                            del config[key]

              lowerCAmelCase    :     Optional[int]					=						[k.replace("""UNetRes""", """""") for k in config["""down_block_types"""]]
              lowerCAmelCase    :     Dict					=						[k.replace("""UNetRes""", """""") for k in config["""up_block_types"""]]

       if do_only_weights:
              lowerCAmelCase    :     Tuple					=						torch.load(os.path.join(args.repo_path, subfolder, """diffusion_pytorch_model.bin"""))

              lowerCAmelCase    :     str					=						{}
              for param_key, param_value in state_dict.items():
                     if param_key.endswith(""".op.bias""") or param_key.endswith(""".op.weight"""):
                            continue
                     lowerCAmelCase    :     str					=						False
                     for key, new_key in key_parameters_to_change.items():
                            if not has_changed and param_key.split(""".""")[0] == key:
                                   lowerCAmelCase    :     Dict					=						param_value
                                   lowerCAmelCase    :     Tuple					=						True
                     if not has_changed:
                            lowerCAmelCase    :     Tuple					=						param_value

              model.load_state_dict(new_state_dict)
              model.save_pretrained(os.path.join(args.repo_path, subfolder))

