



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 		lowercase						:






	"""simple docstring"""


	def __init__(     self     :						Optional[int]						,  __UpperCAmelCase     :						List[str]						,  __UpperCAmelCase     :						Tuple=13						,  __UpperCAmelCase     :						int=7						,  __UpperCAmelCase     :						Optional[int]=False						,  __UpperCAmelCase     :						Union[str, Any]=True						,  __UpperCAmelCase     :						str=False						,  __UpperCAmelCase     :						int=False						,  __UpperCAmelCase     :						Dict=19						,  __UpperCAmelCase     :						Optional[int]=32						,  __UpperCAmelCase     :						Union[str, Any]=5						,  __UpperCAmelCase     :						List[Any]=4						,  __UpperCAmelCase     :						Optional[Any]=37						,  __UpperCAmelCase     :						int="gelu"						,  __UpperCAmelCase     :						Any=0.1						,  __UpperCAmelCase     :						Optional[int]=0.1						,  __UpperCAmelCase     :						Any=512						,  __UpperCAmelCase     :						Tuple=16						,  __UpperCAmelCase     :						Union[str, Any]=2						,  __UpperCAmelCase     :						Any=0.02						,  __UpperCAmelCase     :						Tuple=3						,  __UpperCAmelCase     :						Dict=4						,  __UpperCAmelCase     :						Dict=None						,  )      ->			str:
							UpperCAmelCase_=       parent
							UpperCAmelCase_=       batch_size
							UpperCAmelCase_=       seq_length
							UpperCAmelCase_=       is_training
							UpperCAmelCase_=       use_input_mask
							UpperCAmelCase_=       use_token_type_ids
							UpperCAmelCase_=       use_labels
							UpperCAmelCase_=       vocab_size
							UpperCAmelCase_=       hidden_size
							UpperCAmelCase_=       num_hidden_layers
							UpperCAmelCase_=       num_attention_heads
							UpperCAmelCase_=       intermediate_size
							UpperCAmelCase_=       hidden_act
							UpperCAmelCase_=       hidden_dropout_prob
							UpperCAmelCase_=       attention_probs_dropout_prob
							UpperCAmelCase_=       max_position_embeddings
							UpperCAmelCase_=       type_vocab_size
							UpperCAmelCase_=       type_sequence_label_size
							UpperCAmelCase_=       initializer_range
							UpperCAmelCase_=       num_labels
							UpperCAmelCase_=       num_choices
							UpperCAmelCase_=       scope




	def   _SCREAMING_SNAKE_CASE		(     self     :						List[str]					)      ->			Dict:
							UpperCAmelCase_=       ids_tensor([self.batch_size, self.seq_length]						,  self.vocab_size					)

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

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

							UpperCAmelCase_=       self.get_config()

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




	def   _SCREAMING_SNAKE_CASE		(     self     :						List[Any]					)      ->			Union[str, Any]:
							UpperCAmelCase_=       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   _SCREAMING_SNAKE_CASE		(     self     :						List[str]						,  __UpperCAmelCase     :						List[str]						,  __UpperCAmelCase     :						Any						,  __UpperCAmelCase     :						List[Any]						,  __UpperCAmelCase     :						Dict						,  __UpperCAmelCase     :						Union[str, Any]						,  __UpperCAmelCase     :						Optional[int]					)      ->			Union[str, Any]:
							UpperCAmelCase_=       EsmForProteinFolding(config=__UpperCAmelCase					).float()
							model.to(__UpperCAmelCase					)
							model.eval()
							UpperCAmelCase_=       model(__UpperCAmelCase						,  attention_mask=__UpperCAmelCase					)
							UpperCAmelCase_=       model(__UpperCAmelCase					)
							UpperCAmelCase_=       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   _SCREAMING_SNAKE_CASE		(     self     :						Union[str, Any]					)      ->			int:
							UpperCAmelCase_=       self.prepare_config_and_inputs()
							(
							    (
							    UpperCAmelCase_       
							),       (
							    UpperCAmelCase_       
							),       (
							    UpperCAmelCase_       
							),       (
							    UpperCAmelCase_       
							),       (
							    UpperCAmelCase_       
							),       (
							    UpperCAmelCase_       
							),              
							)=       config_and_inputs
							UpperCAmelCase_=       {"""input_ids""": input_ids, """attention_mask""": input_mask}
							return config, inputs_dict






