





'''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 			UpperCamelCase_		:




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




					def 	lowerCAmelCase			(		self  )		->							Union[str, 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									=				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 	lowerCAmelCase			(		self  )		->							Dict:
										(
										    (
										    _snake_case      
										)							,     (
										    _snake_case      
										)							,     (
										    _snake_case      
										)							,     (
										    _snake_case      
										)							,     (
										    _snake_case      
										)							,     (
										    _snake_case      
										)							,           
										)									=				self.prepare_config_and_inputs()

										_snake_case									=				True
										_snake_case									=				floats_tensor([self.batch_size, self.seq_length, self.hidden_size]  )
										_snake_case									=				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 	lowerCAmelCase			(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_  )		->							Dict:
										_snake_case									=				TFEsmModel(config=A_  )
										_snake_case									=				{'input_ids': input_ids, 'attention_mask': input_mask}
										_snake_case									=				model(A_  )

										_snake_case									=				[input_ids, input_mask]
										_snake_case									=				model(A_  )

										_snake_case									=				model(A_  )

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




					def 	lowerCAmelCase			(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						)		->							Optional[int]:
										_snake_case									=				True

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

										_snake_case									=				[input_ids, input_mask]
										_snake_case									=				model(A_				,						encoder_hidden_states=A_  )

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

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




					def 	lowerCAmelCase			(		self				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_				,						lowerCAmelCase_  )		->							Dict:
										_snake_case									=				TFEsmForMaskedLM(config=A_  )
										_snake_case									=				model([input_ids, input_mask]  )
										self.parent.assertEqual(result.logits.shape				,						(self.batch_size, self.seq_length, self.vocab_size)  )




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






					def 	lowerCAmelCase			(		self  )		->							List[Any]:
										_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_tf
class 			UpperCamelCase_		(  _lowercase							,  _lowercase							,  unittest.TestCase					):
					lowerCAmelCase_							=	(
					    (
					        TFEsmModel,
					        TFEsmForMaskedLM,
					        TFEsmForSequenceClassification,
					        TFEsmForTokenClassification,
					    )
					    if is_tf_available()
					    else ()
					)
					lowerCAmelCase_							=	(
					    {
					        '''feature-extraction''': TFEsmModel,
					        '''fill-mask''': TFEsmForMaskedLM,
					        '''text-classification''': TFEsmForSequenceClassification,
					        '''token-classification''': TFEsmForTokenClassification,
					        '''zero-shot''': TFEsmForSequenceClassification,
					    }
					    if is_tf_available()
					    else {}
					)
					lowerCAmelCase_							=	False
					lowerCAmelCase_							=	False




					def 	lowerCAmelCase			(		self  )		->							int:
										_snake_case									=				TFEsmModelTester(self  )
										_snake_case									=				ConfigTester(self				,						config_class=A_				,						hidden_size=37  )




					def 	lowerCAmelCase			(		self  )		->							Optional[int]:
										self.config_tester.run_common_tests()




					def 	lowerCAmelCase			(		self  )		->							List[str]:
										_snake_case									=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*A_  )




					def 	lowerCAmelCase			(		self  )		->							List[Any]:
										_snake_case									=				self.model_tester.prepare_config_and_inputs_for_decoder()
										self.model_tester.create_and_check_model_as_decoder(*A_  )




					def 	lowerCAmelCase			(		self  )		->							Dict:
										_snake_case									=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_lm(*A_  )




					def 	lowerCAmelCase			(		self  )		->							List[Any]:
										_snake_case									=				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_token_classification(*A_  )




					@slow
					def 	lowerCAmelCase			(		self  )		->							Tuple:
										for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															_snake_case									=				TFEsmModel.from_pretrained(A_  )
															self.assertIsNotNone(A_  )




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




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






					def 	lowerCAmelCase			(		self  )		->							Dict:
										_snake_case							,     _snake_case									=				self.model_tester.prepare_config_and_inputs_for_common()

										for model_class in self.all_model_classes:
															_snake_case									=				model_class(A_  )
															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
																				_snake_case									=				model.get_bias()
																				assert isinstance(A_				,						A_  )
																				for k, v in name.items():
																									assert isinstance(A_				,						tf.Variable  )
															else:
																				_snake_case									=				model.get_output_embeddings()
																				assert x is None
																				_snake_case									=				model.get_bias()
																				assert name is None





@require_tf
class 			UpperCamelCase_		(  unittest.TestCase					):




					@slow
					def 	lowerCAmelCase			(		self  )		->							Tuple:
										_snake_case									=				TFEsmForMaskedLM.from_pretrained('facebook/esm2_t6_8M_UR50D'  )

										_snake_case									=				tf.constant([[0, 1, 2, 3, 4, 5]]  )
										_snake_case									=				model(A_  )[0]
										_snake_case									=				[1, 6, 33]
										self.assertEqual(list(output.numpy().shape  )				,						A_  )
										# compare the actual values for a slice.
										_snake_case									=				tf.constant(
										    [
										        [
										            [8.92_15_18, -10.58_98_14, -6.4_67_13_07],
										            [-6.3_96_71_56, -13.91_13_77, -1.1_21_19_15],
										            [-7.78_12_47, -13.95_15_57, -3.74_05_92],
										        ]
										    ]  )
										self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()				,						expected_slice.numpy()				,						atol=1E-2  )  )






					@slow
					def 	lowerCAmelCase			(		self  )		->							int:
										_snake_case									=				TFEsmModel.from_pretrained('facebook/esm2_t6_8M_UR50D'  )

										_snake_case									=				tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]  )
										_snake_case									=				model(A_  )[0]
										# compare the actual values for a slice.
										_snake_case									=				tf.constant(
										    [
										        [
										            [0.14_44_30_92, 0.54_12_53_27, 0.3_24_77_39],
										            [0.30_34_04_84, 0.00_52_66_76, 0.31_07_77_22],
										            [0.32_27_80_43, -0.24_98_70_96, 0.3_41_46_28],
										        ]
										    ]  )
										self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()				,						expected_slice.numpy()				,						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)



