


'''simple docstring'''


import inspect
import unittest

from huggingface_hub import hf_hub_download

from transformers import ConvNextConfig, UperNetConfig
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
						import torch

						from transformers import UperNetForSemanticSegmentation
						from transformers.models.upernet.modeling_upernet import UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
						from PIL import Image

						from transformers import AutoImageProcessor






class 	lowerCAmelCase__			:




			def __init__(		self  :     List[Any]	,  lowerCamelCase__  :     Optional[Any]	,  lowerCamelCase__  :     List[str]=13	,  lowerCamelCase__  :     List[Any]=32	,  lowerCamelCase__  :     str=3	,  lowerCamelCase__  :     Union[str, Any]=4	,  lowerCamelCase__  :     str=[10, 20, 30, 40]	,  lowerCamelCase__  :     Optional[Any]=[2, 2, 3, 2]	,  lowerCamelCase__  :     Optional[Any]=True	,  lowerCamelCase__  :     List[str]=True	,  lowerCamelCase__  :     List[Any]=37	,  lowerCamelCase__  :     List[str]="gelu"	,  lowerCamelCase__  :     List[str]=10	,  lowerCamelCase__  :     Union[str, Any]=0.0_2	,  lowerCamelCase__  :     int=["stage2", "stage3", "stage4"]	,  lowerCamelCase__  :     Optional[Any]=3	,  lowerCamelCase__  :     Tuple=None	,  )   ->Dict:



					'''simple docstring'''







					_UpperCAmelCase		:       List[Any]								=      parent
					_UpperCAmelCase		:       Dict								=      batch_size
					_UpperCAmelCase		:       List[str]								=      image_size
					_UpperCAmelCase		:       Tuple								=      num_channels
					_UpperCAmelCase		:       Any								=      num_stages
					_UpperCAmelCase		:       List[str]								=      hidden_sizes
					_UpperCAmelCase		:       Optional[Any]								=      depths
					_UpperCAmelCase		:       List[str]								=      is_training
					_UpperCAmelCase		:       Optional[int]								=      use_labels
					_UpperCAmelCase		:       List[Any]								=      intermediate_size
					_UpperCAmelCase		:       Union[str, Any]								=      hidden_act
					_UpperCAmelCase		:       Optional[int]								=      type_sequence_label_size
					_UpperCAmelCase		:       Optional[int]								=      initializer_range
					_UpperCAmelCase		:       str								=      out_features
					_UpperCAmelCase		:       Optional[int]								=      num_labels
					_UpperCAmelCase		:       Any								=      scope
					_UpperCAmelCase		:       Optional[int]								=      num_stages




			def 							lowerCAmelCase__    (		self  :     Tuple )   ->Optional[Any]:



					'''simple docstring'''







					_UpperCAmelCase		:       Optional[Any]								=      floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

					_UpperCAmelCase		:       Any								=      None
					if self.use_labels:
							_UpperCAmelCase		:       List[Any]								=      ids_tensor([self.batch_size]	,  self.type_sequence_label_size )

					_UpperCAmelCase		:       Optional[Any]								=      self.get_config()

					return config, pixel_values, labels




			def 							lowerCAmelCase__    (		self  :     Dict )   ->Any:



					'''simple docstring'''







					return ConvNextConfig(
					    num_channels=self.num_channels	,  num_stages=self.num_stages	,  hidden_sizes=self.hidden_sizes	,  depths=self.depths	,  is_training=self.is_training	,  intermediate_size=self.intermediate_size	,  hidden_act=self.hidden_act	,  out_features=self.out_features	,  )




			def 							lowerCAmelCase__    (		self  :     Dict )   ->Dict:



					'''simple docstring'''







					return UperNetConfig(
					    backbone_config=self.get_backbone_config()	,  hidden_size=5_12	,  pool_scales=[1, 2, 3, 6]	,  use_auxiliary_head=lowerCamelCase__	,  auxiliary_loss_weight=0.4	,  auxiliary_in_channels=40	,  auxiliary_channels=2_56	,  auxiliary_num_convs=1	,  auxiliary_concat_input=lowerCamelCase__	,  loss_ignore_index=2_55	,  num_labels=self.num_labels	,  )




			def 							lowerCAmelCase__    (		self  :     Union[str, Any]	,  lowerCamelCase__  :     List[Any]	,  lowerCamelCase__  :     str	,  lowerCamelCase__  :     Optional[Any] )   ->Tuple:



					'''simple docstring'''







					_UpperCAmelCase		:       List[Any]								=      UperNetForSemanticSegmentation(config=lowerCamelCase__ )
					model.to(lowerCamelCase__ )
					model.eval()
					_UpperCAmelCase		:       str								=      model(lowerCamelCase__ )
					self.parent.assertEqual(
					    result.logits.shape	,  (self.batch_size, self.num_labels, self.image_size, self.image_size) )






			def 							lowerCAmelCase__    (		self  :     Optional[Any] )   ->Tuple:



					'''simple docstring'''







					_UpperCAmelCase		:       List[Any]								=      self.prepare_config_and_inputs()
					(
					    (
					    _UpperCAmelCase
					)					,    (
					    _UpperCAmelCase
					)					,    (
					    _UpperCAmelCase
					)					,    
					)		:       int								=      config_and_inputs
					_UpperCAmelCase		:       Tuple								=      {"pixel_values": pixel_values}
					return config, inputs_dict






