




"""simple docstring"""







from __future__ import annotations

import unittest

from transformers import RoFormerConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
						import tensorflow as tf

						from transformers import (
						    TFRoFormerForCausalLM,
						    TFRoFormerForMaskedLM,
						    TFRoFormerForMultipleChoice,
						    TFRoFormerForQuestionAnswering,
						    TFRoFormerForSequenceClassification,
						    TFRoFormerForTokenClassification,
						    TFRoFormerModel,
						)
						from transformers.models.roformer.modeling_tf_roformer import (
						    TFRoFormerSelfAttention,
						    TFRoFormerSinusoidalPositionalEmbedding,
						)




class       lowerCAmelCase__       :


				"""simple docstring"""





				def __init__(     self					:	Any    ,   lowercase__					:	Optional[int]    ,   lowercase__					:	Union[str, Any]=1_3    ,   lowercase__					:	List[str]=7    ,   lowercase__					:	Any=True    ,   lowercase__					:	List[str]=True    ,   lowercase__					:	List[Any]=True    ,   lowercase__					:	Optional[Any]=True    ,   lowercase__					:	Dict=9_9    ,   lowercase__					:	int=3_2    ,   lowercase__					:	Tuple=2    ,   lowercase__					:	Optional[Any]=4    ,   lowercase__					:	List[str]=3_7    ,   lowercase__					:	int="gelu"    ,   lowercase__					:	Optional[Any]=0.1    ,   lowercase__					:	int=0.1    ,   lowercase__					:	int=5_1_2    ,   lowercase__					:	str=1_6    ,   lowercase__					:	Union[str, Any]=2    ,   lowercase__					:	Dict=0.0_2    ,   lowercase__					:	Dict=3    ,   lowercase__					:	Tuple=4    ,   lowercase__					:	Dict=None    ,   ):
						__lowercase				:		List[str]												=				parent
						__lowercase				:		Any												=				1_3
						__lowercase				:		Dict												=				7
						__lowercase				:		Optional[int]												=				True
						__lowercase				:		List[Any]												=				True
						__lowercase				:		Optional[int]												=				True
						__lowercase				:		List[Any]												=				True
						__lowercase				:		Tuple												=				9_9
						__lowercase				:		Tuple												=				3_2
						__lowercase				:		Optional[int]												=				2
						__lowercase				:		Union[str, Any]												=				4
						__lowercase				:		str												=				3_7
						__lowercase				:		Any												=				"gelu"
						__lowercase				:		List[Any]												=				0.1
						__lowercase				:		Optional[int]												=				0.1
						__lowercase				:		List[Any]												=				5_1_2
						__lowercase				:		Dict												=				1_6
						__lowercase				:		Union[str, Any]												=				2
						__lowercase				:		Any												=				0.0_2
						__lowercase				:		List[Any]												=				3
						__lowercase				:		int												=				4
						__lowercase				:		str												=				None



				def       snake_case    (     self					:	str						):
						__lowercase				:		Union[str, Any]												=				ids_tensor([self.batch_size, self.seq_length]    ,   self.vocab_size						)

						__lowercase				:		Any												=				None
						if self.use_input_mask:
								__lowercase				:		Tuple												=				random_attention_mask([self.batch_size, self.seq_length]						)

						__lowercase				:		Any												=				None
						if self.use_token_type_ids:
								__lowercase				:		List[str]												=				ids_tensor([self.batch_size, self.seq_length]    ,   self.type_vocab_size						)

						__lowercase				:		List[str]												=				None
						__lowercase				:		int												=				None
						__lowercase				:		Dict												=				None
						if self.use_labels:
								__lowercase				:		Optional[Any]												=				ids_tensor([self.batch_size]    ,   self.type_sequence_label_size						)
								__lowercase				:		Tuple												=				ids_tensor([self.batch_size, self.seq_length]    ,   self.num_labels						)
								__lowercase				:		int												=				ids_tensor([self.batch_size]    ,   self.num_choices						)

						__lowercase				:		Union[str, Any]												=				RoFormerConfig(
						    vocab_size=self.vocab_size    ,   hidden_size=self.hidden_size    ,   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    ,   return_dict=lowercase__    ,   )

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



				def       snake_case    (     self					:	Dict    ,   lowercase__					:	List[Any]    ,   lowercase__					:	Union[str, Any]    ,   lowercase__					:	Optional[Any]    ,   lowercase__					:	Optional[int]    ,   lowercase__					:	Dict    ,   lowercase__					:	List[Any]    ,   lowercase__					:	List[str]						):
						__lowercase				:		List[Any]												=				TFRoFormerModel(config=lowercase__						)
						__lowercase				:		List[str]												=				{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}

						__lowercase				:		Tuple												=				[input_ids, input_mask]
						__lowercase				:		List[Any]												=				model(lowercase__						)

						__lowercase				:		Tuple												=				model(lowercase__						)

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



				def       snake_case    (     self					:	Any    ,   lowercase__					:	Union[str, Any]    ,   lowercase__					:	Dict    ,   lowercase__					:	Optional[Any]    ,   lowercase__					:	str    ,   lowercase__					:	Optional[int]    ,   lowercase__					:	Dict    ,   lowercase__					:	List[str]						):
						__lowercase				:		str												=				True
						__lowercase				:		List[Any]												=				TFRoFormerForCausalLM(config=lowercase__						)
						__lowercase				:		Any												=				{
						    "input_ids": input_ids,
						    "attention_mask": input_mask,
						    "token_type_ids": token_type_ids,
						}
						__lowercase				:		List[str]												=				model(lowercase__						)["logits"]
						self.parent.assertListEqual(
						    list(prediction_scores.numpy().shape						)    ,   [self.batch_size, self.seq_length, self.vocab_size]						)



				def       snake_case    (     self					:	str    ,   lowercase__					:	Tuple    ,   lowercase__					:	Dict    ,   lowercase__					:	str    ,   lowercase__					:	Union[str, Any]    ,   lowercase__					:	Optional[int]    ,   lowercase__					:	Union[str, Any]    ,   lowercase__					:	Dict						):
						__lowercase				:		Tuple												=				TFRoFormerForMaskedLM(config=lowercase__						)
						__lowercase				:		List[Any]												=				{
						    "input_ids": input_ids,
						    "attention_mask": input_mask,
						    "token_type_ids": token_type_ids,
						}
						__lowercase				:		List[Any]												=				model(lowercase__						)
						self.parent.assertEqual(result.logits.shape    ,   (self.batch_size, self.seq_length, self.vocab_size)						)



				def       snake_case    (     self					:	Dict    ,   lowercase__					:	List[Any]    ,   lowercase__					:	List[str]    ,   lowercase__					:	int    ,   lowercase__					:	Dict    ,   lowercase__					:	int    ,   lowercase__					:	int    ,   lowercase__					:	Union[str, Any]						):
						__lowercase				:		Any												=				self.num_labels
						__lowercase				:		Dict												=				TFRoFormerForSequenceClassification(config=lowercase__						)
						__lowercase				:		Any												=				{
						    "input_ids": input_ids,
						    "attention_mask": input_mask,
						    "token_type_ids": token_type_ids,
						}

						__lowercase				:		Any												=				model(lowercase__						)
						self.parent.assertEqual(result.logits.shape    ,   (self.batch_size, self.num_labels)						)



				def       snake_case    (     self					:	Tuple    ,   lowercase__					:	List[Any]    ,   lowercase__					:	Optional[Any]    ,   lowercase__					:	Optional[int]    ,   lowercase__					:	List[str]    ,   lowercase__					:	Optional[int]    ,   lowercase__					:	Optional[int]    ,   lowercase__					:	Optional[int]						):
						__lowercase				:		Union[str, Any]												=				self.num_choices
						__lowercase				:		int												=				TFRoFormerForMultipleChoice(config=lowercase__						)
						__lowercase				:		Any												=				tf.tile(tf.expand_dims(lowercase__    ,   1						)    ,   (1, self.num_choices, 1)						)
						__lowercase				:		Dict												=				tf.tile(tf.expand_dims(lowercase__    ,   1						)    ,   (1, self.num_choices, 1)						)
						__lowercase				:		Any												=				tf.tile(tf.expand_dims(lowercase__    ,   1						)    ,   (1, self.num_choices, 1)						)
						__lowercase				:		Any												=				{
						    "input_ids": multiple_choice_inputs_ids,
						    "attention_mask": multiple_choice_input_mask,
						    "token_type_ids": multiple_choice_token_type_ids,
						}
						__lowercase				:		List[Any]												=				model(lowercase__						)
						self.parent.assertEqual(result.logits.shape    ,   (self.batch_size, self.num_choices)						)



				def       snake_case    (     self					:	Union[str, Any]    ,   lowercase__					:	Dict    ,   lowercase__					:	int    ,   lowercase__					:	str    ,   lowercase__					:	Optional[Any]    ,   lowercase__					:	Optional[Any]    ,   lowercase__					:	Union[str, Any]    ,   lowercase__					:	str						):
						__lowercase				:		List[str]												=				self.num_labels
						__lowercase				:		Optional[int]												=				TFRoFormerForTokenClassification(config=lowercase__						)
						__lowercase				:		Union[str, Any]												=				{
						    "input_ids": input_ids,
						    "attention_mask": input_mask,
						    "token_type_ids": token_type_ids,
						}
						__lowercase				:		int												=				model(lowercase__						)
						self.parent.assertEqual(result.logits.shape    ,   (self.batch_size, self.seq_length, self.num_labels)						)



				def       snake_case    (     self					:	Dict    ,   lowercase__					:	Tuple    ,   lowercase__					:	List[Any]    ,   lowercase__					:	Tuple    ,   lowercase__					:	Tuple    ,   lowercase__					:	str    ,   lowercase__					:	int    ,   lowercase__					:	Optional[int]						):
						__lowercase				:		List[str]												=				TFRoFormerForQuestionAnswering(config=lowercase__						)
						__lowercase				:		List[str]												=				{
						    "input_ids": input_ids,
						    "attention_mask": input_mask,
						    "token_type_ids": token_type_ids,
						}

						__lowercase				:		Tuple												=				model(lowercase__						)
						self.parent.assertEqual(result.start_logits.shape    ,   (self.batch_size, self.seq_length)						)
						self.parent.assertEqual(result.end_logits.shape    ,   (self.batch_size, self.seq_length)						)



				def       snake_case    (     self					:	Tuple						):
						__lowercase				:		Dict												=				self.prepare_config_and_inputs()
						(
						    (
						    __lowercase
						)						,(
						    __lowercase
						)						,(
						    __lowercase
						)						,(
						    __lowercase
						)						,(
						    __lowercase
						)						,(
						    __lowercase
						)						,(
						    __lowercase
						)						,
						)				:		Union[str, Any]												=				config_and_inputs
						__lowercase				:		Tuple												=				{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
						return config, inputs_dict





@require_tf
class       lowerCAmelCase__       (			lowerCAmelCase_  ,			lowerCAmelCase_  ,			unittest.TestCase    ):


				"""simple docstring"""





				__UpperCAmelCase :     List[str]       =      (
				    (
				        TFRoFormerModel,
				        TFRoFormerForCausalLM,
				        TFRoFormerForMaskedLM,
				        TFRoFormerForQuestionAnswering,
				        TFRoFormerForSequenceClassification,
				        TFRoFormerForTokenClassification,
				        TFRoFormerForMultipleChoice,
				    )
				    if is_tf_available()
				    else ()
				)
				__UpperCAmelCase :     List[str]       =      (
				    {
				        "feature-extraction": TFRoFormerModel,
				        "fill-mask": TFRoFormerForMaskedLM,
				        "question-answering": TFRoFormerForQuestionAnswering,
				        "text-classification": TFRoFormerForSequenceClassification,
				        "text-generation": TFRoFormerForCausalLM,
				        "token-classification": TFRoFormerForTokenClassification,
				        "zero-shot": TFRoFormerForSequenceClassification,
				    }
				    if is_tf_available()
				    else {}
				)

				__UpperCAmelCase :     int       =      False
				__UpperCAmelCase :     Any       =      False



				def       snake_case    (     self					:	Union[str, Any]    ,   lowercase__					:	Tuple    ,   lowercase__					:	int    ,   lowercase__					:	List[Any]    ,   lowercase__					:	Optional[int]    ,   lowercase__					:	Dict						):
						if pipeline_test_casse_name == "TextGenerationPipelineTests":
								return True

						return False



				def       snake_case    (     self					:	Optional[Any]						):
						__lowercase				:		Tuple												=				TFRoFormerModelTester(self						)
						__lowercase				:		Any												=				ConfigTester(self    ,   config_class=lowercase__    ,   hidden_size=3_7						)



				def       snake_case    (     self					:	Optional[Any]						):
						self.config_tester.run_common_tests()



				def       snake_case    (     self					:	int						):
						__lowercase				:		List[Any]												=				self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*lowercase__						)



				def       snake_case    (     self					:	Optional[Any]						):
						__lowercase				:		Union[str, Any]												=				self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_masked_lm(*lowercase__						)



				def       snake_case    (     self					:	Any						):
						__lowercase				:		Tuple												=				self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_lm_head(*lowercase__						)



				def       snake_case    (     self					:	Any						):
						__lowercase				:		List[Any]												=				self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_multiple_choice(*lowercase__						)



				def       snake_case    (     self					:	Tuple						):
						__lowercase				:		int												=				self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_question_answering(*lowercase__						)



				def       snake_case    (     self					:	List[Any]						):
						__lowercase				:		Any												=				self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_sequence_classification(*lowercase__						)



				def       snake_case    (     self					:	Optional[Any]						):
						__lowercase				:		List[str]												=				self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_for_token_classification(*lowercase__						)



				@slow
				def       snake_case    (     self					:	Tuple						):
						__lowercase				:		Dict												=				TFRoFormerModel.from_pretrained("junnyu/roformer_chinese_base"						)
						self.assertIsNotNone(lowercase__						)





