

'''simple docstring'''


import unittest

from diffusers.models.unet_ad_blocks import *  # noqa F403
from diffusers.utils import torch_device

from .test_unet_blocks_common import UNetBlockTesterMixin




class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    DownBlockaD  # noqa F405
						__UpperCAmelCase			     =    """down"""







						def       __magic_name__							(							self					:      int  ):



								'''simple docstring'''



								snake_case__						:       Optional[int]       						=  [-0.0_2_3_2, -0.9_8_6_9, 0.8_0_5_4, -0.0_6_3_7, -0.1_6_8_8, -1.4_2_6_4, 0.4_4_7_0, -1.3_3_9_4, 0.0_9_0_4]
								super().test_output(snake_case_  )





class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    ResnetDownsampleBlockaD  # noqa F405
						__UpperCAmelCase			     =    """down"""







						def       __magic_name__							(							self					:      List[Any]  ):



								'''simple docstring'''



								snake_case__						:       Optional[Any]       						=  [0.0_7_1_0, 0.2_4_1_0, -0.7_3_2_0, -1.0_7_5_7, -1.1_3_4_3, 0.3_5_4_0, -0.0_1_3_3, -0.2_5_7_6, 0.0_9_4_8]
								super().test_output(snake_case_  )





class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    AttnDownBlockaD  # noqa F405
						__UpperCAmelCase			     =    """down"""







						def       __magic_name__							(							self					:      List[Any]  ):



								'''simple docstring'''



								snake_case__						:       Dict       						=  [0.0_6_3_6, 0.8_9_6_4, -0.6_2_3_4, -1.0_1_3_1, 0.0_8_4_4, 0.4_9_3_5, 0.3_4_3_7, 0.0_9_1_1, -0.2_9_5_7]
								super().test_output(snake_case_  )





class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    CrossAttnDownBlockaD  # noqa F405
						__UpperCAmelCase			     =    """down"""







						def       __magic_name__							(							self					:      str  ):



								'''simple docstring'''



								snake_case__						,       snake_case__						:       Union[str, Any]       						=  super().prepare_init_args_and_inputs_for_common()
								snake_case__						:       List[str]       						=  3_2
								return init_dict, inputs_dict







						def       __magic_name__							(							self					:      Optional[Any]  ):



								'''simple docstring'''



								snake_case__						:       str       						=  [0.2_2_3_8, -0.7_3_9_6, -0.2_2_5_5, -0.3_8_2_9, 0.1_9_2_5, 1.1_6_6_5, 0.0_6_0_3, -0.7_2_9_5, 0.1_9_8_3]
								super().test_output(snake_case_  )





class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    SimpleCrossAttnDownBlockaD  # noqa F405
						__UpperCAmelCase			     =    """down"""







						@property
						def       __magic_name__							(							self					:      List[Any]  ):



								'''simple docstring'''



								return super().get_dummy_input(include_encoder_hidden_states=snake_case_  )







						def       __magic_name__							(							self					:      Dict  ):



								'''simple docstring'''



								snake_case__						,       snake_case__						:       List[str]       						=  super().prepare_init_args_and_inputs_for_common()
								snake_case__						:       Union[str, Any]       						=  3_2
								return init_dict, inputs_dict







						@unittest.skipIf(torch_device == '''mps'''			,			'''MPS result is not consistent'''  )
						def       __magic_name__							(							self					:      Tuple  ):



								'''simple docstring'''



								snake_case__						:       Optional[int]       						=  [0.7_9_2_1, -0.0_9_9_2, -0.1_9_6_2, -0.7_6_9_5, -0.4_2_4_2, 0.7_8_0_4, 0.4_7_3_7, 0.2_7_6_5, 0.3_3_3_8]
								super().test_output(snake_case_  )