@require_torch
class 	lowerCAmelCase__			(			UpperCAmelCase__    ,       UpperCAmelCase__    ,       unittest.TestCase	):
			lowerCAmelCase :	str          =     (UperNetForSemanticSegmentation,) if is_torch_available() else ()
			lowerCAmelCase :	Dict          =     {"image-segmentation": UperNetForSemanticSegmentation} if is_torch_available() else {}
			lowerCAmelCase :	Optional[Any]          =     False
			lowerCAmelCase :	List[Any]          =     False
			lowerCAmelCase :	Optional[int]          =     False
			lowerCAmelCase :	Optional[int]          =     False
			lowerCAmelCase :	Any          =     False
			lowerCAmelCase :	Optional[int]          =     False




			def 							lowerCAmelCase__    (		self  :     Optional[int] )   ->List[str]:



					'''simple docstring'''







					_UpperCAmelCase		:       List[Any]								=      UperNetModelTester(self )
					_UpperCAmelCase		:       Tuple								=      ConfigTester(self	,  config_class=lowerCamelCase__	,  has_text_modality=lowerCamelCase__	,  hidden_size=37 )




			def 							lowerCAmelCase__    (		self  :     Optional[int] )   ->str:



					'''simple docstring'''







					self.create_and_test_config_common_properties()
					self.config_tester.create_and_test_config_to_json_string()
					self.config_tester.create_and_test_config_to_json_file()
					self.config_tester.create_and_test_config_from_and_save_pretrained()
					self.config_tester.create_and_test_config_with_num_labels()
					self.config_tester.check_config_can_be_init_without_params()
					self.config_tester.check_config_arguments_init()




			def 							lowerCAmelCase__    (		self  :     Tuple )   ->Tuple:



					'''simple docstring'''







					return




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



					'''simple docstring'''







					_UpperCAmelCase					,    _UpperCAmelCase		:       Any								=      self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
							_UpperCAmelCase		:       Tuple								=      model_class(lowerCamelCase__ )
							_UpperCAmelCase		:       str								=      inspect.signature(model.forward )
							# signature.parameters is an OrderedDict => so arg_names order is deterministic
							_UpperCAmelCase		:       Optional[int]								=      [*signature.parameters.keys()]

							_UpperCAmelCase		:       List[Any]								=      ["pixel_values"]
							self.assertListEqual(arg_names[:1]	,  lowerCamelCase__ )




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



					'''simple docstring'''







					_UpperCAmelCase		:       List[str]								=      self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_semantic_segmentation(*lowerCamelCase__ )




			@unittest.skip(reason="UperNet does not use inputs_embeds" )
			def 							lowerCAmelCase__    (		self  :     Optional[Any] )   ->int:



					'''simple docstring'''







					pass




			@unittest.skip(reason="UperNet does not support input and output embeddings" )
			def 							lowerCAmelCase__    (		self  :     List[Any] )   ->Tuple:



					'''simple docstring'''







					pass




			@unittest.skip(reason="UperNet does not have a base model" )
			def 							lowerCAmelCase__    (		self  :     List[str] )   ->Tuple:



					'''simple docstring'''







					pass




			@unittest.skip(reason="UperNet does not have a base model" )
			def 							lowerCAmelCase__    (		self  :     Optional[int] )   ->Tuple:



					'''simple docstring'''







					pass




			@require_torch_multi_gpu
			@unittest.skip(reason="UperNet has some layers using `add_module` which doesn't work well with `nn.DataParallel`" )
			def 							lowerCAmelCase__    (		self  :     Optional[Any] )   ->List[str]:



					'''simple docstring'''







					pass




			@unittest.skip("Will be fixed soon by reducing the size of the model used for common tests." )
			def 							lowerCAmelCase__    (		self  :     str )   ->Optional[Any]:



					'''simple docstring'''







					pass




			def 							lowerCAmelCase__    (		self  :     Optional[int] )   ->Any:



					'''simple docstring'''







					def check_hidden_states_output(lowerCamelCase__  :     List[str]	,  lowerCamelCase__  :     Tuple	,  lowerCamelCase__  :     Optional[Any] ):
							_UpperCAmelCase		:       Dict								=      model_class(lowerCamelCase__ )
							model.to(lowerCamelCase__ )
							model.eval()

							with torch.no_grad():
									_UpperCAmelCase		:       List[str]								=      model(**self._prepare_for_class(lowerCamelCase__	,  lowerCamelCase__ ) )

							_UpperCAmelCase		:       int								=      outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

							_UpperCAmelCase		:       List[str]								=      self.model_tester.num_stages
							self.assertEqual(len(lowerCamelCase__ )	,  expected_num_stages + 1 )

							# ConvNext's feature maps are of shape (batch_size, num_channels, height, width)
							self.assertListEqual(
							    list(hidden_states[0].shape[-2:] )	,  [self.model_tester.image_size // 4, self.model_tester.image_size // 4]	,  )

					_UpperCAmelCase					,    _UpperCAmelCase		:       str								=      self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
							_UpperCAmelCase		:       int								=      True
							check_hidden_states_output(lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__ )

							# check that output_hidden_states also work using config
							del inputs_dict["output_hidden_states"]
							_UpperCAmelCase		:       Optional[Any]								=      True

							check_hidden_states_output(lowerCamelCase__	,  lowerCamelCase__	,  lowerCamelCase__ )




			def 							lowerCAmelCase__    (		self  :     Union[str, Any] )   ->Optional[int]:



					'''simple docstring'''







					_UpperCAmelCase					,    _UpperCAmelCase		:       Union[str, Any]								=      self.model_tester.prepare_config_and_inputs_for_common()

					_UpperCAmelCase		:       Union[str, Any]								=      _config_zero_init(lowerCamelCase__ )
					_UpperCAmelCase		:       str								=      _config_zero_init(configs_no_init.backbone_config )
					for model_class in self.all_model_classes:
							_UpperCAmelCase		:       Tuple								=      model_class(config=lowerCamelCase__ )
							for name, param in model.named_parameters():
									if param.requires_grad:
											self.assertIn(
											    ((param.data.mean() * 1E9).round() / 1E9).item()	,  [0.0, 1.0]	,  msg=F"""Parameter {name} of model {model_class} seems not properly initialized"""	,  )




			@unittest.skip(reason="UperNet does not have tied weights" )
			def 							lowerCAmelCase__    (		self  :     Dict )   ->List[Any]:



					'''simple docstring'''







					pass






			@slow
			def 							lowerCAmelCase__    (		self  :     int )   ->int:



					'''simple docstring'''







					for model_name in UPERNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
							_UpperCAmelCase		:       str								=      UperNetForSemanticSegmentation.from_pretrained(lowerCamelCase__ )
							self.assertIsNotNone(lowerCamelCase__ )






def  __lowerCAmelCase ():
		_UpperCAmelCase		:       Any								=      hf_hub_download(
		    repo_id="hf-internal-testing/fixtures_ade20k"      , repo_type="dataset"      , filename="ADE_val_00000001.jpg"   )
		_UpperCAmelCase		:       int								=      Image.open(__lowerCAmelCase   ).convert("RGB"   )
		return image






@require_torch
@require_vision
@slow
class 	lowerCAmelCase__			(			unittest.TestCase	):




			def 							lowerCAmelCase__    (		self  :     Tuple )   ->Any:



					'''simple docstring'''







					_UpperCAmelCase		:       List[Any]								=      AutoImageProcessor.from_pretrained("openmmlab/upernet-swin-tiny" )
					_UpperCAmelCase		:       Union[str, Any]								=      UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-swin-tiny" ).to(lowerCamelCase__ )

					_UpperCAmelCase		:       str								=      prepare_img()
					_UpperCAmelCase		:       Any								=      processor(images=lowerCamelCase__	,  return_tensors="pt" ).to(lowerCamelCase__ )

					with torch.no_grad():
							_UpperCAmelCase		:       Union[str, Any]								=      model(**lowerCamelCase__ )

					_UpperCAmelCase		:       List[Any]								=      torch.Size((1, model.config.num_labels, 5_12, 5_12) )
					self.assertEqual(outputs.logits.shape	,  lowerCamelCase__ )

					_UpperCAmelCase		:       Any								=      torch.tensor(
					    [[-7.5_9_5_8, -7.5_9_5_8, -7.4_3_0_2], [-7.5_9_5_8, -7.5_9_5_8, -7.4_3_0_2], [-7.4_7_9_7, -7.4_7_9_7, -7.3_0_6_8]] ).to(lowerCamelCase__ )
					self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]	,  lowerCamelCase__	,  atol=1E-4 ) )






			def 							lowerCAmelCase__    (		self  :     List[str] )   ->Any:



					'''simple docstring'''







					_UpperCAmelCase		:       Dict								=      AutoImageProcessor.from_pretrained("openmmlab/upernet-convnext-tiny" )
					_UpperCAmelCase		:       Dict								=      UperNetForSemanticSegmentation.from_pretrained("openmmlab/upernet-convnext-tiny" ).to(lowerCamelCase__ )

					_UpperCAmelCase		:       int								=      prepare_img()
					_UpperCAmelCase		:       List[Any]								=      processor(images=lowerCamelCase__	,  return_tensors="pt" ).to(lowerCamelCase__ )

					with torch.no_grad():
							_UpperCAmelCase		:       Dict								=      model(**lowerCamelCase__ )

					_UpperCAmelCase		:       int								=      torch.Size((1, model.config.num_labels, 5_12, 5_12) )
					self.assertEqual(outputs.logits.shape	,  lowerCamelCase__ )

					_UpperCAmelCase		:       Dict								=      torch.tensor(
					    [[-8.8_1_1_0, -8.8_1_1_0, -8.6_5_2_1], [-8.8_1_1_0, -8.8_1_1_0, -8.6_5_2_1], [-8.7_7_4_6, -8.7_7_4_6, -8.6_1_3_0]] ).to(lowerCamelCase__ )
					self.assertTrue(torch.allclose(outputs.logits[0, 0, :3, :3]	,  lowerCamelCase__	,  atol=1E-4 ) )







'''simple docstring'''


import argparse

import torch

from transformers import BertForMaskedLM


if __name__ == "__main__":
						lowerCamelCase__    	=     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')
						lowerCamelCase__    	=     parser.parse_args()

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

						lowerCamelCase__    	=     model.state_dict()
						lowerCamelCase__    	=     {}

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

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

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

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

						lowerCamelCase__    	=     state_dict['cls.predictions.decoder.weight']
						lowerCamelCase__    	=     state_dict['cls.predictions.bias']
						if args.vocab_transform:
												for w in ["weight", "bias"]:
																		lowerCamelCase__    	=     state_dict[F'''cls.predictions.transform.dense.{w}''']
																		lowerCamelCase__    	=     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)



