


import unittest

from transformers import EsmConfig, is_torch_available
from transformers.testing_utils import TestCasePlus, require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
							import torch

							from transformers.models.esm.modeling_esmfold import EsmForProteinFolding







class 							UpperCamelCase__ :

					def __init__(  self							: Optional[Any]				,       UpperCamelCase__							: Dict				,       UpperCamelCase__							: int=13				,       UpperCamelCase__							: Any=7				,       UpperCamelCase__							: Any=False				,       UpperCamelCase__							: Tuple=True				,       UpperCamelCase__							: List[Any]=False				,       UpperCamelCase__							: List[Any]=False				,       UpperCamelCase__							: Optional[int]=19				,       UpperCamelCase__							: Dict=32				,       UpperCamelCase__							: Tuple=5				,       UpperCamelCase__							: Optional[int]=4				,       UpperCamelCase__							: Union[str, Any]=37				,       UpperCamelCase__							: str="gelu"				,       UpperCamelCase__							: Dict=0.1				,       UpperCamelCase__							: str=0.1				,       UpperCamelCase__							: int=512				,       UpperCamelCase__							: Tuple=16				,       UpperCamelCase__							: Tuple=2				,       UpperCamelCase__							: Tuple=0.02				,       UpperCamelCase__							: Union[str, Any]=3				,       UpperCamelCase__							: Union[str, Any]=4				,       UpperCamelCase__							: int=None				,       ):






										'''simple docstring'''






										lowercase_					=      parent
										lowercase_					=      batch_size
										lowercase_					=      seq_length
										lowercase_					=      is_training
										lowercase_					=      use_input_mask
										lowercase_					=      use_token_type_ids
										lowercase_					=      use_labels
										lowercase_					=      vocab_size
										lowercase_					=      hidden_size
										lowercase_					=      num_hidden_layers
										lowercase_					=      num_attention_heads
										lowercase_					=      intermediate_size
										lowercase_					=      hidden_act
										lowercase_					=      hidden_dropout_prob
										lowercase_					=      attention_probs_dropout_prob
										lowercase_					=      max_position_embeddings
										lowercase_					=      type_vocab_size
										lowercase_					=      type_sequence_label_size
										lowercase_					=      initializer_range
										lowercase_					=      num_labels
										lowercase_					=      num_choices
										lowercase_					=      scope

					def 	UpperCAmelCase__							(  self							: Tuple							):






										'''simple docstring'''






										lowercase_					=      ids_tensor([self.batch_size, self.seq_length]				,       self.vocab_size							)

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

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

										lowercase_					=      self.get_config()

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

					def 	UpperCAmelCase__							(  self							: Optional[Any]							):






										'''simple docstring'''






										lowercase_					=      EsmConfig(
										    vocab_size=33				,       hidden_size=self.hidden_size				,       pad_token_id=1				,       num_hidden_layers=self.num_hidden_layers				,       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				,       is_folding_model=UpperCamelCase__				,       esmfold_config={"""trunk""": {"""num_blocks""": 2}, """fp16_esm""": False}				,       )
										return config

					def 	UpperCAmelCase__							(  self							: List[str]				,       UpperCamelCase__							: Tuple				,       UpperCamelCase__							: Dict				,       UpperCamelCase__							: Optional[Any]				,       UpperCamelCase__							: Tuple				,       UpperCamelCase__							: Optional[int]				,       UpperCamelCase__							: int							):






										'''simple docstring'''






										lowercase_					=      EsmForProteinFolding(config=UpperCamelCase__							).float()
										model.to(UpperCamelCase__							)
										model.eval()
										lowercase_					=      model(UpperCamelCase__				,       attention_mask=UpperCamelCase__							)
										lowercase_					=      model(UpperCamelCase__							)
										lowercase_					=      model(UpperCamelCase__							)

										self.parent.assertEqual(result.positions.shape				,       (8, self.batch_size, self.seq_length, 14, 3)							)
										self.parent.assertEqual(result.angles.shape				,       (8, self.batch_size, self.seq_length, 7, 2)							)







					def 	UpperCAmelCase__							(  self							: str							):






										'''simple docstring'''






										lowercase_					=      self.prepare_config_and_inputs()
										(
										    (
										    lowercase_
										)				,       (
										    lowercase_
										)				,       (
										    lowercase_
										)				,       (
										    lowercase_
										)				,       (
										    lowercase_
										)				,       (
										    lowercase_
										)				,       
										)					=      config_and_inputs
										lowercase_					=      {"""input_ids""": input_ids, """attention_mask""": input_mask}
										return config, inputs_dict