class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    SkipDownBlockaD  # noqa F405
						__UpperCAmelCase			     =    """down"""







						@property
						def       __magic_name__							(							self					:      Optional[int]  ):



								'''simple docstring'''



								return super().get_dummy_input(include_skip_sample=snake_case_  )







						def       __magic_name__							(							self					:      Optional[int]  ):



								'''simple docstring'''



								snake_case__						:       int       						=  [-0.0_8_4_5, -0.2_0_8_7, -0.2_4_6_5, 0.0_9_7_1, 0.1_9_0_0, -0.0_4_8_4, 0.2_6_6_4, 0.4_1_7_9, 0.5_0_6_9]
								super().test_output(snake_case_  )





class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    AttnSkipDownBlockaD  # noqa F405
						__UpperCAmelCase			     =    """down"""







						@property
						def       __magic_name__							(							self					:      Dict  ):



								'''simple docstring'''



								return super().get_dummy_input(include_skip_sample=snake_case_  )







						def       __magic_name__							(							self					:      List[str]  ):



								'''simple docstring'''



								snake_case__						:       Any       						=  [0.5_5_3_9, 0.1_6_0_9, 0.4_9_2_4, 0.0_5_3_7, -0.1_9_9_5, 0.4_0_5_0, 0.0_9_7_9, -0.2_7_2_1, -0.0_6_4_2]
								super().test_output(snake_case_  )





class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    DownEncoderBlockaD  # noqa F405
						__UpperCAmelCase			     =    """down"""







						@property
						def       __magic_name__							(							self					:      int  ):



								'''simple docstring'''



								return super().get_dummy_input(include_temb=snake_case_  )







						def       __magic_name__							(							self					:      Optional[Any]  ):



								'''simple docstring'''



								snake_case__						:       List[str]       						=  {
								    '''in_channels''': 3_2,
								    '''out_channels''': 3_2,
								}
								snake_case__						:       Tuple       						=  self.dummy_input
								return init_dict, inputs_dict







						def       __magic_name__							(							self					:      List[Any]  ):



								'''simple docstring'''



								snake_case__						:       Optional[Any]       						=  [1.1_1_0_2, 0.5_3_0_2, 0.4_8_7_2, -0.0_0_2_3, -0.8_0_4_2, 0.0_4_8_3, -0.3_4_8_9, -0.5_6_3_2, 0.7_6_2_6]
								super().test_output(snake_case_  )





class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    AttnDownEncoderBlockaD  # noqa F405
						__UpperCAmelCase			     =    """down"""







						@property
						def       __magic_name__							(							self					:      str  ):



								'''simple docstring'''



								return super().get_dummy_input(include_temb=snake_case_  )







						def       __magic_name__							(							self					:      List[str]  ):



								'''simple docstring'''



								snake_case__						:       List[Any]       						=  {
								    '''in_channels''': 3_2,
								    '''out_channels''': 3_2,
								}
								snake_case__						:       Dict       						=  self.dummy_input
								return init_dict, inputs_dict







						def       __magic_name__							(							self					:      Optional[int]  ):



								'''simple docstring'''



								snake_case__						:       Dict       						=  [0.8_9_6_6, -0.1_4_8_6, 0.8_5_6_8, 0.8_1_4_1, -0.9_0_4_6, -0.1_3_4_2, -0.0_9_7_2, -0.7_4_1_7, 0.1_5_3_8]
								super().test_output(snake_case_  )





class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    UNetMidBlockaD  # noqa F405
						__UpperCAmelCase			     =    """mid"""







						def       __magic_name__							(							self					:      int  ):



								'''simple docstring'''



								snake_case__						:       int       						=  {
								    '''in_channels''': 3_2,
								    '''temb_channels''': 1_2_8,
								}
								snake_case__						:       Optional[int]       						=  self.dummy_input
								return init_dict, inputs_dict







						def       __magic_name__							(							self					:      Union[str, Any]  ):



								'''simple docstring'''



								snake_case__						:       Union[str, Any]       						=  [-0.1_0_6_2, 1.7_2_4_8, 0.3_4_9_4, 1.4_5_6_9, -0.0_9_1_0, -1.2_4_2_1, -0.9_9_8_4, 0.6_7_3_6, 1.0_0_2_8]
								super().test_output(snake_case_  )