@require_torch
class 		lowercase						(			snake_case__			,   snake_case__			,   unittest.TestCase):






	"""simple docstring"""


	a__				: Dict		       =							False

	a__				: str		       =							(EsmForProteinFolding,) if is_torch_available() else ()
	a__				: Dict		       =							()
	a__				: Any		       =							{} if is_torch_available() else {}
	a__				: Union[str, Any]		       =							False




	def   _SCREAMING_SNAKE_CASE		(     self     :						Tuple					)      ->			Tuple:
							UpperCAmelCase_=       EsmFoldModelTester(self					)
							UpperCAmelCase_=       ConfigTester(self						,  config_class=__UpperCAmelCase						,  hidden_size=37					)




	def   _SCREAMING_SNAKE_CASE		(     self     :						Tuple					)      ->			List[str]:
							self.config_tester.run_common_tests()




	def   _SCREAMING_SNAKE_CASE		(     self     :						Dict					)      ->			int:
							UpperCAmelCase_=       self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*__UpperCAmelCase					)




	@unittest.skip("""Does not support attention outputs"""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						Optional[Any]					)      ->			List[str]:
							pass




	@unittest.skip
	def   _SCREAMING_SNAKE_CASE		(     self     :						str					)      ->			Any:
							pass




	@unittest.skip("""Esm does not support embedding resizing"""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						Union[str, Any]					)      ->			Optional[Any]:
							pass




	@unittest.skip("""Esm does not support embedding resizing"""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						int					)      ->			int:
							pass




	@unittest.skip("""ESMFold does not support passing input embeds!"""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						Union[str, Any]					)      ->			Optional[int]:
							pass




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




	@unittest.skip("""ESMFold does not support head pruning."""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						List[str]					)      ->			Tuple:
							pass




	@unittest.skip("""ESMFold does not support head pruning."""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						List[str]					)      ->			Union[str, Any]:
							pass




	@unittest.skip("""ESMFold does not support head pruning."""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						Dict					)      ->			str:
							pass




	@unittest.skip("""ESMFold does not support head pruning."""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						str					)      ->			str:
							pass




	@unittest.skip("""ESMFold does not output hidden states in the normal way."""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						str					)      ->			List[Any]:
							pass




	@unittest.skip("""ESMfold does not output hidden states in the normal way."""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						int					)      ->			Optional[int]:
							pass




	@unittest.skip("""ESMFold only has one output format."""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						int					)      ->			int:
							pass




	@unittest.skip("""This test doesn't work for ESMFold and doesn't test core functionality"""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						Tuple					)      ->			Optional[Any]:
							pass




	@unittest.skip("""ESMFold does not support input chunking."""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						Optional[Any]					)      ->			List[str]:
							pass




	@unittest.skip("""ESMFold doesn't respect you and it certainly doesn't respect your initialization arguments."""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						Union[str, Any]					)      ->			Tuple:
							pass




	@unittest.skip("""ESMFold doesn't support torchscript compilation."""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						int					)      ->			List[Any]:
							pass




	@unittest.skip("""ESMFold doesn't support torchscript compilation."""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						Optional[Any]					)      ->			Union[str, Any]:
							pass




	@unittest.skip("""ESMFold doesn't support torchscript compilation."""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						str					)      ->			str:
							pass




	@unittest.skip("""ESMFold doesn't support data parallel."""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						Optional[Any]					)      ->			List[Any]:
							pass




	@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""					)
	def   _SCREAMING_SNAKE_CASE		(     self     :						Dict					)      ->			Any:
							pass






