'''simple docstring'''

from __future__ import annotations

import unittest

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

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


if is_tf_available():
							import numpy
							import tensorflow as tf

							from transformers.models.esm.modeling_tf_esm import (
							    TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
							    TFEsmForMaskedLM,
							    TFEsmForSequenceClassification,
							    TFEsmForTokenClassification,
							    TFEsmModel,
							)
class   SCREAMING_SNAKE_CASE			:







		"""simple docstring"""






		def __init__(   self		,						__A		,						) ->     Optional[Any]:
								_lowerCAmelCase										=parent
								_lowerCAmelCase										=13
								_lowerCAmelCase										=7
								_lowerCAmelCase										=True
								_lowerCAmelCase										=True
								_lowerCAmelCase										=True
								_lowerCAmelCase										=99
								_lowerCAmelCase										=32
								_lowerCAmelCase										=2
								_lowerCAmelCase										=4
								_lowerCAmelCase										=37
								_lowerCAmelCase										='gelu'
								_lowerCAmelCase										=0.1
								_lowerCAmelCase										=0.1
								_lowerCAmelCase										=512
								_lowerCAmelCase										=16
								_lowerCAmelCase										=2
								_lowerCAmelCase										=0.02
								_lowerCAmelCase										=3
								_lowerCAmelCase										=4
								_lowerCAmelCase										=None






		def 		UpperCamelCase__       (   self				) ->     List[str]:
								_lowerCAmelCase										=ids_tensor([self.batch_size, self.seq_length]		,						self.vocab_size				)

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

								_lowerCAmelCase										=None
								_lowerCAmelCase										=None
								_lowerCAmelCase										=None
								if self.use_labels:
														_lowerCAmelCase										=ids_tensor([self.batch_size]		,						self.type_sequence_label_size				)
														_lowerCAmelCase										=ids_tensor([self.batch_size, self.seq_length]		,						self.num_labels				)
														_lowerCAmelCase										=ids_tensor([self.batch_size]		,						self.num_choices				)

								_lowerCAmelCase										=EsmConfig(
								    vocab_size=self.vocab_size		,						hidden_size=self.hidden_size		,						num_hidden_layers=self.num_hidden_layers		,						pad_token_id=1		,						num_attention_heads=self.num_attention_heads		,						intermediate_size=self.intermediate_size		,						hidden_act=self.hidden_act		,						hidden_dropout_prob=self.hidden_dropout_prob		,						attention_probs_dropout_prob=self.attention_probs_dropout_prob		,						max_position_embeddings=self.max_position_embeddings		,						type_vocab_size=self.type_vocab_size		,						initializer_range=self.initializer_range		,						)

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






		def 		UpperCamelCase__       (   self				) ->     str:
								(
								    (
								    _lowerCAmelCase 
								)						,				(
								    _lowerCAmelCase 
								)						,				(
								    _lowerCAmelCase 
								)						,				(
								    _lowerCAmelCase 
								)						,				(
								    _lowerCAmelCase 
								)						,				(
								    _lowerCAmelCase 
								)						,				 
								)										=self.prepare_config_and_inputs()

								_lowerCAmelCase										=True
								_lowerCAmelCase										=floats_tensor([self.batch_size, self.seq_length, self.hidden_size]				)
								_lowerCAmelCase										=ids_tensor([self.batch_size, self.seq_length]		,						vocab_size=2				)

								return (
								    config,
								    input_ids,
								    input_mask,
								    sequence_labels,
								    token_labels,
								    choice_labels,
								    encoder_hidden_states,
								    encoder_attention_mask,
								)






		def 		UpperCamelCase__       (   self		,						__A		,						__A		,						__A		,						__A		,						__A		,						__A				) ->     List[str]:
								_lowerCAmelCase										=TFEsmModel(config=UpperCamelCase__				)
								_lowerCAmelCase										={'input_ids': input_ids, 'attention_mask': input_mask}
								_lowerCAmelCase										=model(UpperCamelCase__				)

								_lowerCAmelCase										=[input_ids, input_mask]
								_lowerCAmelCase										=model(UpperCamelCase__				)

								_lowerCAmelCase										=model(UpperCamelCase__				)

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






		def 		UpperCamelCase__       (   self		,						__A		,						__A		,						__A		,						__A		,						__A		,						__A		,						__A		,						__A		,						) ->     Any:
								_lowerCAmelCase										=True

								_lowerCAmelCase										=TFEsmModel(config=UpperCamelCase__				)
								_lowerCAmelCase										={
								    'input_ids': input_ids,
								    'attention_mask': input_mask,
								    'encoder_hidden_states': encoder_hidden_states,
								    'encoder_attention_mask': encoder_attention_mask,
								}
								_lowerCAmelCase										=model(UpperCamelCase__				)

								_lowerCAmelCase										=[input_ids, input_mask]
								_lowerCAmelCase										=model(UpperCamelCase__		,						encoder_hidden_states=UpperCamelCase__				)

								# Also check the case where encoder outputs are not passed
								_lowerCAmelCase										=model(UpperCamelCase__		,						attention_mask=UpperCamelCase__				)

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






		def 		UpperCamelCase__       (   self		,						__A		,						__A		,						__A		,						__A		,						__A		,						__A				) ->     Union[str, Any]:
								_lowerCAmelCase										=TFEsmForMaskedLM(config=UpperCamelCase__				)
								_lowerCAmelCase										=model([input_ids, input_mask]				)
								self.parent.assertEqual(result.logits.shape		,						(self.batch_size, self.seq_length, self.vocab_size)				)






		def 		UpperCamelCase__       (   self		,						__A		,						__A		,						__A		,						__A		,						__A		,						__A				) ->     Union[str, Any]:
								_lowerCAmelCase										=self.num_labels
								_lowerCAmelCase										=TFEsmForTokenClassification(config=UpperCamelCase__				)
								_lowerCAmelCase										={'input_ids': input_ids, 'attention_mask': input_mask}
								_lowerCAmelCase										=model(UpperCamelCase__				)
								self.parent.assertEqual(result.logits.shape		,						(self.batch_size, self.seq_length, self.num_labels)				)






		def 		UpperCamelCase__       (   self				) ->     Optional[int]:
								_lowerCAmelCase										=self.prepare_config_and_inputs()
								(
								    (
								    _lowerCAmelCase 
								)						,				(
								    _lowerCAmelCase 
								)						,				(
								    _lowerCAmelCase 
								)						,				(
								    _lowerCAmelCase 
								)						,				(
								    _lowerCAmelCase 
								)						,				(
								    _lowerCAmelCase 
								)						,				 
								)										=config_and_inputs
								_lowerCAmelCase										={'input_ids': input_ids, 'attention_mask': input_mask}
								return config, inputs_dict

