






"""simple docstring"""



import argparse
import json
import os
from collections import OrderedDict

import numpy as np
import tensorflow as tf
import torch



def 						lowercase    (				__snake_case				: List[Any]	):
		lowercase_ : int							       =   os.path.join(args.tf_model_dir		,					'''parameters.json'''	)
		lowercase_ : Any							       =   json.loads(open(__snake_case	).read()	)
		if not params:
				raise ValueError(
				    F'''It seems that the json file at {parameter_file} is empty. Make sure you have a correct json file.'''	)
		if not args.output.endswith('''.pt'''	):
				lowercase_ : Dict							       =   args.output + '''.pt'''
		lowercase_ : Any							       =   OrderedDict()
		with tf.device('''/CPU:0'''	):
				lowercase_ : int							       =   tf.train.load_checkpoint(args.tf_model_dir	)
				lowercase_ : Optional[int]							       =   reader.get_variable_to_shape_map()
				for key_name in shapes.keys():
						lowercase_ : int							       =   reader.get_tensor(__snake_case	).astype(np.floataa	)
						if key_name.endswith('''/adam_m'''	) or key_name.endswith('''/adam_v'''	):
								continue
						if key_name.startswith('''pasts/'''	):
								if key_name.startswith('''pasts/mlp'''	):
										lowercase_ : Optional[Any]							       =   int(key_name[9]	)
								elif key_name.startswith('''pasts/out'''	):
										lowercase_ : Tuple							       =   8
								lowercase_ : str							       =   '''model.sqout.%d.weight''' % (player * 2)  # enter to nn.Sequencial with Tanh, so 2 at a time
								lowercase_ : Optional[Any]							       =   vnp.transpose([1, 0]	).copy()  # Mesh-Tensorflow is a diagonal matrix
								lowercase_ : List[Any]							       =   torch.tensor(__snake_case	)
						elif key_name.startswith('''model/moe'''	):
								lowercase_ : Optional[int]							       =   int(key_name[9:].split('''/'''	)[0]	)
								if key_name.endswith('''/switch_gating/kernel'''	):
										lowercase_ : List[str]							       =   '''model.blocks.%d.feed_forward.mlp.router.classifier.weight''' % player
										lowercase_ : str							       =   vnp.transpose([1, 0]	).copy()  # Mesh-Tensorflow is a diagonal matrix
										lowercase_ : Tuple							       =   torch.tensor(__snake_case	)
								elif key_name.endswith('''/softmlp/kernel'''	):
										lowercase_ : List[str]							       =   '''model.blocks.%d.feed_forward.soft_bypass_mlp.weight''' % player
										lowercase_ : List[str]							       =   vnp.transpose([1, 0]	).copy()  # Mesh-Tensorflow is a diagonal matrix
										lowercase_ : List[str]							       =   torch.tensor(__snake_case	)
								elif key_name.endswith('''/wo/kernel'''	) or key_name.endswith('''/wi/kernel'''	):
										lowercase_ : Union[str, Any]							       =   key_name[-9:-7]
										for i in range(1_6	):
												lowercase_ : Tuple							       =   '''model.blocks.%d.feed_forward.mlp.experts.expert_%d.%s.weight''' % (player, i, nlayer)
												lowercase_ : str							       =   (
												    vnp[i].transpose([1, 0]	).copy()
												)  # In Mesh-Tensorflow, it is one array, so it is divided
												lowercase_ : Tuple							       =   torch.tensor(__snake_case	)
						elif key_name.startswith('''model/mlp'''	):
								lowercase_ : Optional[int]							       =   int(key_name[9:].split('''/'''	)[0]	)
								if key_name.endswith('''/p1/kernel'''	):
										lowercase_ : List[Any]							       =   '''model.blocks.%d.feed_forward.mlp.wi.weight''' % player
										lowercase_ : str							       =   vnp.transpose([1, 0]	).copy()  # Mesh-Tensorflow is a diagonal matrix
										lowercase_ : Dict							       =   torch.tensor(__snake_case	)
								elif key_name.endswith('''/p1/bias'''	):
										lowercase_ : Any							       =   '''model.blocks.%d.feed_forward.mlp.wi.bias''' % player
										lowercase_ : Any							       =   vnp.copy()  # same because it is one dimensional
										lowercase_ : int							       =   torch.tensor(__snake_case	)
								elif key_name.endswith('''/p2/kernel'''	):
										lowercase_ : Optional[Any]							       =   '''model.blocks.%d.feed_forward.mlp.wo.weight''' % player
										lowercase_ : int							       =   vnp.transpose([1, 0]	).copy()  # Mesh-Tensorflow is a diagonal matrix
										lowercase_ : int							       =   torch.tensor(__snake_case	)
								elif key_name.endswith('''/p2/bias'''	):
										lowercase_ : Optional[int]							       =   '''model.blocks.%d.feed_forward.mlp.wo.bias''' % player
										lowercase_ : List[str]							       =   vnp.copy()  # same because it is one dimensional
										lowercase_ : int							       =   torch.tensor(__snake_case	)
						elif key_name.startswith('''model/ln'''	):
								lowercase_ : Optional[Any]							       =   int(key_name[8:].split('''/'''	)[0]	)
								if key_name.endswith('''/b'''	):
										lowercase_ : str							       =   '''model.blocks.%d.feed_forward.norm.bias''' % player
										lowercase_ : Any							       =   vnp.copy()  # same because it is one dimensional
										lowercase_ : int							       =   torch.tensor(__snake_case	)
								elif key_name.endswith('''/g'''	):
										lowercase_ : Union[str, Any]							       =   '''model.blocks.%d.feed_forward.norm.weight''' % player
										lowercase_ : Union[str, Any]							       =   vnp.copy()  # same because it is one dimensional
										lowercase_ : int							       =   torch.tensor(__snake_case	)
						elif key_name.startswith('''model/att'''	):
								lowercase_ : Optional[int]							       =   int(key_name[9:].split('''/'''	)[0]	)
								if key_name.endswith('''/qkv/kernel'''	):
										lowercase_ : Dict							       =   vnp.copy()  # Compute same dimension as Mesh-tensorflow using einsum
										lowercase_ : Tuple							       =   state[:, 0, :, :]
										lowercase_ : Dict							       =   state[:, 1, :, :]
										lowercase_ : Union[str, Any]							       =   state[:, 2, :, :]
										lowercase_ : int							       =   (
										    state_q.reshape([state_q.shape[0], state_q.shape[1] * state_q.shape[2]]	)
										    .transpose([1, 0]	)
										    .copy()
										)  # Mesh-Tensorflow is a diagonal matrix
										lowercase_ : Optional[Any]							       =   (
										    state_k.reshape([state_k.shape[0], state_k.shape[1] * state_k.shape[2]]	)
										    .transpose([1, 0]	)
										    .copy()
										)  # Mesh-Tensorflow is a diagonal matrix
										lowercase_ : Union[str, Any]							       =   (
										    state_v.reshape([state_v.shape[0], state_v.shape[1] * state_v.shape[2]]	)
										    .transpose([1, 0]	)
										    .copy()
										)  # Mesh-Tensorflow is a diagonal matrix
										lowercase_ : List[str]							       =   '''model.blocks.%d.self_attn.self_attn.q_proj.weight''' % player
										lowercase_ : str							       =   torch.tensor(__snake_case	)
										lowercase_ : str							       =   '''model.blocks.%d.self_attn.self_attn.k_proj.weight''' % player
										lowercase_ : Any							       =   torch.tensor(__snake_case	)
										lowercase_ : List[Any]							       =   '''model.blocks.%d.self_attn.self_attn.v_proj.weight''' % player
										lowercase_ : Any							       =   torch.tensor(__snake_case	)
								elif key_name.endswith('''/o/kernel'''	):
										lowercase_ : Dict							       =   '''model.blocks.%d.self_attn.self_attn.out_proj.weight''' % player
										lowercase_ : Optional[int]							       =   (
										    vnp.reshape([vnp.shape[0] * vnp.shape[1], vnp.shape[2]]	).transpose([1, 0]	).copy()
										)  # Mesh-Tensorflow is a diagonal matrix
										lowercase_ : Any							       =   torch.tensor(__snake_case	)
						elif key_name.startswith('''model/an'''	):
								lowercase_ : str							       =   int(key_name[8:].split('''/'''	)[0]	)
								if key_name.endswith('''/b'''	):
										lowercase_ : Dict							       =   '''model.blocks.%d.self_attn.norm.bias''' % player
										lowercase_ : Union[str, Any]							       =   vnp.copy()  # same because it is one dimensional
										lowercase_ : str							       =   torch.tensor(__snake_case	)
								elif key_name.endswith('''/g'''	):
										lowercase_ : str							       =   '''model.blocks.%d.self_attn.norm.weight''' % player
										lowercase_ : Optional[Any]							       =   vnp.copy()  # same because it is one dimensional
										lowercase_ : Any							       =   torch.tensor(__snake_case	)
						elif (
						    key_name.startswith('''model/wte'''	)
						    or key_name.startswith('''model/wpe'''	)
						    or key_name.startswith('''model/ete'''	)
						):
								lowercase_ : int							       =   {'''wte''': '''embed_tokens''', '''wpe''': '''position_embeddings''', '''ete''': '''extra_position_embeddings'''}[
								    key_name[-3:]
								]
								lowercase_ : str							       =   '''model.%s.weight''' % nlayer
								lowercase_ : int							       =   vnp.copy()  # same in embedded
								lowercase_ : int							       =   torch.tensor(__snake_case	)
								if key_name.startswith('''model/wte'''	):
										lowercase_ : Dict							       =   '''lm_head.weight'''
										lowercase_ : Tuple							       =   vnp.copy()  # same in embedded
										lowercase_ : Dict							       =   torch.tensor(__snake_case	)
						elif key_name.startswith('''model/wob'''	):
								lowercase_ : int							       =   '''final_logits_bias'''
								lowercase_ : Any							       =   vnp.copy()  # same in embedded
								lowercase_ : Optional[Any]							       =   state.reshape((1, -1)	)
								lowercase_ : Any							       =   torch.tensor(__snake_case	)
						elif key_name == "model/dense/kernel":
								lowercase_ : List[str]							       =   '''model.last_project.weight'''
								lowercase_ : str							       =   vnp.transpose([1, 0]	).copy()  # Mesh-Tensorflow is a diagonal matrix
								lowercase_ : Optional[Any]							       =   torch.tensor(__snake_case	)
						elif key_name == "model/dense_1/bias":
								lowercase_ : Dict							       =   '''model.last_project.bias'''
								lowercase_ : Tuple							       =   vnp.copy()  # same because it is one dimensional
								lowercase_ : Any							       =   torch.tensor(__snake_case	)
		torch.save(__snake_case		,					args.output	)


