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



	"""simple docstring"""




	def __init__(       self  :			Tuple       ,							SCREAMING_SNAKE_CASE_  :			Union[str, Any]       ,							)       ->     Any:
					lowercase_       =					parent
					lowercase_       =					1_3
					lowercase_       =					7
					lowercase_       =					True
					lowercase_       =					True
					lowercase_       =					True
					lowercase_       =					9_9
					lowercase_       =					3_2
					lowercase_       =					2
					lowercase_       =					4
					lowercase_       =					3_7
					lowercase_       =					'''gelu'''
					lowercase_       =					0.1
					lowercase_       =					0.1
					lowercase_       =					5_1_2
					lowercase_       =					1_6
					lowercase_       =					2
					lowercase_       =					0.02
					lowercase_       =					3
					lowercase_       =					4
					lowercase_       =					None
	def 						_lowercase  (       self  :			Optional[int]       )       ->     Tuple:
					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_       =					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 						_lowercase  (       self  :			Any       )       ->     Optional[Any]:
					(
					    (
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	(
					    lowercase_
					)  ,	
					)       =					self.prepare_config_and_inputs()

					lowercase_       =					True
					lowercase_       =					floats_tensor([self.batch_size, self.seq_length, self.hidden_size]       )
					lowercase_       =					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 						_lowercase  (       self  :			Tuple       ,							SCREAMING_SNAKE_CASE_  :			Any       ,							SCREAMING_SNAKE_CASE_  :			List[str]       ,							SCREAMING_SNAKE_CASE_  :			str       ,							SCREAMING_SNAKE_CASE_  :			List[str]       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]       ,							SCREAMING_SNAKE_CASE_  :			str       )       ->     List[str]:
					lowercase_       =					TFEsmModel(config=SCREAMING_SNAKE_CASE_       )
					lowercase_       =					{'''input_ids''': input_ids, '''attention_mask''': input_mask}
					lowercase_       =					model(SCREAMING_SNAKE_CASE_       )

					lowercase_       =					[input_ids, input_mask]
					lowercase_       =					model(SCREAMING_SNAKE_CASE_       )

					lowercase_       =					model(SCREAMING_SNAKE_CASE_       )

					self.parent.assertEqual(result.last_hidden_state.shape       ,							(self.batch_size, self.seq_length, self.hidden_size)       )
	def 						_lowercase  (       self  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]       ,							SCREAMING_SNAKE_CASE_  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			Dict       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]       ,							SCREAMING_SNAKE_CASE_  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			int       ,							SCREAMING_SNAKE_CASE_  :			int       ,							SCREAMING_SNAKE_CASE_  :			str       ,							)       ->     Union[str, Any]:
					lowercase_       =					True

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

					lowercase_       =					[input_ids, input_mask]
					lowercase_       =					model(SCREAMING_SNAKE_CASE_       ,							encoder_hidden_states=SCREAMING_SNAKE_CASE_       )

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

					self.parent.assertEqual(result.last_hidden_state.shape       ,							(self.batch_size, self.seq_length, self.hidden_size)       )
	def 						_lowercase  (       self  :			Optional[int]       ,							SCREAMING_SNAKE_CASE_  :			List[str]       ,							SCREAMING_SNAKE_CASE_  :			str       ,							SCREAMING_SNAKE_CASE_  :			Dict       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]       ,							SCREAMING_SNAKE_CASE_  :			Any       ,							SCREAMING_SNAKE_CASE_  :			Any       )       ->     Any:
					lowercase_       =					TFEsmForMaskedLM(config=SCREAMING_SNAKE_CASE_       )
					lowercase_       =					model([input_ids, input_mask]       )
					self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.seq_length, self.vocab_size)       )
	def 						_lowercase  (       self  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]       ,							SCREAMING_SNAKE_CASE_  :			Union[str, Any]       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]       ,							SCREAMING_SNAKE_CASE_  :			str       )       ->     int:
					lowercase_       =					self.num_labels
					lowercase_       =					TFEsmForTokenClassification(config=SCREAMING_SNAKE_CASE_       )
					lowercase_       =					{'''input_ids''': input_ids, '''attention_mask''': input_mask}
					lowercase_       =					model(SCREAMING_SNAKE_CASE_       )
					self.parent.assertEqual(result.logits.shape       ,							(self.batch_size, self.seq_length, self.num_labels)       )




	def 						_lowercase  (       self  :			Tuple       )       ->     str:
					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_tf