class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    UNetMidBlockaDCrossAttn  # noqa F405
						__UpperCAmelCase			     =    """mid"""







						def       __magic_name__							(							self					:      Union[str, Any]  ):



								'''simple docstring'''



								snake_case__						,       snake_case__						:       Dict       						=  super().prepare_init_args_and_inputs_for_common()
								snake_case__						:       List[Any]       						=  3_2
								return init_dict, inputs_dict







						def       __magic_name__							(							self					:      Dict  ):



								'''simple docstring'''



								snake_case__						:       List[Any]       						=  [0.0_1_8_7, 2.4_2_2_0, 0.4_4_8_4, 1.1_2_0_3, -0.6_1_2_1, -1.5_1_2_2, -0.8_2_7_0, 0.7_8_5_1, 1.8_3_3_5]
								super().test_output(snake_case_  )





class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    UNetMidBlockaDSimpleCrossAttn  # noqa F405
						__UpperCAmelCase			     =    """mid"""







						@property
						def       __magic_name__							(							self					:      Dict  ):



								'''simple docstring'''



								return super().get_dummy_input(include_encoder_hidden_states=snake_case_  )







						def       __magic_name__							(							self					:      Any  ):



								'''simple docstring'''



								snake_case__						,       snake_case__						:       Tuple       						=  super().prepare_init_args_and_inputs_for_common()
								snake_case__						:       List[str]       						=  3_2
								return init_dict, inputs_dict







						def       __magic_name__							(							self					:      List[Any]  ):



								'''simple docstring'''



								snake_case__						:       Dict       						=  [0.7_1_4_3, 1.9_9_7_4, 0.5_4_4_8, 1.3_9_7_7, 0.1_2_8_2, -1.1_2_3_7, -1.4_2_3_8, 0.5_5_3_0, 0.8_8_8_0]
								super().test_output(snake_case_  )





class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    UpBlockaD  # noqa F405
						__UpperCAmelCase			     =    """up"""







						@property
						def       __magic_name__							(							self					:      Any  ):



								'''simple docstring'''



								return super().get_dummy_input(include_res_hidden_states_tuple=snake_case_  )







						def       __magic_name__							(							self					:      str  ):



								'''simple docstring'''



								snake_case__						:       Optional[int]       						=  [-0.2_0_4_1, -0.4_1_6_5, -0.3_0_2_2, 0.0_0_4_1, -0.6_6_2_8, -0.7_0_5_3, 0.1_9_2_8, -0.0_3_2_5, 0.0_5_2_3]
								super().test_output(snake_case_  )





class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    ResnetUpsampleBlockaD  # noqa F405
						__UpperCAmelCase			     =    """up"""







						@property
						def       __magic_name__							(							self					:      str  ):



								'''simple docstring'''



								return super().get_dummy_input(include_res_hidden_states_tuple=snake_case_  )







						def       __magic_name__							(							self					:      Tuple  ):



								'''simple docstring'''



								snake_case__						:       Any       						=  [0.2_2_8_7, 0.3_5_4_9, -0.1_3_4_6, 0.4_7_9_7, -0.1_7_1_5, -0.9_6_4_9, 0.7_3_0_5, -0.5_8_6_4, -0.6_2_4_4]
								super().test_output(snake_case_  )