@require_tf
class       lowerCAmelCase__       (			unittest.TestCase    ):


				"""simple docstring"""





				@slow
				def       snake_case    (     self					:	str						):
						__lowercase				:		List[Any]												=				TFRoFormerForMaskedLM.from_pretrained("junnyu/roformer_chinese_base"						)
						__lowercase				:		List[Any]												=				tf.constant([[0, 1, 2, 3, 4, 5]]						)
						__lowercase				:		Dict												=				model(lowercase__						)[0]

						# TODO Replace vocab size
						__lowercase				:		Any												=				5_0_0_0_0

						__lowercase				:		Optional[int]												=				[1, 6, vocab_size]
						self.assertEqual(output.shape    ,   lowercase__						)

						print(output[:, :3, :3]						)

						# TODO Replace values below with what was printed above.
						__lowercase				:		str												=				tf.constant(
						    [
						        [
						            [-0.1_2_0_5_3_3_4_1, -1.0_2_6_4_9_0_1, 0.2_9_2_2_1_9_4_6],
						            [-1.5_1_3_3_7_8_3, 0.1_9_7_4_3_3, 0.1_5_1_9_0_6_0_7],
						            [-5.0_1_3_5_4_0_3, -3.9_0_0_2_5_6, -0.8_4_0_3_8_7_6_4],
						        ]
						    ]						)
						tf.debugging.assert_near(output[:, :3, :3]    ,   lowercase__    ,   atol=1e-4						)