if __name__ == "__main__":
						__A		:       Tuple  				=		argparse.ArgumentParser(
						    description='''model converter.''', formatter_class=argparse.ArgumentDefaultsHelpFormatter
						)
						parser.add_argument('''--tf_model_dir''', metavar='''PATH''', type=str, required=True, help='''import model''')
						parser.add_argument('''--output''', metavar='''PATH''', type=str, required=True, help='''output model''')
						__A		:       Any  				=		parser.parse_args()
						convert_tf_gptsan_to_pt(args)











"""simple docstring"""



from __future__ import annotations

import inspect
import unittest

import numpy as np

from transformers import ResNetConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

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


if is_tf_available():
						import tensorflow as tf

						from transformers import TFResNetForImageClassification, TFResNetModel
						from transformers.models.resnet.modeling_tf_resnet import TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
						from PIL import Image

						from transformers import AutoImageProcessor





class 					_UpperCAmelCase       :




					def __init__(				self     :						int    ,						A     :						Tuple    ,						A     :						int=3    ,						A     :						List[str]=32    ,						A     :						Dict=3    ,						A     :						Any=10    ,						A     :						Dict=[10, 20, 30, 40]    ,						A     :						Optional[Any]=[1, 1, 2, 1]    ,						A     :						Union[str, Any]=True    ,						A     :						Optional[Any]=True    ,						A     :						Any="relu"    ,						A     :						Optional[Any]=3    ,						A     :						Tuple=None    ,						)  ->  Dict:
							lowercase_ : str							       =   parent
							lowercase_ : List[Any]							       =   batch_size
							lowercase_ : Optional[int]							       =   image_size
							lowercase_ : int							       =   num_channels
							lowercase_ : int							       =   embeddings_size
							lowercase_ : str							       =   hidden_sizes
							lowercase_ : List[str]							       =   depths
							lowercase_ : Dict							       =   is_training
							lowercase_ : int							       =   use_labels
							lowercase_ : Any							       =   hidden_act
							lowercase_ : List[Any]							       =   num_labels
							lowercase_ : Tuple							       =   scope
							lowercase_ : Optional[Any]							       =   len(A   )




					def   A  (				self     :						str   )  ->  Tuple:
							lowercase_ : Optional[int]							       =   floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

							lowercase_ : Union[str, Any]							       =   None
							if self.use_labels:
									lowercase_ : List[str]							       =   ids_tensor([self.batch_size]    ,						self.num_labels   )

							lowercase_ : Optional[int]							       =   self.get_config()

							return config, pixel_values, labels




					def   A  (				self     :						Dict   )  ->  int:
							return ResNetConfig(
							    num_channels=self.num_channels    ,						embeddings_size=self.embeddings_size    ,						hidden_sizes=self.hidden_sizes    ,						depths=self.depths    ,						hidden_act=self.hidden_act    ,						num_labels=self.num_labels    ,						image_size=self.image_size    ,						)




					def   A  (				self     :						str    ,						A     :						Tuple    ,						A     :						str    ,						A     :						str   )  ->  str:
							lowercase_ : str							       =   TFResNetModel(config=A   )
							lowercase_ : Union[str, Any]							       =   model(A   )
							# expected last hidden states: B, C, H // 32, W // 32
							self.parent.assertEqual(
							    result.last_hidden_state.shape    ,						(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)    ,						)




					def   A  (				self     :						Any    ,						A     :						int    ,						A     :						List[Any]    ,						A     :						Optional[Any]   )  ->  Optional[Any]:
							lowercase_ : Tuple							       =   self.num_labels
							lowercase_ : Union[str, Any]							       =   TFResNetForImageClassification(A   )
							lowercase_ : Tuple							       =   model(A    ,						labels=A   )
							self.parent.assertEqual(result.logits.shape    ,						(self.batch_size, self.num_labels)   )






					def   A  (				self     :						Union[str, Any]   )  ->  Tuple:
							lowercase_ : Tuple							       =   self.prepare_config_and_inputs()
							lowercase_       ,					lowercase_       ,					lowercase_ : Dict							       =   config_and_inputs
							lowercase_ : int							       =   {'''pixel_values''': pixel_values}
							return config, inputs_dict





