





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				,						lowerCAmelCase_				,						lowerCAmelCase_=13				,						lowerCAmelCase_=7				,						lowerCAmelCase_=False				,						lowerCAmelCase_=True				,						lowerCAmelCase_=False				,						lowerCAmelCase_=False				,						lowerCAmelCase_=19				,						lowerCAmelCase_=32				,						lowerCAmelCase_=5				,						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				,						)		->							str:
										_snake_case									=				parent
										_snake_case									=				batch_size
										_snake_case									=				seq_length
										_snake_case									=				is_training
										_snake_case									=				use_input_mask
										_snake_case									=				use_token_type_ids
										_snake_case									=				use_labels
										_snake_case									=				vocab_size
										_snake_case									=				hidden_size
										_snake_case									=				num_hidden_layers
										_snake_case									=				num_attention_heads
										_snake_case									=				intermediate_size
										_snake_case									=				hidden_act
										_snake_case									=				hidden_dropout_prob
										_snake_case									=				attention_probs_dropout_prob
										_snake_case									=				max_position_embeddings
										_snake_case									=				type_vocab_size
										_snake_case									=				type_sequence_label_size
										_snake_case									=				initializer_range
										_snake_case									=				num_labels
										_snake_case									=				num_choices
										_snake_case									=				scope




					def 	lowerCAmelCase			(		self  )		->							List[Any]:
										_snake_case									=				ids_tensor([self.batch_size, self.seq_length]				,						self.vocab_size  )

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

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

										_snake_case									=				self.get_config()

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




					def 	lowerCAmelCase			(		self  )		->							List[str]:
										_snake_case									=				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=lowerCAmelCase_				,						esmfold_config={'trunk': {'num_blocks': 2}, 'fp16_esm': False}				,						)
										return config




					def 	lowerCAmelCase			(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_  )		->							int:
										_snake_case									=				EsmForProteinFolding(config=lowerCAmelCase_  ).float()
										model.to(lowerCAmelCase_  )
										model.eval()
										_snake_case									=				model(lowerCAmelCase_				,						attention_mask=lowerCAmelCase_  )
										_snake_case									=				model(lowerCAmelCase_  )
										_snake_case									=				model(lowerCAmelCase_  )

										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 	lowerCAmelCase			(		self  )		->							Tuple:
										_snake_case									=				self.prepare_config_and_inputs()
										(
										    (
										    _snake_case
										)							,     (
										    _snake_case
										)							,     (
										    _snake_case
										)							,     (
										    _snake_case
										)							,     (
										    _snake_case
										)							,     (
										    _snake_case
										)							,     
										)									=				config_and_inputs
										_snake_case									=				{'input_ids': input_ids, 'attention_mask': input_mask}
										return config, inputs_dict