@require_tf
class       lowerCAmelCase__       (			unittest.TestCase    ):


				"""simple docstring"""





				__UpperCAmelCase :     Any       =      1E-4



				def       snake_case    (     self					:	Tuple						):
						__lowercase				:		str												=				tf.constant([[4, 1_0]]						)
						__lowercase				:		str												=				TFRoFormerSinusoidalPositionalEmbedding(num_positions=6    ,   embedding_dim=6						)

						__lowercase				:		List[Any]												=				emba(input_ids.shape						)
						__lowercase				:		List[Any]												=				tf.constant(
						    [[0.0_0_0_0, 0.0_0_0_0, 0.0_0_0_0, 1.0_0_0_0, 1.0_0_0_0, 1.0_0_0_0], [0.8_4_1_5, 0.0_4_6_4, 0.0_0_2_2, 0.5_4_0_3, 0.9_9_8_9, 1.0_0_0_0]]						)

						tf.debugging.assert_near(lowercase__    ,   lowercase__    ,   atol=self.tolerance						)



				def       snake_case    (     self					:	List[str]						):
						__lowercase				:		Optional[int]												=				tf.constant(
						    [
						        [0.0_0_0_0, 0.0_0_0_0, 0.0_0_0_0, 0.0_0_0_0, 0.0_0_0_0],
						        [0.8_4_1_5, 0.8_2_1_9, 0.8_0_2_0, 0.7_8_1_9, 0.7_6_1_7],
						        [0.9_0_9_3, 0.9_3_6_4, 0.9_5_8_1, 0.9_7_4_9, 0.9_8_7_0],
						    ]						)
						__lowercase				:		List[str]												=				TFRoFormerSinusoidalPositionalEmbedding(num_positions=5_1_2    ,   embedding_dim=5_1_2						)
						emba([2, 1_6, 5_1_2]						)
						__lowercase				:		Optional[int]												=				emba.weight[:3, :5]

						tf.debugging.assert_near(lowercase__    ,   lowercase__    ,   atol=self.tolerance						)





