




'''simple docstring'''





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 import EsmForMaskedLM, EsmForSequenceClassification, EsmForTokenClassification, EsmModel
		from transformers.models.esm.modeling_esm import (
		    ESM_PRETRAINED_MODEL_ARCHIVE_LIST,
		    EsmEmbeddings,
		    create_position_ids_from_input_ids,
		)




class 			_snake_case		:




			def __init__( self    ,							a__    ,							a__=13    ,							a__=7    ,							a__=False    ,							a__=True    ,							a__=False    ,							a__=True    ,							a__=33    ,							a__=32    ,							a__=5    ,							a__=4    ,							a__=37    ,							a__="gelu"    ,							a__=0.1    ,							a__=0.1    ,							a__=512    ,							a__=16    ,							a__=2    ,							a__=0.0_2    ,							a__=3    ,							a__=4    ,							a__=None    ,							)    ->    Dict:

										'''simple docstring'''




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

										'''simple docstring'''




										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							)    ->    Tuple:

										'''simple docstring'''




										return EsmConfig(
										    vocab_size=self.vocab_size    ,							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    ,							)




			def        lowerCAmelCase__    ( self    ,							a__    ,							a__    ,							a__    ,							a__    ,							a__    ,							a__							)    ->    str:

										'''simple docstring'''




										snake_case_      =				EsmModel(config=__A							)
										model.to(__A							)
										model.eval()
										snake_case_      =				model(__A    ,							attention_mask=__A							)
										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)							)
										self.parent.assertEqual(result.pooler_output.shape    ,							(self.batch_size, self.hidden_size)							)




			def        lowerCAmelCase__    ( self    ,							a__    ,							a__    ,							a__    ,							a__    ,							a__    ,							a__							)    ->    List[str]:

										'''simple docstring'''




										snake_case_      =				EsmForMaskedLM(config=__A							)
										model.to(__A							)
										model.eval()
										snake_case_      =				model(__A    ,							attention_mask=__A    ,							labels=__A							)
										self.parent.assertEqual(result.logits.shape    ,							(self.batch_size, self.seq_length, self.vocab_size)							)




			def        lowerCAmelCase__    ( self    ,							a__    ,							a__    ,							a__    ,							a__    ,							a__    ,							a__							)    ->    int:

										'''simple docstring'''




										snake_case_      =				self.num_labels
										snake_case_      =				EsmForTokenClassification(config=__A							)
										model.to(__A							)
										model.eval()
										snake_case_      =				model(__A    ,							attention_mask=__A    ,							labels=__A							)
										self.parent.assertEqual(result.logits.shape    ,							(self.batch_size, self.seq_length, self.num_labels)							)




			def        lowerCAmelCase__    ( self							)    ->    List[Any]:

										'''simple docstring'''




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




@require_torch
class 			_snake_case		(     A__						,						A__						,						unittest.TestCase							):
			lowerCAmelCase_					:				Dict							=   False

			lowerCAmelCase_					:				List[Any]							=   (
			    (
			        EsmForMaskedLM,
			        EsmModel,
			        EsmForSequenceClassification,
			        EsmForTokenClassification,
			    )
			    if is_torch_available()
			    else ()
			)
			lowerCAmelCase_					:				Union[str, Any]							=   ()
			lowerCAmelCase_					:				List[str]							=   (
			    {
			        "feature-extraction": EsmModel,
			        "fill-mask": EsmForMaskedLM,
			        "text-classification": EsmForSequenceClassification,
			        "token-classification": EsmForTokenClassification,
			        "zero-shot": EsmForSequenceClassification,
			    }
			    if is_torch_available()
			    else {}
			)
			lowerCAmelCase_					:				Any							=   True




			def        lowerCAmelCase__    ( self							)    ->    str:

										'''simple docstring'''




										snake_case_      =				EsmModelTester(self							)
										snake_case_      =				ConfigTester(self    ,							config_class=__A    ,							hidden_size=37							)




			def        lowerCAmelCase__    ( self							)    ->    List[Any]:

										'''simple docstring'''




										self.config_tester.run_common_tests()




			def        lowerCAmelCase__    ( self							)    ->    str:

										'''simple docstring'''




										snake_case_      =				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_model(*__A							)




			def        lowerCAmelCase__    ( self							)    ->    Dict:

										'''simple docstring'''




										snake_case_      =				self.model_tester.prepare_config_and_inputs()
										for type in ["absolute", "relative_key", "relative_key_query"]:
																	snake_case_      =				type
																	self.model_tester.create_and_check_model(*__A							)




			def        lowerCAmelCase__    ( self							)    ->    Tuple:

										'''simple docstring'''




										snake_case_      =				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_masked_lm(*__A							)




			def        lowerCAmelCase__    ( self							)    ->    Any:

										'''simple docstring'''




										snake_case_      =				self.model_tester.prepare_config_and_inputs()
										self.model_tester.create_and_check_for_token_classification(*__A							)




			@slow
			def        lowerCAmelCase__    ( self							)    ->    Union[str, Any]:

										'''simple docstring'''




										for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																	snake_case_      =				EsmModel.from_pretrained(__A							)
																	self.assertIsNotNone(__A							)




			def        lowerCAmelCase__    ( self							)    ->    str:

										'''simple docstring'''




										snake_case_      =				self.model_tester.prepare_config_and_inputs()[0]
										snake_case_      =				EsmEmbeddings(config=__A							)

										snake_case_      =				torch.as_tensor([[12, 31, 13, model.padding_idx]]							)
										snake_case_      =				torch.as_tensor(
										    [
										        [
										            0 + model.padding_idx + 1,
										            1 + model.padding_idx + 1,
										            2 + model.padding_idx + 1,
										            model.padding_idx,
										        ]
										    ]							)
										snake_case_      =				create_position_ids_from_input_ids(__A    ,							model.padding_idx							)
										self.assertEqual(position_ids.shape    ,							expected_positions.shape							)
										self.assertTrue(torch.all(torch.eq(__A    ,							__A							)							)							)




			def        lowerCAmelCase__    ( self							)    ->    List[str]:

										'''simple docstring'''




										snake_case_      =				self.model_tester.prepare_config_and_inputs()[0]
										snake_case_      =				EsmEmbeddings(config=__A							)

										snake_case_      =				torch.empty(2    ,							4    ,							30							)
										snake_case_      =				[
										    0 + embeddings.padding_idx + 1,
										    1 + embeddings.padding_idx + 1,
										    2 + embeddings.padding_idx + 1,
										    3 + embeddings.padding_idx + 1,
										]
										snake_case_      =				torch.as_tensor([expected_single_positions, expected_single_positions]							)
										snake_case_      =				embeddings.create_position_ids_from_inputs_embeds(__A							)
										self.assertEqual(position_ids.shape    ,							expected_positions.shape							)
										self.assertTrue(torch.all(torch.eq(__A    ,							__A							)							)							)




			@unittest.skip("Esm does not support embedding resizing"							)
			def        lowerCAmelCase__    ( self							)    ->    Dict:

										'''simple docstring'''




										pass




			@unittest.skip("Esm does not support embedding resizing"							)
			def        lowerCAmelCase__    ( self							)    ->    Dict:

										'''simple docstring'''




										pass




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

										'''simple docstring'''




										pass