@require_tf
class   SCREAMING_SNAKE_CASE			(						lowercase_    ,			lowercase_    ,			unittest.TestCase):







		"""simple docstring"""
		lowercase   :      Tuple	       =							(
		    (
		        TFEsmModel,
		        TFEsmForMaskedLM,
		        TFEsmForSequenceClassification,
		        TFEsmForTokenClassification,
		    )
		    if is_tf_available()
		    else ()
		)
		lowercase   :      List[str]	       =							(
		    {
		        '''feature-extraction''': TFEsmModel,
		        '''fill-mask''': TFEsmForMaskedLM,
		        '''text-classification''': TFEsmForSequenceClassification,
		        '''token-classification''': TFEsmForTokenClassification,
		        '''zero-shot''': TFEsmForSequenceClassification,
		    }
		    if is_tf_available()
		    else {}
		)
		lowercase   :      Tuple	       =							False
		lowercase   :      Optional[int]	       =							False






		def 		UpperCamelCase__       (   self				) ->     str:
								_lowerCAmelCase										=TFEsmModelTester(self				)
								_lowerCAmelCase										=ConfigTester(self		,						config_class=UpperCamelCase__		,						hidden_size=37				)






		def 		UpperCamelCase__       (   self				) ->     Any:
								self.config_tester.run_common_tests()






		def 		UpperCamelCase__       (   self				) ->     Union[str, Any]:
								_lowerCAmelCase										=self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_model(*UpperCamelCase__				)






		def 		UpperCamelCase__       (   self				) ->     int:
								_lowerCAmelCase										=self.model_tester.prepare_config_and_inputs_for_decoder()
								self.model_tester.create_and_check_model_as_decoder(*UpperCamelCase__				)






		def 		UpperCamelCase__       (   self				) ->     List[Any]:
								_lowerCAmelCase										=self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_masked_lm(*UpperCamelCase__				)






		def 		UpperCamelCase__       (   self				) ->     Tuple:
								_lowerCAmelCase										=self.model_tester.prepare_config_and_inputs()
								self.model_tester.create_and_check_for_token_classification(*UpperCamelCase__				)






		@slow
		def 		UpperCamelCase__       (   self				) ->     int:







								for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
														_lowerCAmelCase										=TFEsmModel.from_pretrained(UpperCamelCase__				)
														self.assertIsNotNone(UpperCamelCase__				)






		@unittest.skip('Protein models do not support embedding resizing.'				)
		def 		UpperCamelCase__       (   self				) ->     Optional[int]:
								pass






		@unittest.skip('Protein models do not support embedding resizing.'				)
		def 		UpperCamelCase__       (   self				) ->     Tuple:
								pass






		def 		UpperCamelCase__       (   self				) ->     str:
								_lowerCAmelCase						,				_lowerCAmelCase										=self.model_tester.prepare_config_and_inputs_for_common()

								for model_class in self.all_model_classes:
														_lowerCAmelCase										=model_class(UpperCamelCase__				)
														assert isinstance(model.get_input_embeddings()		,						tf.keras.layers.Layer				)
														if model_class is TFEsmForMaskedLM:
																				# Output embedding test differs from the main test because they're a matrix, not a layer
																				_lowerCAmelCase										=model.get_bias()
																				assert isinstance(UpperCamelCase__		,						UpperCamelCase__				)
																				for k, v in name.items():
																										assert isinstance(UpperCamelCase__		,						tf.Variable				)
														else:
																				_lowerCAmelCase										=model.get_output_embeddings()
																				assert x is None
																				_lowerCAmelCase										=model.get_bias()
																				assert name is None

