


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




			def __init__(						self     :   Union[str, Any]		,     lowerCamelCase__     :   Optional[Any]		,     lowerCamelCase__     :   Tuple=13		,     lowerCamelCase__     :   int=7		,     lowerCamelCase__     :   List[str]=False		,     lowerCamelCase__     :   Optional[Any]=True		,     lowerCamelCase__     :   List[str]=False		,     lowerCamelCase__     :   str=False		,     lowerCamelCase__     :   Union[str, Any]=19		,     lowerCamelCase__     :   Union[str, Any]=32		,     lowerCamelCase__     :   Optional[Any]=5		,     lowerCamelCase__     :   List[str]=4		,     lowerCamelCase__     :   Dict=37		,     lowerCamelCase__     :   Tuple="gelu"		,     lowerCamelCase__     :   Any=0.1		,     lowerCamelCase__     :   int=0.1		,     lowerCamelCase__     :   Any=5_12		,     lowerCamelCase__     :   Dict=16		,     lowerCamelCase__     :   Optional[int]=2		,     lowerCamelCase__     :   Optional[Any]=0.0_2		,     lowerCamelCase__     :   Union[str, Any]=3		,     lowerCamelCase__     :   Optional[int]=4		,     lowerCamelCase__     :   Optional[int]=None		,     )      ->int:



					'''simple docstring'''







					_UpperCAmelCase						:				List[str]	    =					parent
					_UpperCAmelCase						:				int	    =					batch_size
					_UpperCAmelCase						:				Any	    =					seq_length
					_UpperCAmelCase						:				Dict	    =					is_training
					_UpperCAmelCase						:				Any	    =					use_input_mask
					_UpperCAmelCase						:				Optional[Any]	    =					use_token_type_ids
					_UpperCAmelCase						:				Optional[Any]	    =					use_labels
					_UpperCAmelCase						:				List[Any]	    =					vocab_size
					_UpperCAmelCase						:				List[str]	    =					hidden_size
					_UpperCAmelCase						:				List[str]	    =					num_hidden_layers
					_UpperCAmelCase						:				Optional[Any]	    =					num_attention_heads
					_UpperCAmelCase						:				Dict	    =					intermediate_size
					_UpperCAmelCase						:				Dict	    =					hidden_act
					_UpperCAmelCase						:				List[Any]	    =					hidden_dropout_prob
					_UpperCAmelCase						:				Dict	    =					attention_probs_dropout_prob
					_UpperCAmelCase						:				Optional[Any]	    =					max_position_embeddings
					_UpperCAmelCase						:				Union[str, Any]	    =					type_vocab_size
					_UpperCAmelCase						:				Dict	    =					type_sequence_label_size
					_UpperCAmelCase						:				Optional[Any]	    =					initializer_range
					_UpperCAmelCase						:				Any	    =					num_labels
					_UpperCAmelCase						:				Tuple	    =					num_choices
					_UpperCAmelCase						:				List[Any]	    =					scope




			def 							lowerCAmelCase__			(						self     :   Dict       )      ->str:



					'''simple docstring'''







					_UpperCAmelCase						:				Dict	    =					ids_tensor([self.batch_size, self.seq_length]		,     self.vocab_size       )

					_UpperCAmelCase						:				Optional[Any]	    =					None
					if self.use_input_mask:
							_UpperCAmelCase						:				Optional[int]	    =					random_attention_mask([self.batch_size, self.seq_length]       )

					_UpperCAmelCase						:				str	    =					None
					_UpperCAmelCase						:				str	    =					None
					_UpperCAmelCase						:				Union[str, Any]	    =					None
					if self.use_labels:
							_UpperCAmelCase						:				Any	    =					ids_tensor([self.batch_size]		,     self.type_sequence_label_size       )
							_UpperCAmelCase						:				Any	    =					ids_tensor([self.batch_size, self.seq_length]		,     self.num_labels       )
							_UpperCAmelCase						:				List[Any]	    =					ids_tensor([self.batch_size]		,     self.num_choices       )

					_UpperCAmelCase						:				int	    =					self.get_config()

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




			def 							lowerCAmelCase__			(						self     :   str       )      ->Tuple:



					'''simple docstring'''







					_UpperCAmelCase						:				int	    =					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     :   str		,     lowerCamelCase__     :   List[str]		,     lowerCamelCase__     :   List[str]		,     lowerCamelCase__     :   Union[str, Any]		,     lowerCamelCase__     :   Any		,     lowerCamelCase__     :   Optional[Any]		,     lowerCamelCase__     :   Dict       )      ->List[str]:



					'''simple docstring'''







					_UpperCAmelCase						:				Dict	    =					EsmForProteinFolding(config=lowerCamelCase__       ).float()
					model.to(lowerCamelCase__       )
					model.eval()
					_UpperCAmelCase						:				List[Any]	    =					model(lowerCamelCase__		,     attention_mask=lowerCamelCase__       )
					_UpperCAmelCase						:				int	    =					model(lowerCamelCase__       )
					_UpperCAmelCase						:				Optional[Any]	    =					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       )      ->Union[str, Any]:



					'''simple docstring'''







					_UpperCAmelCase						:				Dict	    =					self.prepare_config_and_inputs()
					(
					    (
					    _UpperCAmelCase
					)    ,			(
					    _UpperCAmelCase
					)    ,			(
					    _UpperCAmelCase
					)    ,			(
					    _UpperCAmelCase
					)    ,			(
					    _UpperCAmelCase
					)    ,			(
					    _UpperCAmelCase
					)    ,			
					)						:				Optional[int]	    =					config_and_inputs
					_UpperCAmelCase						:				Dict	    =					{"input_ids": input_ids, "attention_mask": input_mask}
					return config, inputs_dict






