
"""simple docstring"""





import unittest

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

from diffusers import DDIMScheduler, LDMPipeline, UNetaDModel, VQModel
from diffusers.utils.testing_utils import enable_full_determinism, require_torch, slow, torch_device


enable_full_determinism()






class        _UpperCAmelCase(       unittest.TestCase      ):



		@property
		def 		UpperCAmelCase				(   self)			->			List[str]:






							'''simple docstring'''
							torch.manual_seed(0)
							_UpperCamelCase									=    UNetaDModel(
							    block_out_channels=(32, 64)							,						layers_per_block=2							,						sample_size=32							,						in_channels=3							,						out_channels=3							,						down_block_types=('''DownBlock2D''', '''AttnDownBlock2D''')							,						up_block_types=('''AttnUpBlock2D''', '''UpBlock2D''')							,						)
							return model



		@property
		def 		UpperCAmelCase				(   self)			->			Union[str, Any]:






							'''simple docstring'''
							torch.manual_seed(0)
							_UpperCamelCase									=    VQModel(
							    block_out_channels=[32, 64]							,						in_channels=3							,						out_channels=3							,						down_block_types=['''DownEncoderBlock2D''', '''DownEncoderBlock2D''']							,						up_block_types=['''UpDecoderBlock2D''', '''UpDecoderBlock2D''']							,						latent_channels=3							,						)
							return model



		@property
		def 		UpperCAmelCase				(   self)			->			int:






							'''simple docstring'''
							torch.manual_seed(0)
							_UpperCamelCase									=    CLIPTextConfig(
							    bos_token_id=0							,						eos_token_id=2							,						hidden_size=32							,						intermediate_size=37							,						layer_norm_eps=1e-05							,						num_attention_heads=4							,						num_hidden_layers=5							,						pad_token_id=1							,						vocab_size=10_00							,						)
							return CLIPTextModel(__a)




		def 		UpperCAmelCase				(   self)			->			str:






							'''simple docstring'''
							_UpperCamelCase									=    self.dummy_uncond_unet
							_UpperCamelCase									=    DDIMScheduler()
							_UpperCamelCase									=    self.dummy_vq_model

							_UpperCamelCase									=    LDMPipeline(unet=__a							,						vqvae=__a							,						scheduler=__a)
							ldm.to(__a)
							ldm.set_progress_bar_config(disable=__a)

							_UpperCamelCase									=    torch.manual_seed(0)
							_UpperCamelCase									=    ldm(generator=__a							,						num_inference_steps=2							,						output_type='''numpy''').images

							_UpperCamelCase									=    torch.manual_seed(0)
							_UpperCamelCase									=    ldm(generator=__a							,						num_inference_steps=2							,						output_type='''numpy'''							,						return_dict=__a)[0]

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

							assert image.shape == (1, 64, 64, 3)
							_UpperCamelCase									=    np.array([0.8512, 0.818, 0.6411, 0.6808, 0.4465, 0.5618, 0.46, 0.6231, 0.5172])
							_UpperCamelCase									=    1e-2 if torch_device != '''mps''' else 3e-2

							assert np.abs(image_slice.flatten() - expected_slice).max() < tolerance
							assert np.abs(image_from_tuple_slice.flatten() - expected_slice).max() < tolerance






@slow
@require_torch
class        _UpperCAmelCase(       unittest.TestCase      ):



		def 		UpperCAmelCase				(   self)			->			Optional[Any]:






							'''simple docstring'''
							_UpperCamelCase									=    LDMPipeline.from_pretrained('''CompVis/ldm-celebahq-256''')
							ldm.to(__a)
							ldm.set_progress_bar_config(disable=__a)

							_UpperCamelCase									=    torch.manual_seed(0)
							_UpperCamelCase									=    ldm(generator=__a							,						num_inference_steps=5							,						output_type='''numpy''').images

							_UpperCamelCase									=    image[0, -3:, -3:, -1]

							assert image.shape == (1, 2_56, 2_56, 3)
							_UpperCamelCase									=    np.array([0.4399, 0.4_4975, 0.4_6825, 0.474, 0.4359, 0.4581, 0.4_5095, 0.4341, 0.4447])
							_UpperCamelCase									=    1e-2 if torch_device != '''mps''' else 3e-2

							assert np.abs(image_slice.flatten() - expected_slice).max() < tolerance








"""simple docstring"""





import argparse

import torch

from transformers import BertForMaskedLM


if __name__ == "__main__":
				_a								=						argparse.ArgumentParser(
				    description=(
				        """Extraction some layers of the full BertForMaskedLM or RObertaForMaskedLM for Transfer Learned"""
				        """ Distillation"""
				    )
				)
				parser.add_argument("""--model_type""", default="""bert""", choices=["""bert"""])
				parser.add_argument("""--model_name""", default="""bert-base-uncased""", type=str)
				parser.add_argument("""--dump_checkpoint""", default="""serialization_dir/tf_bert-base-uncased_0247911.pth""", type=str)
				parser.add_argument("""--vocab_transform""", action="""store_true""")
				_a								=						parser.parse_args()

				if args.model_type == "bert":
								_a								=						BertForMaskedLM.from_pretrained(args.model_name)
								_a								=						"""bert"""
				else:
								raise ValueError("""args.model_type should be \"bert\".""")

				_a								=						model.state_dict()
				_a								=						{}

				for w in ["word_embeddings", "position_embeddings"]:
								_a								=						state_dict[F"""{prefix}.embeddings.{w}.weight"""]
				for w in ["weight", "bias"]:
								_a								=						state_dict[F"""{prefix}.embeddings.LayerNorm.{w}"""]

				_a								=						0
				for teacher_idx in [0, 2, 4, 7, 9, 11]:
								for w in ["weight", "bias"]:
												_a								=						state_dict[
												    F"""{prefix}.encoder.layer.{teacher_idx}.attention.self.query.{w}"""
												]
												_a								=						state_dict[
												    F"""{prefix}.encoder.layer.{teacher_idx}.attention.self.key.{w}"""
												]
												_a								=						state_dict[
												    F"""{prefix}.encoder.layer.{teacher_idx}.attention.self.value.{w}"""
												]

												_a								=						state_dict[
												    F"""{prefix}.encoder.layer.{teacher_idx}.attention.output.dense.{w}"""
												]
												_a								=						state_dict[
												    F"""{prefix}.encoder.layer.{teacher_idx}.attention.output.LayerNorm.{w}"""
												]

												_a								=						state_dict[
												    F"""{prefix}.encoder.layer.{teacher_idx}.intermediate.dense.{w}"""
												]
												_a								=						state_dict[
												    F"""{prefix}.encoder.layer.{teacher_idx}.output.dense.{w}"""
												]
												_a								=						state_dict[
												    F"""{prefix}.encoder.layer.{teacher_idx}.output.LayerNorm.{w}"""
												]
								std_idx += 1

				_a								=						state_dict["""cls.predictions.decoder.weight"""]
				_a								=						state_dict["""cls.predictions.bias"""]
				if args.vocab_transform:
								for w in ["weight", "bias"]:
												_a								=						state_dict[F"""cls.predictions.transform.dense.{w}"""]
												_a								=						state_dict[F"""cls.predictions.transform.LayerNorm.{w}"""]

				print(F"""N layers selected for distillation: {std_idx}""")
				print(F"""Number of params transferred for distillation: {len(compressed_sd.keys())}""")

				print(F"""Save transferred checkpoint to {args.dump_checkpoint}.""")
				torch.save(compressed_sd, args.dump_checkpoint)