@require_tf
class   SCREAMING_SNAKE_CASE			(						unittest.TestCase):







		"""simple docstring"""






		@slow
		def 		UpperCamelCase__       (   self				) ->     Optional[Any]:
								_lowerCAmelCase										=TFEsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D'				)

								_lowerCAmelCase										=tf.constant([[0, 1, 2, 3, 4, 5]]				)
								_lowerCAmelCase										=model(UpperCamelCase__				)[0]
								_lowerCAmelCase										=[1, 6, 33]
								self.assertEqual(list(output.numpy().shape				)		,						UpperCamelCase__				)
								# compare the actual values for a slice.
								_lowerCAmelCase										=tf.constant(
								    [
								        [
								            [8.921_518, -10.589_814, -6.4_671_307],
								            [-6.3_967_156, -13.911_377, -1.1_211_915],
								            [-7.781_247, -13.951_557, -3.740_592],
								        ]
								    ]				)
								self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()		,						expected_slice.numpy()		,						atol=1E-2				)				)






		@slow
		def 		UpperCamelCase__       (   self				) ->     str:
								_lowerCAmelCase										=TFEsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D'				)

								_lowerCAmelCase										=tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]				)
								_lowerCAmelCase										=model(UpperCamelCase__				)[0]
								# compare the actual values for a slice.
								_lowerCAmelCase										=tf.constant(
								    [
								        [
								            [0.14_443_092, 0.54_125_327, 0.3_247_739],
								            [0.30_340_484, 0.00_526_676, 0.31_077_722],
								            [0.32_278_043, -0.24_987_096, 0.3_414_628],
								        ]
								    ]				)
								self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()		,						expected_slice.numpy()		,						atol=1E-4				)				)






'''simple docstring'''

import argparse
import json
import os
from collections import OrderedDict

import numpy as np
import tensorflow as tf
import torch