@require_tf
class       lowerCAmelCase__       (			unittest.TestCase    ):


				"""simple docstring"""





				__UpperCAmelCase :     Optional[int]       =      1E-4



				def       snake_case    (     self					:	Tuple						):
						# 2,12,16,64
						__lowercase				:		List[str]												=				tf.reshape(tf.range(2 * 1_2 * 1_6 * 6_4    ,   dtype=tf.floataa						)    ,   shape=(2, 1_2, 1_6, 6_4)						) / 1_0_0

						__lowercase				:		List[Any]												=				-tf.reshape(tf.range(2 * 1_2 * 1_6 * 6_4    ,   dtype=tf.floataa						)    ,   shape=(2, 1_2, 1_6, 6_4)						) / 1_0_0

						__lowercase				:		Optional[Any]												=				TFRoFormerSinusoidalPositionalEmbedding(num_positions=3_2    ,   embedding_dim=6_4						)
						__lowercase				:		List[str]												=				embed_positions([2, 1_6, 7_6_8]						)[None, None, :, :]

						__lowercase						,__lowercase				:		Any												=				TFRoFormerSelfAttention.apply_rotary_position_embeddings(
						    lowercase__    ,   lowercase__    ,   lowercase__						)

						__lowercase				:		Optional[int]												=				tf.constant(
						    [
						        [0.0_0_0_0, 0.0_1_0_0, 0.0_2_0_0, 0.0_3_0_0, 0.0_4_0_0, 0.0_5_0_0, 0.0_6_0_0, 0.0_7_0_0],
						        [-0.2_0_1_2, 0.8_8_9_7, 0.0_2_6_3, 0.9_4_0_1, 0.2_0_7_4, 0.9_4_6_3, 0.3_4_8_1, 0.9_3_4_3],
						        [-1.7_0_5_7, 0.6_2_7_1, -1.2_1_4_5, 1.3_8_9_7, -0.6_3_0_3, 1.7_6_4_7, -0.1_1_7_3, 1.8_9_8_5],
						        [-2.1_7_3_1, -1.6_3_9_7, -2.7_3_5_8, 0.2_8_5_4, -2.1_8_4_0, 1.7_1_8_3, -1.3_0_1_8, 2.4_8_7_1],
						        [0.2_7_1_7, -3.6_1_7_3, -2.9_2_0_6, -2.1_9_8_8, -3.6_6_3_8, 0.3_8_5_8, -2.9_1_5_5, 2.2_9_8_0],
						        [3.9_8_5_9, -2.1_5_8_0, -0.7_9_8_4, -4.4_9_0_4, -4.1_1_8_1, -2.0_2_5_2, -4.4_7_8_2, 1.1_2_5_3],
						    ]						)
						__lowercase				:		str												=				tf.constant(
						    [
						        [0.0_0_0_0, -0.0_1_0_0, -0.0_2_0_0, -0.0_3_0_0, -0.0_4_0_0, -0.0_5_0_0, -0.0_6_0_0, -0.0_7_0_0],
						        [0.2_0_1_2, -0.8_8_9_7, -0.0_2_6_3, -0.9_4_0_1, -0.2_0_7_4, -0.9_4_6_3, -0.3_4_8_1, -0.9_3_4_3],
						        [1.7_0_5_7, -0.6_2_7_1, 1.2_1_4_5, -1.3_8_9_7, 0.6_3_0_3, -1.7_6_4_7, 0.1_1_7_3, -1.8_9_8_5],
						        [2.1_7_3_1, 1.6_3_9_7, 2.7_3_5_8, -0.2_8_5_4, 2.1_8_4_0, -1.7_1_8_3, 1.3_0_1_8, -2.4_8_7_1],
						        [-0.2_7_1_7, 3.6_1_7_3, 2.9_2_0_6, 2.1_9_8_8, 3.6_6_3_8, -0.3_8_5_8, 2.9_1_5_5, -2.2_9_8_0],
						        [-3.9_8_5_9, 2.1_5_8_0, 0.7_9_8_4, 4.4_9_0_4, 4.1_1_8_1, 2.0_2_5_2, 4.4_7_8_2, -1.1_2_5_3],
						    ]						)

						tf.debugging.assert_near(query_layer[0, 0, :6, :8]    ,   lowercase__    ,   atol=self.tolerance						)
						tf.debugging.assert_near(key_layer[0, 0, :6, :8]    ,   lowercase__    ,   atol=self.tolerance						)