@require_torch
class 		lowerCAmelCase__	(	UpperCAmelCase__			,				UpperCAmelCase__			,				unittest.TestCase		):
			lowerCAmelCase       :		Any				      =   False

			lowerCAmelCase       :		Dict				      =   (EsmForProteinFolding,) if is_torch_available() else ()
			lowerCAmelCase       :		int				      =   ()
			lowerCAmelCase       :		Any				      =   {} if is_torch_available() else {}
			lowerCAmelCase       :		Optional[int]				      =   False




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



					'''simple docstring'''







					_UpperCAmelCase						:				Any	    =					EsmFoldModelTester(self       )
					_UpperCAmelCase						:				Optional[int]	    =					ConfigTester(self		,     config_class=lowerCamelCase__		,     hidden_size=37       )




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



					'''simple docstring'''







					self.config_tester.run_common_tests()




			def 							lowerCAmelCase__			(						self     :   str       )      ->int:



					'''simple docstring'''







					_UpperCAmelCase						:				List[str]	    =					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     :   Dict       )      ->int:



					'''simple docstring'''







					pass




			@unittest.skip
			def 							lowerCAmelCase__			(						self     :   Tuple       )      ->Optional[Any]:



					'''simple docstring'''







					pass




			@unittest.skip("Esm does not support embedding resizing"       )
			def 							lowerCAmelCase__			(						self     :   List[Any]       )      ->List[str]:



					'''simple docstring'''







					pass




			@unittest.skip("Esm does not support embedding resizing"       )
			def 							lowerCAmelCase__			(						self     :   Dict       )      ->Union[str, Any]:



					'''simple docstring'''







					pass




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



					'''simple docstring'''







					pass




			@unittest.skip("ESMFold does not support head pruning."       )
			def 							lowerCAmelCase__			(						self     :   Optional[Any]       )      ->Dict:



					'''simple docstring'''







					pass




			@unittest.skip("ESMFold does not support head pruning."       )
			def 							lowerCAmelCase__			(						self     :   Optional[Any]       )      ->Dict:



					'''simple docstring'''







					pass




			@unittest.skip("ESMFold does not support head pruning."       )
			def 							lowerCAmelCase__			(						self     :   Tuple       )      ->Union[str, Any]:



					'''simple docstring'''







					pass




			@unittest.skip("ESMFold does not support head pruning."       )
			def 							lowerCAmelCase__			(						self     :   Any       )      ->Optional[int]:



					'''simple docstring'''







					pass




			@unittest.skip("ESMFold does not support head pruning."       )
			def 							lowerCAmelCase__			(						self     :   Any       )      ->Union[str, Any]:



					'''simple docstring'''







					pass




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



					'''simple docstring'''







					pass




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



					'''simple docstring'''







					pass




			@unittest.skip("ESMFold only has one output format."       )
			def 							lowerCAmelCase__			(						self     :   Any       )      ->List[str]:



					'''simple docstring'''







					pass




			@unittest.skip("This test doesn't work for ESMFold and doesn't test core functionality"       )
			def 							lowerCAmelCase__			(						self     :   Union[str, Any]       )      ->Tuple:



					'''simple docstring'''







					pass




			@unittest.skip("ESMFold does not support input chunking."       )
			def 							lowerCAmelCase__			(						self     :   Union[str, Any]       )      ->int:



					'''simple docstring'''







					pass




			@unittest.skip("ESMFold doesn't respect you and it certainly doesn't respect your initialization arguments."       )
			def 							lowerCAmelCase__			(						self     :   List[Any]       )      ->Tuple:



					'''simple docstring'''







					pass




			@unittest.skip("ESMFold doesn't support torchscript compilation."       )
			def 							lowerCAmelCase__			(						self     :   Optional[int]       )      ->Union[str, Any]:



					'''simple docstring'''







					pass




			@unittest.skip("ESMFold doesn't support torchscript compilation."       )
			def 							lowerCAmelCase__			(						self     :   List[Any]       )      ->Optional[Any]:



					'''simple docstring'''







					pass




			@unittest.skip("ESMFold doesn't support torchscript compilation."       )
			def 							lowerCAmelCase__			(						self     :   Optional[int]       )      ->Optional[int]:



					'''simple docstring'''







					pass




			@unittest.skip("ESMFold doesn't support data parallel."       )
			def 							lowerCAmelCase__			(						self     :   Optional[int]       )      ->int:



					'''simple docstring'''







					pass






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



					'''simple docstring'''







					pass