@require_torch
class 							UpperCamelCase__ (						lowercase_				,					lowercase_				,					unittest.TestCase ):
					__SCREAMING_SNAKE_CASE							:	Any            = False

					__SCREAMING_SNAKE_CASE							:	Optional[Any]            = (EsmForProteinFolding,) if is_torch_available() else ()
					__SCREAMING_SNAKE_CASE							:	List[Any]            = ()
					__SCREAMING_SNAKE_CASE							:	Optional[int]            = {} if is_torch_available() else {}
					__SCREAMING_SNAKE_CASE							:	Any            = False

					def 	UpperCAmelCase__							(  self							: Dict							):






										'''simple docstring'''






										lowercase_					=      EsmFoldModelTester(self							)
										lowercase_					=      ConfigTester(self				,       config_class=UpperCamelCase__				,       hidden_size=37							)

					def 	UpperCAmelCase__							(  self							: Any							):






										'''simple docstring'''






										self.config_tester.run_common_tests()

					def 	UpperCAmelCase__							(  self							: Any							):






										'''simple docstring'''






										lowercase_					=      self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*UpperCamelCase__							)

					@unittest.skip("""Does not support attention outputs"""							)
					def 	UpperCAmelCase__							(  self							: Optional[Any]							):






										'''simple docstring'''






										pass

					@unittest.skip
					def 	UpperCAmelCase__							(  self							: List[Any]							):






										'''simple docstring'''






										pass

					@unittest.skip("""Esm does not support embedding resizing"""							)
					def 	UpperCAmelCase__							(  self							: str							):






										'''simple docstring'''






										pass

					@unittest.skip("""Esm does not support embedding resizing"""							)
					def 	UpperCAmelCase__							(  self							: Tuple							):






										'''simple docstring'''






										pass

					@unittest.skip("""ESMFold does not support passing input embeds!"""							)
					def 	UpperCAmelCase__							(  self							: Optional[int]							):






										'''simple docstring'''






										pass

					@unittest.skip("""ESMFold does not support head pruning."""							)
					def 	UpperCAmelCase__							(  self							: Tuple							):






										'''simple docstring'''






										pass

					@unittest.skip("""ESMFold does not support head pruning."""							)
					def 	UpperCAmelCase__							(  self							: Tuple							):






										'''simple docstring'''






										pass

					@unittest.skip("""ESMFold does not support head pruning."""							)
					def 	UpperCAmelCase__							(  self							: int							):






										'''simple docstring'''






										pass

					@unittest.skip("""ESMFold does not support head pruning."""							)
					def 	UpperCAmelCase__							(  self							: Optional[Any]							):






										'''simple docstring'''






										pass

					@unittest.skip("""ESMFold does not support head pruning."""							)
					def 	UpperCAmelCase__							(  self							: Optional[int]							):






										'''simple docstring'''






										pass

					@unittest.skip("""ESMFold does not output hidden states in the normal way."""							)
					def 	UpperCAmelCase__							(  self							: Tuple							):






										'''simple docstring'''






										pass

					@unittest.skip("""ESMfold does not output hidden states in the normal way."""							)
					def 	UpperCAmelCase__							(  self							: int							):






										'''simple docstring'''






										pass

					@unittest.skip("""ESMFold only has one output format."""							)
					def 	UpperCAmelCase__							(  self							: List[str]							):






										'''simple docstring'''






										pass

					@unittest.skip("""This test doesn\'t work for ESMFold and doesn\'t test core functionality"""							)
					def 	UpperCAmelCase__							(  self							: Tuple							):






										'''simple docstring'''






										pass

					@unittest.skip("""ESMFold does not support input chunking."""							)
					def 	UpperCAmelCase__							(  self							: List[str]							):






										'''simple docstring'''






										pass

					@unittest.skip("""ESMFold doesn\'t respect you and it certainly doesn\'t respect your initialization arguments."""							)
					def 	UpperCAmelCase__							(  self							: Tuple							):






										'''simple docstring'''






										pass

					@unittest.skip("""ESMFold doesn\'t support torchscript compilation."""							)
					def 	UpperCAmelCase__							(  self							: List[str]							):






										'''simple docstring'''






										pass

					@unittest.skip("""ESMFold doesn\'t support torchscript compilation."""							)
					def 	UpperCAmelCase__							(  self							: Optional[int]							):






										'''simple docstring'''






										pass

					@unittest.skip("""ESMFold doesn\'t support torchscript compilation."""							)
					def 	UpperCAmelCase__							(  self							: str							):






										'''simple docstring'''






										pass

					@unittest.skip("""ESMFold doesn\'t support data parallel."""							)
					def 	UpperCAmelCase__							(  self							: Optional[Any]							):






										'''simple docstring'''






										pass







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






										'''simple docstring'''






										pass







@require_torch
class 							UpperCamelCase__ (						lowercase_ ):

					@slow
					def 	UpperCAmelCase__							(  self							: str							):






										'''simple docstring'''






										lowercase_					=      EsmForProteinFolding.from_pretrained("""facebook/esmfold_v1"""							).float()
										model.eval()
										lowercase_					=      torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]							)
										lowercase_					=      model(UpperCamelCase__							)["""positions"""]
										lowercase_					=      torch.tensor([2.5_828, 0.7_993, -10.9_334]				,       dtype=torch.floataa							)
										self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]				,       UpperCamelCase__				,       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 _SCREAMING_SNAKE_CASE (A			)			->			Dict:






       """simple docstring"""







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


if __name__ == "__main__":
 lowerCamelCase				:  Optional[int]    		=      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')
 lowerCamelCase				:  Dict    		=      parser.parse_args()
 convert_tf_gptsan_to_pt(args)