"""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.models.esm.modeling_esmfold import EsmForProteinFolding




class       lowerCAmelCase__       :


				"""simple docstring"""





				def __init__(     self					:	Tuple    ,   lowercase__					:	Tuple    ,   lowercase__					:	Tuple=1_3    ,   lowercase__					:	Optional[int]=7    ,   lowercase__					:	List[str]=False    ,   lowercase__					:	Tuple=True    ,   lowercase__					:	int=False    ,   lowercase__					:	List[str]=False    ,   lowercase__					:	Optional[Any]=1_9    ,   lowercase__					:	int=3_2    ,   lowercase__					:	List[Any]=5    ,   lowercase__					:	Optional[int]=4    ,   lowercase__					:	Any=3_7    ,   lowercase__					:	Tuple="gelu"    ,   lowercase__					:	int=0.1    ,   lowercase__					:	Tuple=0.1    ,   lowercase__					:	List[Any]=5_1_2    ,   lowercase__					:	List[Any]=1_6    ,   lowercase__					:	Union[str, Any]=2    ,   lowercase__					:	List[str]=0.0_2    ,   lowercase__					:	List[Any]=3    ,   lowercase__					:	Any=4    ,   lowercase__					:	Optional[Any]=None    ,   ):
						__lowercase				:		int												=				parent
						__lowercase				:		Tuple												=				batch_size
						__lowercase				:		Optional[int]												=				seq_length
						__lowercase				:		str												=				is_training
						__lowercase				:		List[Any]												=				use_input_mask
						__lowercase				:		Any												=				use_token_type_ids
						__lowercase				:		str												=				use_labels
						__lowercase				:		Dict												=				vocab_size
						__lowercase				:		Optional[int]												=				hidden_size
						__lowercase				:		List[Any]												=				num_hidden_layers
						__lowercase				:		int												=				num_attention_heads
						__lowercase				:		Union[str, Any]												=				intermediate_size
						__lowercase				:		Dict												=				hidden_act
						__lowercase				:		Union[str, Any]												=				hidden_dropout_prob
						__lowercase				:		Optional[int]												=				attention_probs_dropout_prob
						__lowercase				:		Optional[Any]												=				max_position_embeddings
						__lowercase				:		List[str]												=				type_vocab_size
						__lowercase				:		str												=				type_sequence_label_size
						__lowercase				:		List[str]												=				initializer_range
						__lowercase				:		Optional[Any]												=				num_labels
						__lowercase				:		Tuple												=				num_choices
						__lowercase				:		Optional[Any]												=				scope



				def       snake_case    (     self					:	Optional[int]						):
						__lowercase				:		Tuple												=				ids_tensor([self.batch_size, self.seq_length]    ,   self.vocab_size						)

						__lowercase				:		Union[str, Any]												=				None
						if self.use_input_mask:
								__lowercase				:		List[str]												=				random_attention_mask([self.batch_size, self.seq_length]						)

						__lowercase				:		Dict												=				None
						__lowercase				:		Any												=				None
						__lowercase				:		Optional[Any]												=				None
						if self.use_labels:
								__lowercase				:		Tuple												=				ids_tensor([self.batch_size]    ,   self.type_sequence_label_size						)
								__lowercase				:		List[str]												=				ids_tensor([self.batch_size, self.seq_length]    ,   self.num_labels						)
								__lowercase				:		Optional[Any]												=				ids_tensor([self.batch_size]    ,   self.num_choices						)

						__lowercase				:		Optional[Any]												=				self.get_config()

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



				def       snake_case    (     self					:	Any						):
						__lowercase				:		Dict												=				EsmConfig(
						    vocab_size=3_3    ,   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=lowercase__    ,   esmfold_config={"trunk": {"num_blocks": 2}, "fp16_esm": False}    ,   )
						return config



				def       snake_case    (     self					:	str    ,   lowercase__					:	List[Any]    ,   lowercase__					:	Optional[int]    ,   lowercase__					:	List[Any]    ,   lowercase__					:	int    ,   lowercase__					:	List[Any]    ,   lowercase__					:	str						):
						__lowercase				:		List[Any]												=				EsmForProteinFolding(config=lowercase__						).float()
						model.to(lowercase__						)
						model.eval()
						__lowercase				:		List[str]												=				model(lowercase__    ,   attention_mask=lowercase__						)
						__lowercase				:		Any												=				model(lowercase__						)
						__lowercase				:		Any												=				model(lowercase__						)

						self.parent.assertEqual(result.positions.shape    ,   (8, self.batch_size, self.seq_length, 1_4, 3)						)
						self.parent.assertEqual(result.angles.shape    ,   (8, self.batch_size, self.seq_length, 7, 2)						)



				def       snake_case    (     self					:	Tuple						):
						__lowercase				:		List[str]												=				self.prepare_config_and_inputs()
						(
						    (
						    __lowercase
						)						,(
						    __lowercase
						)						,(
						    __lowercase
						)						,(
						    __lowercase
						)						,(
						    __lowercase
						)						,(
						    __lowercase
						)						,
						)				:		Any												=				config_and_inputs
						__lowercase				:		int												=				{"input_ids": input_ids, "attention_mask": input_mask}
						return config, inputs_dict





@require_torch
class       lowerCAmelCase__       (			lowerCAmelCase_  ,			lowerCAmelCase_  ,			unittest.TestCase    ):


				"""simple docstring"""





				__UpperCAmelCase :     Union[str, Any]       =      False

				__UpperCAmelCase :     List[Any]       =      (EsmForProteinFolding,) if is_torch_available() else ()
				__UpperCAmelCase :     int       =      ()
				__UpperCAmelCase :     Tuple       =      {} if is_torch_available() else {}
				__UpperCAmelCase :     Optional[Any]       =      False



				def       snake_case    (     self					:	Tuple						):
						__lowercase				:		Tuple												=				EsmFoldModelTester(self						)
						__lowercase				:		Dict												=				ConfigTester(self    ,   config_class=lowercase__    ,   hidden_size=3_7						)



				def       snake_case    (     self					:	Dict						):
						self.config_tester.run_common_tests()



				def       snake_case    (     self					:	Optional[int]						):
						__lowercase				:		Tuple												=				self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_model(*lowercase__						)



				@unittest.skip("Does not support attention outputs"						)
				def       snake_case    (     self					:	Union[str, Any]						):
						pass



				@unittest.skip
				def       snake_case    (     self					:	Tuple						):
						pass



				@unittest.skip("Esm does not support embedding resizing"						)
				def       snake_case    (     self					:	Optional[int]						):
						pass



				@unittest.skip("Esm does not support embedding resizing"						)
				def       snake_case    (     self					:	List[str]						):
						pass



				@unittest.skip("ESMFold does not support passing input embeds!"						)
				def       snake_case    (     self					:	int						):
						pass



				@unittest.skip("ESMFold does not support head pruning."						)
				def       snake_case    (     self					:	List[Any]						):
						pass



				@unittest.skip("ESMFold does not support head pruning."						)
				def       snake_case    (     self					:	Any						):
						pass



				@unittest.skip("ESMFold does not support head pruning."						)
				def       snake_case    (     self					:	Optional[Any]						):
						pass



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



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



				@unittest.skip("ESMFold does not output hidden states in the normal way."						)
				def       snake_case    (     self					:	Optional[Any]						):
						pass



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



				@unittest.skip("ESMFold only has one output format."						)
				def       snake_case    (     self					:	Tuple						):
						pass



				@unittest.skip("This test doesn't work for ESMFold and doesn't test core functionality"						)
				def       snake_case    (     self					:	Any						):
						pass



				@unittest.skip("ESMFold does not support input chunking."						)
				def       snake_case    (     self					:	str						):
						pass



				@unittest.skip("ESMFold doesn't respect you and it certainly doesn't respect your initialization arguments."						)
				def       snake_case    (     self					:	Any						):
						pass



				@unittest.skip("ESMFold doesn't support torchscript compilation."						)
				def       snake_case    (     self					:	Any						):
						pass



				@unittest.skip("ESMFold doesn't support torchscript compilation."						)
				def       snake_case    (     self					:	Tuple						):
						pass



				@unittest.skip("ESMFold doesn't support torchscript compilation."						)
				def       snake_case    (     self					:	List[str]						):
						pass



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



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





@require_torch
class       lowerCAmelCase__       (			lowerCAmelCase_    ):


				"""simple docstring"""





				@slow
				def       snake_case    (     self					:	Union[str, Any]						):
						__lowercase				:		Tuple												=				EsmForProteinFolding.from_pretrained("facebook/esmfold_v1"						).float()
						model.eval()
						__lowercase				:		Optional[int]												=				torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]						)
						__lowercase				:		str												=				model(lowercase__						)["positions"]
						__lowercase				:		Union[str, Any]												=				torch.tensor([2.5_8_2_8, 0.7_9_9_3, -1_0.9_3_3_4]    ,   dtype=torch.floataa						)
						self.assertTrue(torch.allclose(position_outputs[0, 0, 0, 0]    ,   lowercase__    ,   atol=1e-4						)						)