@require_torch
class 			UpperCamelCase_		(  _lowerCamelCase							,  _lowerCamelCase							,  unittest.TestCase					):
					lowerCAmelCase_							=	False

					lowerCAmelCase_							=	(EsmForProteinFolding,) if is_torch_available() else ()
					lowerCAmelCase_							=	()
					lowerCAmelCase_							=	{} if is_torch_available() else {}
					lowerCAmelCase_							=	False




					def 	lowerCAmelCase			(		self  )		->							Optional[int]:
										_snake_case									=				EsmFoldModelTester(self  )
										_snake_case									=				ConfigTester(self				,						config_class=lowerCAmelCase_				,						hidden_size=37  )




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




					def 	lowerCAmelCase			(		self  )		->							str:
										_snake_case									=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*lowerCAmelCase_  )




					@unittest.skip('Does not support attention outputs'  )
					def 	lowerCAmelCase			(		self  )		->							Tuple:
										pass




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




					@unittest.skip('Esm does not support embedding resizing'  )
					def 	lowerCAmelCase			(		self  )		->							Optional[Any]:
										pass




					@unittest.skip('Esm does not support embedding resizing'  )
					def 	lowerCAmelCase			(		self  )		->							Optional[int]:
										pass




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




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




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




					@unittest.skip('ESMFold does not support head pruning.'  )
					def 	lowerCAmelCase			(		self  )		->							Tuple:
										pass




					@unittest.skip('ESMFold does not support head pruning.'  )
					def 	lowerCAmelCase			(		self  )		->							Optional[Any]:
										pass




					@unittest.skip('ESMFold does not support head pruning.'  )
					def 	lowerCAmelCase			(		self  )		->							List[Any]:
										pass




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




					@unittest.skip('ESMfold does not output hidden states in the normal way.'  )
					def 	lowerCAmelCase			(		self  )		->							Any:
										pass




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




					@unittest.skip('This test doesn\'t work for ESMFold and doesn\'t test core functionality'  )
					def 	lowerCAmelCase			(		self  )		->							int:
										pass




					@unittest.skip('ESMFold does not support input chunking.'  )
					def 	lowerCAmelCase			(		self  )		->							Optional[int]:
										pass




					@unittest.skip('ESMFold doesn\'t respect you and it certainly doesn\'t respect your initialization arguments.'  )
					def 	lowerCAmelCase			(		self  )		->							Dict:
										pass




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




					@unittest.skip('ESMFold doesn\'t support torchscript compilation.'  )
					def 	lowerCAmelCase			(		self  )		->							List[str]:
										pass




					@unittest.skip('ESMFold doesn\'t support torchscript compilation.'  )
					def 	lowerCAmelCase			(		self  )		->							Optional[Any]:
										pass




					@unittest.skip('ESMFold doesn\'t support data parallel.'  )
					def 	lowerCAmelCase			(		self  )		->							Tuple:
										pass






					@unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'  )
					def 	lowerCAmelCase			(		self  )		->							str:
										pass





@require_torch
class 			UpperCamelCase_		(  _lowerCamelCase					):




					@slow
					def 	lowerCAmelCase			(		self  )		->							List[Any]:
										_snake_case									=				EsmForProteinFolding.from_pretrained('facebook/esmfold_v1'  ).float()
										model.eval()
										_snake_case									=				torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]  )
										_snake_case									=				model(lowerCAmelCase_  )['positions']
										_snake_case									=				torch.tensor([2.58_28, 0.79_93, -10.93_34]				,						dtype=torch.floataa  )
										self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]				,						lowerCAmelCase_				,						atol=1E-4  )  )










import argparse
from collections import defaultdict

import yaml


UpperCAmelCase_      							=     """docs/source/en/_toctree.yml"""





def 						lowerCamelCase__     (     UpperCamelCase__			:       Optional[Any]						)			-> str:



					'''simple docstring'''




					_snake_case									=				defaultdict(UpperCamelCase__						)
					_snake_case									=				[]
					_snake_case									=				[]
					for doc in doc_list:
										if "local" in doc:
															counts[doc["local"]] += 1

										if doc["title"].lower() == "overview":
															overview_doc.append({'local': doc['local'], 'title': doc['title']}						)
										else:
															new_doc_list.append(UpperCamelCase__						)

					_snake_case									=				new_doc_list
					_snake_case									=				[key for key, value in counts.items() if value > 1]

					_snake_case									=				[]
					for duplicate_key in duplicates:
										_snake_case									=				list({doc['title'] for doc in doc_list if doc['local'] == duplicate_key}						)
										if len(UpperCamelCase__						) > 1:
															raise ValueError(
															    F'''{duplicate_key} is present several times in the documentation table of content at '''
															    '`docs/source/en/_toctree.yml` with different *Title* values. Choose one of those and remove the '
															    'others.'						)
										# Only add this once
										new_doc.append({'local': duplicate_key, 'title': titles[0]}						)

					# Add none duplicate-keys
					new_doc.extend([doc for doc in doc_list if 'local' not in counts or counts[doc['local']] == 1]						)
					_snake_case									=				sorted(UpperCamelCase__       ,    key=lambda UpperCamelCase__						: s["title"].lower()						)

					# "overview" gets special treatment and is always first
					if len(UpperCamelCase__						) > 1:
										raise ValueError('{doc_list} has two \'overview\' docs which is not allowed.'						)

					overview_doc.extend(UpperCamelCase__						)

					# Sort
					return overview_doc