@require_torch
class 		lowerCAmelCase__	(	UpperCAmelCase__		):




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



					'''simple docstring'''







					_UpperCAmelCase						:				Any	    =					EsmForProteinFolding.from_pretrained("facebook/esmfold_v1"       ).float()
					model.eval()
					_UpperCAmelCase						:				Dict	    =					torch.tensor([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]       )
					_UpperCAmelCase						:				List[str]	    =					model(lowerCamelCase__       )["positions"]
					_UpperCAmelCase						:				List[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]		,     lowerCamelCase__		,     atol=1E-4       )       )







'''simple docstring'''


import pytest

from datasets import Dataset, DatasetDict, Features, NamedSplit, Value
from datasets.io.text import TextDatasetReader

from ..utils import assert_arrow_memory_doesnt_increase, assert_arrow_memory_increases






def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase      ):
		assert isinstance(__lowerCAmelCase					,       __lowerCAmelCase      )
		assert dataset.num_rows == 4
		assert dataset.num_columns == 1
		assert dataset.column_names == ["text"]
		for feature, expected_dtype in expected_features.items():
				assert dataset.features[feature].dtype == expected_dtype






@pytest.mark.parametrize("keep_in_memory"					,       [False, True]      )
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		_UpperCAmelCase						:				Union[str, Any]	    =					tmp_path / "cache"
		_UpperCAmelCase						:				int	    =					{"text": "string"}
		with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
				_UpperCAmelCase						:				Union[str, Any]	    =					TextDatasetReader(__lowerCAmelCase					,       cache_dir=__lowerCAmelCase					,       keep_in_memory=__lowerCAmelCase      ).read()
		_check_text_dataset(__lowerCAmelCase					,       __lowerCAmelCase      )






@pytest.mark.parametrize(
    "features"					,       [
        None,
        {"text": "string"},
        {"text": "int32"},
        {"text": "float32"},
    ]					,       )
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		_UpperCAmelCase						:				List[Any]	    =					tmp_path / "cache"
		_UpperCAmelCase						:				Any	    =					{"text": "string"}
		_UpperCAmelCase						:				Optional[Any]	    =					features.copy() if features else default_expected_features
		_UpperCAmelCase						:				Union[str, Any]	    =					(
		    Features({feature: Value(__lowerCAmelCase      ) for feature, dtype in features.items()}      ) if features is not None else None
		)
		_UpperCAmelCase						:				Union[str, Any]	    =					TextDatasetReader(__lowerCAmelCase					,       features=__lowerCAmelCase					,       cache_dir=__lowerCAmelCase      ).read()
		_check_text_dataset(__lowerCAmelCase					,       __lowerCAmelCase      )