@require_torch
class 			_snake_case		(     A__							):




			@slow
			def        lowerCAmelCase__    ( self							)    ->    List[str]:

										'''simple docstring'''




										with torch.no_grad():
																	snake_case_      =				EsmForMaskedLM.from_pretrained("facebook/esm2_t6_8M_UR50D"							)
																	model.eval()
																	snake_case_      =				torch.tensor([[0, 1, 2, 3, 4, 5]]							)
																	snake_case_      =				model(__A							)[0]

																	snake_case_      =				33

																	snake_case_      =				torch.Size((1, 6, vocab_size)							)
																	self.assertEqual(output.shape    ,							__A							)

																	snake_case_      =				torch.tensor(
																	    [[[8.9_2_1_5, -10.5_898, -6.4_6_7_1], [-6.3_9_6_7, -13.9_114, -1.1_2_1_2], [-7.7_8_1_2, -13.9_516, -3.7_4_0_6]]]							)
																	self.assertTrue(torch.allclose(output[:, :3, :3]    ,							__A    ,							atol=1e-4							)							)




			@slow
			def        lowerCAmelCase__    ( self							)    ->    List[str]:

										'''simple docstring'''




										with torch.no_grad():
																	snake_case_      =				EsmModel.from_pretrained("facebook/esm2_t6_8M_UR50D"							)
																	model.eval()

																	snake_case_      =				torch.tensor([[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_      =				torch.tensor(
																	    [[[0.1_4_4_4, 0.5_4_1_3, 0.3_2_4_8], [0.3_0_3_4, 0.0_0_5_3, 0.3_1_0_8], [0.3_2_2_8, -0.2_4_9_9, 0.3_4_1_5]]]							)
																	self.assertTrue(torch.allclose(output[:, :3, :3]    ,							__A    ,							atol=1e-4							)							)











'''simple docstring'''





import argparse
from collections import defaultdict

import yaml


_SCREAMING_SNAKE_CASE							:       Optional[Any]						=			"docs/source/en/_toctree.yml"


def 					UpperCamelCase_(     snake_case							:	Optional[Any]					):

							'''simple docstring'''






							snake_case_      =				defaultdict(snake_case					)
							for doc in model_doc:
														counts[doc["local"]] += 1
							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 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 					UpperCamelCase_(     snake_case							:	Optional[int]=False					):

							'''simple docstring'''






							with open(snake_case							,    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[model_idx]["title"] != "Models":
														model_idx += 1

							snake_case_      =				api_doc[model_idx]["sections"]

							snake_case_      =				[(idx, section) for idx, section in enumerate(snake_case					) if "sections" in section]
							snake_case_      =				False
							for idx, modality_doc in modalities_docs:
														snake_case_      =				modality_doc["sections"]
														snake_case_      =				clean_model_doc_toc(snake_case					)

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

							if diff:
														if overwrite:
																					snake_case_      =				model_doc
																					snake_case_      =				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__":
		_SCREAMING_SNAKE_CASE							:       List[str]						=			argparse.ArgumentParser()
		parser.add_argument("--fix_and_overwrite", action="store_true", help="Whether to fix inconsistencies.")
		_SCREAMING_SNAKE_CASE							:       Any						=			parser.parse_args()

		check_model_doc(args.fix_and_overwrite)