@require_tf
class 					_UpperCAmelCase       ( _A   ,						_A   ,						unittest.TestCase							):
					SCREAMING_SNAKE_CASE_		:				Union[str, Any]							    =   (TFResNetModel, TFResNetForImageClassification) if is_tf_available() else ()
					SCREAMING_SNAKE_CASE_		:				List[Any]							    =   (
					    {"feature-extraction": TFResNetModel, "image-classification": TFResNetForImageClassification}
					    if is_tf_available()
					    else {}
					)

					SCREAMING_SNAKE_CASE_		:				Optional[int]							    =   False
					SCREAMING_SNAKE_CASE_		:				Optional[int]							    =   False
					SCREAMING_SNAKE_CASE_		:				str							    =   False
					SCREAMING_SNAKE_CASE_		:				Optional[int]							    =   False
					SCREAMING_SNAKE_CASE_		:				Any							    =   False




					def   A  (				self     :						Union[str, Any]   )  ->  List[Any]:
							lowercase_ : int							       =   TFResNetModelTester(self   )
							lowercase_ : str							       =   ConfigTester(self    ,						config_class=A    ,						has_text_modality=A   )




					def   A  (				self     :						Dict   )  ->  Optional[Any]:
							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   A  (				self     :						Dict   )  ->  List[Any]:
							return




					@unittest.skip(reason='''ResNet does not use inputs_embeds'''   )
					def   A  (				self     :						Any   )  ->  Any:
							pass




					@unittest.skip(reason='''ResNet does not support input and output embeddings'''   )
					def   A  (				self     :						List[str]   )  ->  Optional[Any]:
							pass




					def   A  (				self     :						str   )  ->  Tuple:
							lowercase_       ,					lowercase_ : Optional[Any]							       =   self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
									lowercase_ : int							       =   model_class(A   )
									lowercase_ : str							       =   inspect.signature(model.call   )
									# signature.parameters is an OrderedDict => so arg_names order is deterministic
									lowercase_ : str							       =   [*signature.parameters.keys()]

									lowercase_ : str							       =   ['''pixel_values''']
									self.assertListEqual(arg_names[:1]    ,						A   )




					def   A  (				self     :						List[str]   )  ->  Tuple:
							lowercase_ : Optional[Any]							       =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*A   )




					def   A  (				self     :						List[Any]   )  ->  List[str]:
							def check_hidden_states_output(A     :						Union[str, Any]    ,						A     :						int    ,						A     :						List[Any]   ):
									lowercase_ : int							       =   model_class(A   )
									lowercase_ : Optional[Any]							       =   model(**self._prepare_for_class(A    ,						A   )   )

									lowercase_ : Union[str, Any]							       =   outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

									lowercase_ : Any							       =   self.model_tester.num_stages
									self.assertEqual(len(A   )    ,						expected_num_stages + 1   )

									# ResNet'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]    ,						)

							lowercase_       ,					lowercase_ : List[Any]							       =   self.model_tester.prepare_config_and_inputs_for_common()
							lowercase_ : Union[str, Any]							       =   ['''basic''', '''bottleneck''']
							for model_class in self.all_model_classes:
									for layer_type in layers_type:
											lowercase_ : List[str]							       =   layer_type
											lowercase_ : Tuple							       =   True
											check_hidden_states_output(A    ,						A    ,						A   )

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

											check_hidden_states_output(A    ,						A    ,						A   )




					def   A  (				self     :						Optional[int]   )  ->  Tuple:
							lowercase_ : Optional[int]							       =   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_image_classification(*A   )






					@slow
					def   A  (				self     :						List[str]   )  ->  Optional[int]:
							for model_name in TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									lowercase_ : Tuple							       =   TFResNetModel.from_pretrained(A   )
									self.assertIsNotNone(A   )