@pytest.mark.parametrize("split"					,       [None, NamedSplit("train"      ), "train", "test"]      )
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		_UpperCAmelCase						:				Any	    =					tmp_path / "cache"
		_UpperCAmelCase						:				Dict	    =					{"text": "string"}
		_UpperCAmelCase						:				Union[str, Any]	    =					TextDatasetReader(__lowerCAmelCase					,       cache_dir=__lowerCAmelCase					,       split=__lowerCAmelCase      ).read()
		_check_text_dataset(__lowerCAmelCase					,       __lowerCAmelCase      )
		assert dataset.split == split if split else "train"






@pytest.mark.parametrize("path_type"					,       [str, list]      )
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		if issubclass(__lowerCAmelCase					,       __lowerCAmelCase      ):
				_UpperCAmelCase						:				Union[str, Any]	    =					text_path
		elif issubclass(__lowerCAmelCase					,       __lowerCAmelCase      ):
				_UpperCAmelCase						:				int	    =					[text_path]
		_UpperCAmelCase						:				List[Any]	    =					tmp_path / "cache"
		_UpperCAmelCase						:				Union[str, Any]	    =					{"text": "string"}
		_UpperCAmelCase						:				Optional[int]	    =					TextDatasetReader(__lowerCAmelCase					,       cache_dir=__lowerCAmelCase      ).read()
		_check_text_dataset(__lowerCAmelCase					,       __lowerCAmelCase      )






def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase=("train",)      ):
		assert isinstance(__lowerCAmelCase					,       __lowerCAmelCase      )
		for split in splits:
				_UpperCAmelCase						:				List[str]	    =					dataset_dict[split]
				assert dataset.num_rows == 4
				assert dataset.num_columns == 1
				assert dataset.column_names == ["text"]
				for feature, expected_dtype in expected_features.items():
						assert dataset.features[feature].dtype == expected_dtype






@pytest.mark.parametrize("keep_in_memory"					,       [False, True]      )
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		_UpperCAmelCase						:				List[Any]	    =					tmp_path / "cache"
		_UpperCAmelCase						:				Tuple	    =					{"text": "string"}
		with assert_arrow_memory_increases() if keep_in_memory else assert_arrow_memory_doesnt_increase():
				_UpperCAmelCase						:				Any	    =					TextDatasetReader({"train": text_path}					,       cache_dir=__lowerCAmelCase					,       keep_in_memory=__lowerCAmelCase      ).read()
		_check_text_datasetdict(__lowerCAmelCase					,       __lowerCAmelCase      )






@pytest.mark.parametrize(
    "features"					,       [
        None,
        {"text": "string"},
        {"text": "int32"},
        {"text": "float32"},
    ]					,       )
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		_UpperCAmelCase						:				List[Any]	    =					tmp_path / "cache"
		# CSV file loses col_1 string dtype information: default now is "int64" instead of "string"
		_UpperCAmelCase						:				List[Any]	    =					{"text": "string"}
		_UpperCAmelCase						:				List[str]	    =					features.copy() if features else default_expected_features
		_UpperCAmelCase						:				Optional[int]	    =					(
		    Features({feature: Value(__lowerCAmelCase      ) for feature, dtype in features.items()}      ) if features is not None else None
		)
		_UpperCAmelCase						:				Tuple	    =					TextDatasetReader({"train": text_path}					,       features=__lowerCAmelCase					,       cache_dir=__lowerCAmelCase      ).read()
		_check_text_datasetdict(__lowerCAmelCase					,       __lowerCAmelCase      )






@pytest.mark.parametrize("split"					,       [None, NamedSplit("train"      ), "train", "test"]      )
def        __lowerCAmelCase       (__lowerCAmelCase					,       __lowerCAmelCase					,       __lowerCAmelCase      ):
		if split:
				_UpperCAmelCase						:				int	    =					{split: text_path}
		else:
				_UpperCAmelCase						:				Tuple	    =					"train"
				_UpperCAmelCase						:				List[str]	    =					{"train": text_path, "test": text_path}
		_UpperCAmelCase						:				Optional[Any]	    =					tmp_path / "cache"
		_UpperCAmelCase						:				Optional[int]	    =					{"text": "string"}
		_UpperCAmelCase						:				int	    =					TextDatasetReader(__lowerCAmelCase					,       cache_dir=__lowerCAmelCase      ).read()
		_check_text_datasetdict(__lowerCAmelCase					,       __lowerCAmelCase					,       splits=list(path.keys()      )      )
		assert all(dataset[split].split == split for split in path.keys()      )