class  lowercase__(				UpperCAmelCase ,						UpperCAmelCase ,						unittest.TestCase							):



	"""simple docstring"""




	a			:str         =	(
	    (
	        TFEsmModel,
	        TFEsmForMaskedLM,
	        TFEsmForSequenceClassification,
	        TFEsmForTokenClassification,
	    )
	    if is_tf_available()
	    else ()
	)
	a			:Union[str, Any]         =	(
	    {
	        'feature-extraction': TFEsmModel,
	        'fill-mask': TFEsmForMaskedLM,
	        'text-classification': TFEsmForSequenceClassification,
	        'token-classification': TFEsmForTokenClassification,
	        'zero-shot': TFEsmForSequenceClassification,
	    }
	    if is_tf_available()
	    else {}
	)
	a			:Union[str, Any]         =	False
	a			:Dict         =	False
	def 						_lowercase  (       self  :			List[str]       )       ->     List[Any]:
					lowercase_       =					TFEsmModelTester(self       )
					lowercase_       =					ConfigTester(self       ,							config_class=SCREAMING_SNAKE_CASE_       ,							hidden_size=3_7       )
	def 						_lowercase  (       self  :			Optional[int]       )       ->     str:
					self.config_tester.run_common_tests()
	def 						_lowercase  (       self  :			str       )       ->     int:
					lowercase_       =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_       )
	def 						_lowercase  (       self  :			Any       )       ->     Union[str, Any]:
					lowercase_       =					self.model_tester.prepare_config_and_inputs_for_decoder()
					self.model_tester.create_and_check_model_as_decoder(*SCREAMING_SNAKE_CASE_       )
	def 						_lowercase  (       self  :			str       )       ->     Optional[int]:
					lowercase_       =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_masked_lm(*SCREAMING_SNAKE_CASE_       )
	def 						_lowercase  (       self  :			Optional[Any]       )       ->     List[str]:
					lowercase_       =					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_token_classification(*SCREAMING_SNAKE_CASE_       )
	@slow
	def 						_lowercase  (       self  :			int       )       ->     List[str]:
					for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
									lowercase_       =					TFEsmModel.from_pretrained(SCREAMING_SNAKE_CASE_       )
									self.assertIsNotNone(SCREAMING_SNAKE_CASE_       )
	@unittest.skip('''Protein models do not support embedding resizing.'''       )
	def 						_lowercase  (       self  :			str       )       ->     Tuple:
					pass
	@unittest.skip('''Protein models do not support embedding resizing.'''       )
	def 						_lowercase  (       self  :			List[Any]       )       ->     Any:
					pass




	def 						_lowercase  (       self  :			List[str]       )       ->     Any:
					lowercase_  ,	lowercase_       =					self.model_tester.prepare_config_and_inputs_for_common()

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

@require_tf
class  lowercase__(				unittest.TestCase							):



	"""simple docstring"""




	@slow
	def 						_lowercase  (       self  :			str       )       ->     Optional[int]:
					lowercase_       =					TFEsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''       )

					lowercase_       =					tf.constant([[0, 1, 2, 3, 4, 5]]       )
					lowercase_       =					model(SCREAMING_SNAKE_CASE_       )[0]
					lowercase_       =					[1, 6, 3_3]
					self.assertEqual(list(output.numpy().shape       )       ,							SCREAMING_SNAKE_CASE_       )
					# compare the actual values for a slice.
					lowercase_       =					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 						_lowercase  (       self  :			int       )       ->     Union[str, Any]:
					lowercase_       =					TFEsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''       )

					lowercase_       =					tf.constant([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]       )
					lowercase_       =					model(SCREAMING_SNAKE_CASE_       )[0]
					# compare the actual values for a slice.
					lowercase_       =					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


__a  			=					'docs/source/en/_toctree.yml'




def 				a						(							snake_case__:					Dict					):


				'''simple docstring'''
				lowercase_       =					defaultdict(snake_case__					)
				for doc in model_doc:
								counts[doc["local"]] += 1
				lowercase_       =					[key for key, value in counts.items() if value > 1]

				lowercase_       =					[]
				for duplicate_key in duplicates:
								lowercase_       =					list({doc['''title'''] for doc in model_doc if doc['''local'''] == duplicate_key}					)
								if len(snake_case__					) > 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 model_doc if counts[doc['''local''']] == 1]					)

				# Sort
				return sorted(snake_case__						,   key=lambda snake_case__					: s["title"].lower()					)




def 				a						(							snake_case__:					List[Any]=False					):


				'''simple docstring'''
				with open(snake_case__						,   encoding='''utf-8'''					) as f:
								lowercase_       =					yaml.safe_load(f.read()					)

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

				# Then to the model doc
				lowercase_       =					0
				while api_doc[model_idx]["title"] != "Models":
								model_idx += 1

				lowercase_       =					api_doc[model_idx]['''sections''']

				lowercase_       =					[(idx, section) for idx, section in enumerate(snake_case__					) if '''sections''' in section]
				lowercase_       =					False
				for idx, modality_doc in modalities_docs:
								lowercase_       =					modality_doc['''sections''']
								lowercase_       =					clean_model_doc_toc(snake_case__					)

								if old_modality_doc != new_modality_doc:
												lowercase_       =					True
												if overwrite:
																lowercase_       =					new_modality_doc

				if diff:
								if overwrite:
												lowercase_       =					model_doc
												lowercase_       =					api_doc
												with open(snake_case__						,   '''w'''						,   encoding='''utf-8'''					) as f:
																f.write(yaml.dump(snake_case__						,   allow_unicode=snake_case__					)					)
								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__":
						__a  			=					argparse.ArgumentParser()
						parser.add_argument('--fix_and_overwrite', action='store_true', help='Whether to fix inconsistencies.')
						__a  			=					parser.parse_args()

						check_model_doc(args.fix_and_overwrite)