def 							UpperCamelCase__			(			a__      ):


						'''simple docstring'''






						_lowerCAmelCase										=os.path.join(args.tf_model_dir			,      'parameters.json'      )
						_lowerCAmelCase										=json.loads(open(a__      ).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'      ):
												_lowerCAmelCase										=args.output + '.pt'
						_lowerCAmelCase										=OrderedDict()
						with tf.device('/CPU:0'      ):
												_lowerCAmelCase										=tf.train.load_checkpoint(args.tf_model_dir      )
												_lowerCAmelCase										=reader.get_variable_to_shape_map()
												for key_name in shapes.keys():
																		_lowerCAmelCase										=reader.get_tensor(a__      ).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'      ):
																														_lowerCAmelCase										=int(key_name[9]      )
																								elif key_name.startswith('pasts/out'      ):
																														_lowerCAmelCase										=8
																								_lowerCAmelCase										='model.sqout.%d.weight' % (player * 2)  # enter to nn.Sequencial with Tanh, so 2 at a time
																								_lowerCAmelCase										=vnp.transpose([1, 0]      ).copy()  # Mesh-Tensorflow is a diagonal matrix
																								_lowerCAmelCase										=torch.tensor(a__      )
																		elif key_name.startswith('model/moe'      ):
																								_lowerCAmelCase										=int(key_name[9:].split('/'      )[0]      )
																								if key_name.endswith('/switch_gating/kernel'      ):
																														_lowerCAmelCase										='model.blocks.%d.feed_forward.mlp.router.classifier.weight' % player
																														_lowerCAmelCase										=vnp.transpose([1, 0]      ).copy()  # Mesh-Tensorflow is a diagonal matrix
																														_lowerCAmelCase										=torch.tensor(a__      )
																								elif key_name.endswith('/softmlp/kernel'      ):
																														_lowerCAmelCase										='model.blocks.%d.feed_forward.soft_bypass_mlp.weight' % player
																														_lowerCAmelCase										=vnp.transpose([1, 0]      ).copy()  # Mesh-Tensorflow is a diagonal matrix
																														_lowerCAmelCase										=torch.tensor(a__      )
																								elif key_name.endswith('/wo/kernel'      ) or key_name.endswith('/wi/kernel'      ):
																														_lowerCAmelCase										=key_name[-9:-7]
																														for i in range(1_6      ):
																																				_lowerCAmelCase										='model.blocks.%d.feed_forward.mlp.experts.expert_%d.%s.weight' % (player, i, nlayer)
																																				_lowerCAmelCase										=(
																																				    vnp[i].transpose([1, 0]      ).copy()
																																				)  # In Mesh-Tensorflow, it is one array, so it is divided
																																				_lowerCAmelCase										=torch.tensor(a__      )
																		elif key_name.startswith('model/mlp'      ):
																								_lowerCAmelCase										=int(key_name[9:].split('/'      )[0]      )
																								if key_name.endswith('/p1/kernel'      ):
																														_lowerCAmelCase										='model.blocks.%d.feed_forward.mlp.wi.weight' % player
																														_lowerCAmelCase										=vnp.transpose([1, 0]      ).copy()  # Mesh-Tensorflow is a diagonal matrix
																														_lowerCAmelCase										=torch.tensor(a__      )
																								elif key_name.endswith('/p1/bias'      ):
																														_lowerCAmelCase										='model.blocks.%d.feed_forward.mlp.wi.bias' % player
																														_lowerCAmelCase										=vnp.copy()  # same because it is one dimensional
																														_lowerCAmelCase										=torch.tensor(a__      )
																								elif key_name.endswith('/p2/kernel'      ):
																														_lowerCAmelCase										='model.blocks.%d.feed_forward.mlp.wo.weight' % player
																														_lowerCAmelCase										=vnp.transpose([1, 0]      ).copy()  # Mesh-Tensorflow is a diagonal matrix
																														_lowerCAmelCase										=torch.tensor(a__      )
																								elif key_name.endswith('/p2/bias'      ):
																														_lowerCAmelCase										='model.blocks.%d.feed_forward.mlp.wo.bias' % player
																														_lowerCAmelCase										=vnp.copy()  # same because it is one dimensional
																														_lowerCAmelCase										=torch.tensor(a__      )
																		elif key_name.startswith('model/ln'      ):
																								_lowerCAmelCase										=int(key_name[8:].split('/'      )[0]      )
																								if key_name.endswith('/b'      ):
																														_lowerCAmelCase										='model.blocks.%d.feed_forward.norm.bias' % player
																														_lowerCAmelCase										=vnp.copy()  # same because it is one dimensional
																														_lowerCAmelCase										=torch.tensor(a__      )
																								elif key_name.endswith('/g'      ):
																														_lowerCAmelCase										='model.blocks.%d.feed_forward.norm.weight' % player
																														_lowerCAmelCase										=vnp.copy()  # same because it is one dimensional
																														_lowerCAmelCase										=torch.tensor(a__      )
																		elif key_name.startswith('model/att'      ):
																								_lowerCAmelCase										=int(key_name[9:].split('/'      )[0]      )
																								if key_name.endswith('/qkv/kernel'      ):
																														_lowerCAmelCase										=vnp.copy()  # Compute same dimension as Mesh-tensorflow using einsum
																														_lowerCAmelCase										=state[:, 0, :, :]
																														_lowerCAmelCase										=state[:, 1, :, :]
																														_lowerCAmelCase										=state[:, 2, :, :]
																														_lowerCAmelCase										=(
																														    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
																														_lowerCAmelCase										=(
																														    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
																														_lowerCAmelCase										=(
																														    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
																														_lowerCAmelCase										='model.blocks.%d.self_attn.self_attn.q_proj.weight' % player
																														_lowerCAmelCase										=torch.tensor(a__      )
																														_lowerCAmelCase										='model.blocks.%d.self_attn.self_attn.k_proj.weight' % player
																														_lowerCAmelCase										=torch.tensor(a__      )
																														_lowerCAmelCase										='model.blocks.%d.self_attn.self_attn.v_proj.weight' % player
																														_lowerCAmelCase										=torch.tensor(a__      )
																								elif key_name.endswith('/o/kernel'      ):
																														_lowerCAmelCase										='model.blocks.%d.self_attn.self_attn.out_proj.weight' % player
																														_lowerCAmelCase										=(
																														    vnp.reshape([vnp.shape[0] * vnp.shape[1], vnp.shape[2]]      ).transpose([1, 0]      ).copy()
																														)  # Mesh-Tensorflow is a diagonal matrix
																														_lowerCAmelCase										=torch.tensor(a__      )
																		elif key_name.startswith('model/an'      ):
																								_lowerCAmelCase										=int(key_name[8:].split('/'      )[0]      )
																								if key_name.endswith('/b'      ):
																														_lowerCAmelCase										='model.blocks.%d.self_attn.norm.bias' % player
																														_lowerCAmelCase										=vnp.copy()  # same because it is one dimensional
																														_lowerCAmelCase										=torch.tensor(a__      )
																								elif key_name.endswith('/g'      ):
																														_lowerCAmelCase										='model.blocks.%d.self_attn.norm.weight' % player
																														_lowerCAmelCase										=vnp.copy()  # same because it is one dimensional
																														_lowerCAmelCase										=torch.tensor(a__      )
																		elif (
																		    key_name.startswith('model/wte'      )
																		    or key_name.startswith('model/wpe'      )
																		    or key_name.startswith('model/ete'      )
																		):
																								_lowerCAmelCase										={'wte': 'embed_tokens', 'wpe': 'position_embeddings', 'ete': 'extra_position_embeddings'}[
																								    key_name[-3:]
																								]
																								_lowerCAmelCase										='model.%s.weight' % nlayer
																								_lowerCAmelCase										=vnp.copy()  # same in embedded
																								_lowerCAmelCase										=torch.tensor(a__      )
																								if key_name.startswith('model/wte'      ):
																														_lowerCAmelCase										='lm_head.weight'
																														_lowerCAmelCase										=vnp.copy()  # same in embedded
																														_lowerCAmelCase										=torch.tensor(a__      )
																		elif key_name.startswith('model/wob'      ):
																								_lowerCAmelCase										='final_logits_bias'
																								_lowerCAmelCase										=vnp.copy()  # same in embedded
																								_lowerCAmelCase										=state.reshape((1, -1)      )
																								_lowerCAmelCase										=torch.tensor(a__      )
																		elif key_name == "model/dense/kernel":
																								_lowerCAmelCase										='model.last_project.weight'
																								_lowerCAmelCase										=vnp.transpose([1, 0]      ).copy()  # Mesh-Tensorflow is a diagonal matrix
																								_lowerCAmelCase										=torch.tensor(a__      )
																		elif key_name == "model/dense_1/bias":
																								_lowerCAmelCase										='model.last_project.bias'
																								_lowerCAmelCase										=vnp.copy()  # same because it is one dimensional
																								_lowerCAmelCase										=torch.tensor(a__      )
						torch.save(a__			,      args.output      )


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