class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    CrossAttnUpBlockaD  # noqa F405
						__UpperCAmelCase			     =    """up"""







						@property
						def       __magic_name__							(							self					:      List[Any]  ):



								'''simple docstring'''



								return super().get_dummy_input(include_res_hidden_states_tuple=snake_case_  )







						def       __magic_name__							(							self					:      Dict  ):



								'''simple docstring'''



								snake_case__						,       snake_case__						:       int       						=  super().prepare_init_args_and_inputs_for_common()
								snake_case__						:       Union[str, Any]       						=  3_2
								return init_dict, inputs_dict







						def       __magic_name__							(							self					:      List[Any]  ):



								'''simple docstring'''



								snake_case__						:       str       						=  [-0.1_4_0_3, -0.3_5_1_5, -0.0_4_2_0, -0.1_4_2_5, 0.3_1_6_7, 0.5_0_9_4, -0.2_1_8_1, 0.5_9_3_1, 0.5_5_8_2]
								super().test_output(snake_case_  )





class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    SimpleCrossAttnUpBlockaD  # noqa F405
						__UpperCAmelCase			     =    """up"""







						@property
						def       __magic_name__							(							self					:      Union[str, Any]  ):



								'''simple docstring'''



								return super().get_dummy_input(include_res_hidden_states_tuple=snake_case_			,			include_encoder_hidden_states=snake_case_  )







						def       __magic_name__							(							self					:      Optional[Any]  ):



								'''simple docstring'''



								snake_case__						,       snake_case__						:       Optional[int]       						=  super().prepare_init_args_and_inputs_for_common()
								snake_case__						:       Dict       						=  3_2
								return init_dict, inputs_dict







						def       __magic_name__							(							self					:      Any  ):



								'''simple docstring'''



								snake_case__						:       Optional[Any]       						=  [0.2_6_4_5, 0.1_4_8_0, 0.0_9_0_9, 0.8_0_4_4, -0.9_7_5_8, -0.9_0_8_3, 0.0_9_9_4, -1.1_4_5_3, -0.7_4_0_2]
								super().test_output(snake_case_  )





class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    AttnUpBlockaD  # noqa F405
						__UpperCAmelCase			     =    """up"""







						@property
						def       __magic_name__							(							self					:      int  ):



								'''simple docstring'''



								return super().get_dummy_input(include_res_hidden_states_tuple=snake_case_  )







						@unittest.skipIf(torch_device == '''mps'''			,			'''MPS result is not consistent'''  )
						def       __magic_name__							(							self					:      Optional[int]  ):



								'''simple docstring'''



								snake_case__						:       Any       						=  [0.0_9_7_9, 0.1_3_2_6, 0.0_0_2_1, 0.0_6_5_9, 0.2_2_4_9, 0.0_0_5_9, 0.1_1_3_2, 0.5_9_5_2, 0.1_0_3_3]
								super().test_output(snake_case_  )





class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    SkipUpBlockaD  # noqa F405
						__UpperCAmelCase			     =    """up"""







						@property
						def       __magic_name__							(							self					:      Optional[int]  ):



								'''simple docstring'''



								return super().get_dummy_input(include_res_hidden_states_tuple=snake_case_  )







						def       __magic_name__							(							self					:      Optional[int]  ):



								'''simple docstring'''



								snake_case__						:       Union[str, Any]       						=  [-0.0_8_9_3, -0.1_2_3_4, -0.1_5_0_6, -0.0_3_3_2, 0.0_1_2_3, -0.0_2_1_1, 0.0_5_6_6, 0.0_1_4_3, 0.0_3_6_2]
								super().test_output(snake_case_  )





class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    AttnSkipUpBlockaD  # noqa F405
						__UpperCAmelCase			     =    """up"""







						@property
						def       __magic_name__							(							self					:      Optional[Any]  ):



								'''simple docstring'''



								return super().get_dummy_input(include_res_hidden_states_tuple=snake_case_  )







						def       __magic_name__							(							self					:      Dict  ):



								'''simple docstring'''



								snake_case__						:       Optional[int]       						=  [0.0_3_6_1, 0.0_6_1_7, 0.2_7_8_7, -0.0_3_5_0, 0.0_3_4_2, 0.3_4_2_1, -0.0_8_4_3, 0.0_9_1_3, 0.3_0_1_5]
								super().test_output(snake_case_  )