@require_torch
class 		lowercase						(			snake_case__):






	"""simple docstring"""


	@slow
	def   _SCREAMING_SNAKE_CASE		(     self     :						Optional[Any]					)      ->			Any:
							UpperCAmelCase_=       EsmForProteinFolding.from_pretrained("""facebook/esmfold_v1"""					).float()
							model.eval()
							UpperCAmelCase_=       torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]					)
							UpperCAmelCase_=       model(__UpperCAmelCase					)["""positions"""]
							UpperCAmelCase_=       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					)					)









import argparse
from argparse import Namespace

import torch
from torch import nn

from transformers import XGLMConfig, XGLMForCausalLM



def 	__a					(   lowerCAmelCase_       : Optional[int]    )				->			List[Any]:






						'''simple docstring'''



						UpperCAmelCase_=       [
						    """decoder.version""",
						    """decoder.output_projection.weight""",
						    """_float_tensor""",
						    """decoder.embed_positions._float_tensor""",
						]
						for k in ignore_keys:
												state_dict.pop(lowerCAmelCase_      ,lowerCAmelCase_    )



def 	__a					(   lowerCAmelCase_       : Any    )				->			str:






						'''simple docstring'''



						UpperCAmelCase_,       UpperCAmelCase_=       emb.weight.shape
						UpperCAmelCase_=       nn.Linear(lowerCAmelCase_      ,lowerCAmelCase_      ,bias=lowerCAmelCase_    )
						UpperCAmelCase_=       emb.weight.data
						return lin_layer



def 	__a					(   lowerCAmelCase_       : List[str]    )				->			List[str]:






						'''simple docstring'''



						UpperCAmelCase_=       torch.load(lowerCAmelCase_      ,map_location="""cpu"""    )
						UpperCAmelCase_=       Namespace(**checkpoint["""cfg"""]["""model"""]    )
						UpperCAmelCase_=       checkpoint["""model"""]
						remove_ignore_keys_(lowerCAmelCase_    )
						UpperCAmelCase_=       state_dict["""decoder.embed_tokens.weight"""].shape[0]

						UpperCAmelCase_=       {key.replace("""decoder"""      ,"""model"""    ): val for key, val in state_dict.items()}

						UpperCAmelCase_=       XGLMConfig(
						    vocab_size=lowerCAmelCase_      ,max_position_embeddings=args.max_target_positions      ,num_layers=args.decoder_layers      ,attention_heads=args.decoder_attention_heads      ,ffn_dim=args.decoder_ffn_embed_dim      ,d_model=args.decoder_embed_dim      ,layerdrop=args.decoder_layerdrop      ,dropout=args.dropout      ,attention_dropout=args.attention_dropout      ,activation_dropout=args.activation_dropout      ,activation_function="""gelu"""      ,scale_embedding=not args.no_scale_embedding      ,tie_word_embeddings=args.share_decoder_input_output_embed      ,)

						UpperCAmelCase_=       XGLMForCausalLM(lowerCAmelCase_    )
						UpperCAmelCase_=       model.load_state_dict(lowerCAmelCase_      ,strict=lowerCAmelCase_    )
						print(lowerCAmelCase_    )
						UpperCAmelCase_=       make_linear_from_emb(model.model.embed_tokens    )

						return model


if __name__ == "__main__":
	__A       =						argparse.ArgumentParser()
	# Required parameters
	parser.add_argument('''fairseq_path''', type=str, help='''path to a model.pt on local filesystem.''')
	parser.add_argument('''pytorch_dump_folder_path''', default=None, type=str, help='''Path to the output PyTorch model.''')
	__A       =						parser.parse_args()
	__A       =						convert_fairseq_xglm_checkpoint_from_disk(args.fairseq_path)
	model.save_pretrained(args.pytorch_dump_folder_path)