def 						lowerCamelCase__     (     UpperCamelCase__			:       Dict=False						)			-> Optional[int]:



					'''simple docstring'''




					with open(UpperCamelCase__       ,    encoding='utf-8'						) as f:
										_snake_case									=				yaml.safe_load(f.read()						)

					# Get to the API doc
					_snake_case									=				0
					while content[api_idx]["title"] != "API":
										api_idx += 1
					_snake_case									=				content[api_idx]['sections']

					# Then to the model doc
					_snake_case									=				0
					while api_doc[scheduler_idx]["title"] != "Schedulers":
										scheduler_idx += 1

					_snake_case									=				api_doc[scheduler_idx]['sections']
					_snake_case									=				clean_doc_toc(UpperCamelCase__						)

					_snake_case									=				False
					if new_scheduler_doc != scheduler_doc:
										_snake_case									=				True
										if overwrite:
															_snake_case									=				new_scheduler_doc

					if diff:
										if overwrite:
															_snake_case									=				api_doc
															with open(UpperCamelCase__       ,    'w'       ,    encoding='utf-8'						) as f:
																				f.write(yaml.dump(UpperCamelCase__       ,    allow_unicode=UpperCamelCase__						)						)
										else:
															raise ValueError(
															    'The model doc part of the table of content is not properly sorted, run `make style` to fix this.'						)





def 						lowerCamelCase__     (     UpperCamelCase__			:       Tuple=False						)			-> List[Any]:



					'''simple docstring'''




					with open(UpperCamelCase__       ,    encoding='utf-8'						) as f:
										_snake_case									=				yaml.safe_load(f.read()						)

					# Get to the API doc
					_snake_case									=				0
					while content[api_idx]["title"] != "API":
										api_idx += 1
					_snake_case									=				content[api_idx]['sections']

					# Then to the model doc
					_snake_case									=				0
					while api_doc[pipeline_idx]["title"] != "Pipelines":
										pipeline_idx += 1

					_snake_case									=				False
					_snake_case									=				api_doc[pipeline_idx]['sections']
					_snake_case									=				[]

					# sort sub pipeline docs
					for pipeline_doc in pipeline_docs:
										if "section" in pipeline_doc:
															_snake_case									=				pipeline_doc['section']
															_snake_case									=				clean_doc_toc(UpperCamelCase__						)
															if overwrite:
																				_snake_case									=				new_sub_pipeline_doc
										new_pipeline_docs.append(UpperCamelCase__						)

					# sort overall pipeline doc
					_snake_case									=				clean_doc_toc(UpperCamelCase__						)

					if new_pipeline_docs != pipeline_docs:
										_snake_case									=				True
										if overwrite:
															_snake_case									=				new_pipeline_docs

					if diff:
										if overwrite:
															_snake_case									=				api_doc
															with open(UpperCamelCase__       ,    'w'       ,    encoding='utf-8'						) as f:
																				f.write(yaml.dump(UpperCamelCase__       ,    allow_unicode=UpperCamelCase__						)						)
										else:
															raise ValueError(
															    'The model doc part of the table of content is not properly sorted, run `make style` to fix this.'						)


if __name__ == "__main__":
					UpperCAmelCase_      							=     argparse.ArgumentParser()
					parser.add_argument("""--fix_and_overwrite""", action="""store_true""", help="""Whether to fix inconsistencies.""")
					UpperCAmelCase_      							=     parser.parse_args()

					check_scheduler_doc(args.fix_and_overwrite)
					check_pipeline_doc(args.fix_and_overwrite)