def 						lowercase    (				):
		lowercase_ : List[str]							       =   Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''	)
		return image






@require_tf
@require_vision
class 					_UpperCAmelCase       ( unittest.TestCase							):




					@cached_property
					def   A  (				self     :						Any   )  ->  Optional[int]:
							return (
							    AutoImageProcessor.from_pretrained(TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]   )
							    if is_vision_available()
							    else None
							)






					@slow
					def   A  (				self     :						Any   )  ->  Optional[int]:
							lowercase_ : Optional[int]							       =   TFResNetForImageClassification.from_pretrained(TF_RESNET_PRETRAINED_MODEL_ARCHIVE_LIST[0]   )

							lowercase_ : List[Any]							       =   self.default_image_processor
							lowercase_ : Dict							       =   prepare_img()
							lowercase_ : List[str]							       =   image_processor(images=A    ,						return_tensors='''tf'''   )

							# forward pass
							lowercase_ : Tuple							       =   model(**A   )

							# verify the logits
							lowercase_ : Optional[int]							       =   tf.TensorShape((1, 10_00)   )
							self.assertEqual(outputs.logits.shape    ,						A   )

							lowercase_ : Optional[Any]							       =   tf.constant([-11.1069, -9.7877, -8.3777]   )

							self.assertTrue(np.allclose(outputs.logits[0, :3].numpy()    ,						A    ,						atol=1e-4   )   )