class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    UpDecoderBlockaD  # noqa F405
						__UpperCAmelCase			     =    """up"""







						@property
						def       __magic_name__							(							self					:      List[Any]  ):



								'''simple docstring'''



								return super().get_dummy_input(include_temb=snake_case_  )







						def       __magic_name__							(							self					:      Dict  ):



								'''simple docstring'''



								snake_case__						:       List[str]       						=  {'''in_channels''': 3_2, '''out_channels''': 3_2}

								snake_case__						:       List[str]       						=  self.dummy_input
								return init_dict, inputs_dict







						def       __magic_name__							(							self					:      str  ):



								'''simple docstring'''



								snake_case__						:       List[str]       						=  [0.4_4_0_4, 0.1_9_9_8, -0.9_8_8_6, -0.3_3_2_0, -0.3_1_2_8, -0.7_0_3_4, -0.6_9_5_5, -0.2_3_3_8, -0.3_1_3_7]
								super().test_output(snake_case_  )





class     a			(				SCREAMING_SNAKE_CASE			,			unittest.TestCase       ):






						"""simple docstring"""




						__UpperCAmelCase			     =    AttnUpDecoderBlockaD  # noqa F405
						__UpperCAmelCase			     =    """up"""







						@property
						def       __magic_name__							(							self					:      Optional[Any]  ):



								'''simple docstring'''



								return super().get_dummy_input(include_temb=snake_case_  )







						def       __magic_name__							(							self					:      List[Any]  ):



								'''simple docstring'''



								snake_case__						:       Any       						=  {'''in_channels''': 3_2, '''out_channels''': 3_2}

								snake_case__						:       Union[str, Any]       						=  self.dummy_input
								return init_dict, inputs_dict







						def       __magic_name__							(							self					:      int  ):



								'''simple docstring'''



								snake_case__						:       List[str]       						=  [0.6_7_3_8, 0.4_4_9_1, 0.1_0_5_5, 1.0_7_1_0, 0.7_3_1_6, 0.3_3_3_9, 0.3_3_5_2, 0.1_0_2_3, 0.3_5_6_8]
								super().test_output(snake_case_  )






'''simple docstring'''


import argparse
import os
import shutil

import torch
from emmental.modules import MagnitudeBinarizer, ThresholdBinarizer, TopKBinarizer




def 	_a       (				__lowerCAmelCase		:      Dict     ):

		"""simple docstring"""

		snake_case__						:       int       						=  args.pruning_method
		snake_case__						:       Dict       						=  args.threshold

		snake_case__						:       str       						=  args.model_name_or_path.rstrip('''/'''     )
		snake_case__						:       int       						=  args.target_model_path

		print(F"""Load fine-pruned model from {model_name_or_path}"""     )
		snake_case__						:       Dict       						=  torch.load(os.path.join(__lowerCAmelCase   ,							'''pytorch_model.bin'''     )     )
		snake_case__						:       Any       						=  {}

		for name, tensor in model.items():
				if "embeddings" in name or "LayerNorm" in name or "pooler" in name:
						snake_case__						:       List[Any]       						=  tensor
						print(F"""Copied layer {name}"""     )
				elif "classifier" in name or "qa_output" in name:
						snake_case__						:       int       						=  tensor
						print(F"""Copied layer {name}"""     )
				elif "bias" in name:
						snake_case__						:       Union[str, Any]       						=  tensor
						print(F"""Copied layer {name}"""     )
				else:
						if pruning_method == "magnitude":
								snake_case__						:       Dict       						=  MagnitudeBinarizer.apply(inputs=__lowerCAmelCase   ,							threshold=__lowerCAmelCase     )
								snake_case__						:       Optional[int]       						=  tensor * mask
								print(F"""Pruned layer {name}"""     )
						elif pruning_method == "topK":
								if "mask_scores" in name:
										continue
								snake_case__						:       List[str]       						=  name[:-6]
								snake_case__						:       List[Any]       						=  model[F"""{prefix_}mask_scores"""]
								snake_case__						:       Optional[int]       						=  TopKBinarizer.apply(__lowerCAmelCase   ,							__lowerCAmelCase     )
								snake_case__						:       Any       						=  tensor * mask
								print(F"""Pruned layer {name}"""     )
						elif pruning_method == "sigmoied_threshold":
								if "mask_scores" in name:
										continue
								snake_case__						:       Optional[Any]       						=  name[:-6]
								snake_case__						:       Dict       						=  model[F"""{prefix_}mask_scores"""]
								snake_case__						:       Dict       						=  ThresholdBinarizer.apply(__lowerCAmelCase   ,							__lowerCAmelCase   ,							__lowerCAmelCase     )
								snake_case__						:       int       						=  tensor * mask
								print(F"""Pruned layer {name}"""     )
						elif pruning_method == "l0":
								if "mask_scores" in name:
										continue
								snake_case__						:       Any       						=  name[:-6]
								snake_case__						:       Tuple       						=  model[F"""{prefix_}mask_scores"""]
								snake_case__						,       snake_case__						:       Tuple       						=  -0.1, 1.1
								snake_case__						:       List[str]       						=  torch.sigmoid(__lowerCAmelCase     )
								snake_case__						:       Optional[int]       						=  s * (r - l) + l
								snake_case__						:       Optional[int]       						=  s_bar.clamp(min=0.0   ,							max=1.0     )
								snake_case__						:       Union[str, Any]       						=  tensor * mask
								print(F"""Pruned layer {name}"""     )
						else:
								raise ValueError('''Unknown pruning method'''     )

		if target_model_path is None:
				snake_case__						:       Optional[int]       						=  os.path.join(
				    os.path.dirname(__lowerCAmelCase     )   ,							F"""bertarized_{os.path.basename(__lowerCAmelCase     )}"""     )

		if not os.path.isdir(__lowerCAmelCase     ):
				shutil.copytree(__lowerCAmelCase   ,							__lowerCAmelCase     )
				print(F"""\nCreated folder {target_model_path}"""     )

		torch.save(__lowerCAmelCase   ,							os.path.join(__lowerCAmelCase   ,							'''pytorch_model.bin'''     )     )
		print('''\nPruned model saved! See you later!'''     )


if __name__ == "__main__":
					lowerCAmelCase__						:       str										=						argparse.ArgumentParser()

					parser.add_argument(
					    """--pruning_method""",
					    choices=["""l0""", """magnitude""", """topK""", """sigmoied_threshold"""],
					    type=str,
					    required=True,
					    help=(
					        """Pruning Method (l0 = L0 regularization, magnitude = Magnitude pruning, topK = Movement pruning,"""
					        """ sigmoied_threshold = Soft movement pruning)"""
					    ),
					)
					parser.add_argument(
					    """--threshold""",
					    type=float,
					    required=False,
					    help=(
					        """For `magnitude` and `topK`, it is the level of remaining weights (in %) in the fine-pruned model."""
					        """For `sigmoied_threshold`, it is the threshold \tau against which the (sigmoied) scores are compared."""
					        """Not needed for `l0`"""
					    ),
					)
					parser.add_argument(
					    """--model_name_or_path""",
					    type=str,
					    required=True,
					    help="""Folder containing the model that was previously fine-pruned""",
					)
					parser.add_argument(
					    """--target_model_path""",
					    default=None,
					    type=str,
					    required=False,
					    help="""Folder containing the model that was previously fine-pruned""",
					)

					lowerCAmelCase__						:       Optional[Any]										=						parser.parse_args()

					main(args)



