






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






			def __init__(			self   :      List[str]	,    __SCREAMING_SNAKE_CASE   :      Dict	,    )      ->      Any:
							a_		:	Tuple									=   parent
							a_		:	List[Any]									=   13
							a_		:	List[Any]									=   7
							a_		:	Optional[Any]									=   True
							a_		:	Any									=   True
							a_		:	Any									=   True
							a_		:	List[Any]									=   99
							a_		:	List[Any]									=   32
							a_		:	int									=   2
							a_		:	Optional[Any]									=   4
							a_		:	int									=   37
							a_		:	Tuple									=   '''gelu'''
							a_		:	List[Any]									=   0.1
							a_		:	Union[str, Any]									=   0.1
							a_		:	str									=   512
							a_		:	Optional[Any]									=   16
							a_		:	int									=   2
							a_		:	Optional[Any]									=   0.02
							a_		:	List[str]									=   3
							a_		:	str									=   4
							a_		:	Tuple									=   None






			def       SCREAMING_SNAKE_CASE   (			self   :      int       )      ->      List[str]:
							a_		:	Optional[int]									=   ids_tensor([self.batch_size, self.seq_length]	,    self.vocab_size       )

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

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

							a_		:	Optional[int]									=   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       SCREAMING_SNAKE_CASE   (			self   :      int       )      ->      Union[str, Any]:
							(
							    (
							    a_
							)					,      (
							    a_
							)					,      (
							    a_
							)					,      (
							    a_
							)					,      (
							    a_
							)					,      (
							    a_
							)					,      
							)		:	Optional[int]									=   self.prepare_config_and_inputs()

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

							a_		:	Optional[int]									=   [input_ids, input_mask]
							a_		:	str									=   model(__SCREAMING_SNAKE_CASE       )

							a_		:	Optional[Any]									=   model(__SCREAMING_SNAKE_CASE       )

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






			def       SCREAMING_SNAKE_CASE   (			self   :      Union[str, Any]	,    __SCREAMING_SNAKE_CASE   :      List[Any]	,    __SCREAMING_SNAKE_CASE   :      int	,    __SCREAMING_SNAKE_CASE   :      Any	,    __SCREAMING_SNAKE_CASE   :      List[Any]	,    __SCREAMING_SNAKE_CASE   :      List[Any]	,    __SCREAMING_SNAKE_CASE   :      Tuple	,    __SCREAMING_SNAKE_CASE   :      Any	,    __SCREAMING_SNAKE_CASE   :      List[str]	,    )      ->      List[str]:
							a_		:	List[Any]									=   True

							a_		:	Dict									=   TFEsmModel(config=__SCREAMING_SNAKE_CASE       )
							a_		:	List[str]									=   {
							    '''input_ids''': input_ids,
							    '''attention_mask''': input_mask,
							    '''encoder_hidden_states''': encoder_hidden_states,
							    '''encoder_attention_mask''': encoder_attention_mask,
							}
							a_		:	int									=   model(__SCREAMING_SNAKE_CASE       )

							a_		:	Dict									=   [input_ids, input_mask]
							a_		:	int									=   model(__SCREAMING_SNAKE_CASE	,    encoder_hidden_states=__SCREAMING_SNAKE_CASE       )

							# Also check the case where encoder outputs are not passed
							a_		:	Dict									=   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       SCREAMING_SNAKE_CASE   (			self   :      Dict	,    __SCREAMING_SNAKE_CASE   :      Optional[int]	,    __SCREAMING_SNAKE_CASE   :      List[str]	,    __SCREAMING_SNAKE_CASE   :      List[Any]	,    __SCREAMING_SNAKE_CASE   :      Union[str, Any]	,    __SCREAMING_SNAKE_CASE   :      int	,    __SCREAMING_SNAKE_CASE   :      Any       )      ->      List[str]:
							a_		:	Optional[int]									=   TFEsmForMaskedLM(config=__SCREAMING_SNAKE_CASE       )
							a_		:	int									=   model([input_ids, input_mask]       )
							self.parent.assertEqual(result.logits.shape	,    (self.batch_size, self.seq_length, self.vocab_size)       )






			def       SCREAMING_SNAKE_CASE   (			self   :      List[str]	,    __SCREAMING_SNAKE_CASE   :      Optional[int]	,    __SCREAMING_SNAKE_CASE   :      Optional[Any]	,    __SCREAMING_SNAKE_CASE   :      Optional[Any]	,    __SCREAMING_SNAKE_CASE   :      int	,    __SCREAMING_SNAKE_CASE   :      Dict	,    __SCREAMING_SNAKE_CASE   :      Any       )      ->      Union[str, Any]:
							a_		:	str									=   self.num_labels
							a_		:	int									=   TFEsmForTokenClassification(config=__SCREAMING_SNAKE_CASE       )
							a_		:	List[Any]									=   {'''input_ids''': input_ids, '''attention_mask''': input_mask}
							a_		:	List[str]									=   model(__SCREAMING_SNAKE_CASE       )
							self.parent.assertEqual(result.logits.shape	,    (self.batch_size, self.seq_length, self.num_labels)       )







			def       SCREAMING_SNAKE_CASE   (			self   :      Optional[Any]       )      ->      str:
							a_		:	Any									=   self.prepare_config_and_inputs()
							(
							    (
							    a_
							)					,      (
							    a_
							)					,      (
							    a_
							)					,      (
							    a_
							)					,      (
							    a_
							)					,      (
							    a_
							)					,      
							)		:	Tuple									=   config_and_inputs
							a_		:	List[Any]									=   {'''input_ids''': input_ids, '''attention_mask''': input_mask}
							return config, inputs_dict





@require_tf
class 						SCREAMING_SNAKE_CASE       (       SCREAMING_SNAKE_CASE_						, SCREAMING_SNAKE_CASE_						, unittest.TestCase							):
			snake_case__	=					(
			    (
			        TFEsmModel,
			        TFEsmForMaskedLM,
			        TFEsmForSequenceClassification,
			        TFEsmForTokenClassification,
			    )
			    if is_tf_available()
			    else ()
			)
			snake_case__	=					(
			    {
			        "feature-extraction": TFEsmModel,
			        "fill-mask": TFEsmForMaskedLM,
			        "text-classification": TFEsmForSequenceClassification,
			        "token-classification": TFEsmForTokenClassification,
			        "zero-shot": TFEsmForSequenceClassification,
			    }
			    if is_tf_available()
			    else {}
			)
			snake_case__	=					False
			snake_case__	=					False






			def       SCREAMING_SNAKE_CASE   (			self   :      Any       )      ->      Any:
							a_		:	List[str]									=   TFEsmModelTester(self       )
							a_		:	Tuple									=   ConfigTester(self	,    config_class=__SCREAMING_SNAKE_CASE	,    hidden_size=37       )






			def       SCREAMING_SNAKE_CASE   (			self   :      List[Any]       )      ->      Union[str, Any]:
							self.config_tester.run_common_tests()






			def       SCREAMING_SNAKE_CASE   (			self   :      Union[str, Any]       )      ->      str:
							a_		:	str									=   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_model(*__SCREAMING_SNAKE_CASE       )






			def       SCREAMING_SNAKE_CASE   (			self   :      List[str]       )      ->      str:
							a_		:	Dict									=   self.model_tester.prepare_config_and_inputs_for_decoder()
							self.model_tester.create_and_check_model_as_decoder(*__SCREAMING_SNAKE_CASE       )






			def       SCREAMING_SNAKE_CASE   (			self   :      List[Any]       )      ->      Tuple:
							a_		:	str									=   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_masked_lm(*__SCREAMING_SNAKE_CASE       )






			def       SCREAMING_SNAKE_CASE   (			self   :      str       )      ->      str:
							a_		:	Any									=   self.model_tester.prepare_config_and_inputs()
							self.model_tester.create_and_check_for_token_classification(*__SCREAMING_SNAKE_CASE       )






			@slow
			def       SCREAMING_SNAKE_CASE   (			self   :      int       )      ->      Any:
							for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											a_		:	List[Any]									=   TFEsmModel.from_pretrained(__SCREAMING_SNAKE_CASE       )
											self.assertIsNotNone(__SCREAMING_SNAKE_CASE       )






			@unittest.skip('''Protein models do not support embedding resizing.'''       )
			def       SCREAMING_SNAKE_CASE   (			self   :      Union[str, Any]       )      ->      List[Any]:
							pass






			@unittest.skip('''Protein models do not support embedding resizing.'''       )
			def       SCREAMING_SNAKE_CASE   (			self   :      str       )      ->      List[str]:
							pass







			def       SCREAMING_SNAKE_CASE   (			self   :      Union[str, Any]       )      ->      Optional[Any]:
							a_					,      a_		:	Union[str, Any]									=   self.model_tester.prepare_config_and_inputs_for_common()

							for model_class in self.all_model_classes:
											a_		:	Union[str, Any]									=   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
															a_		:	int									=   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:
															a_		:	Union[str, Any]									=   model.get_output_embeddings()
															assert x is None
															a_		:	int									=   model.get_bias()
															assert name is None





@require_tf
class 						SCREAMING_SNAKE_CASE       (       unittest.TestCase							):






			@slow
			def       SCREAMING_SNAKE_CASE   (			self   :      str       )      ->      List[str]:
							a_		:	str									=   TFEsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''       )

							a_		:	Any									=   tf.constant([[0, 1, 2, 3, 4, 5]]       )
							a_		:	List[Any]									=   model(__SCREAMING_SNAKE_CASE       )[0]
							a_		:	Optional[int]									=   [1, 6, 33]
							self.assertEqual(list(output.numpy().shape       )	,    __SCREAMING_SNAKE_CASE       )
							# compare the actual values for a slice.
							a_		:	Optional[Any]									=   tf.constant(
							    [
							        [
							            [8.92_1518, -10.58_9814, -6.467_1307],
							            [-6.396_7156, -13.91_1377, -1.121_1915],
							            [-7.78_1247, -13.95_1557, -3.74_0592],
							        ]
							    ]       )
							self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()	,    expected_slice.numpy()	,    atol=1e-2       )       )







			@slow
			def       SCREAMING_SNAKE_CASE   (			self   :      Tuple       )      ->      Dict:
							a_		:	Union[str, Any]									=   TFEsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''       )

							a_		:	Tuple									=   tf.constant([[0, 6, 4, 13, 5, 4, 16, 12, 11, 7, 2]]       )
							a_		:	str									=   model(__SCREAMING_SNAKE_CASE       )[0]
							# compare the actual values for a slice.
							a_		:	Dict									=   tf.constant(
							    [
							        [
							            [0.1444_3092, 0.5412_5327, 0.324_7739],
							            [0.3034_0484, 0.0052_6676, 0.3107_7722],
							            [0.3227_8043, -0.2498_7096, 0.341_4628],
							        ]
							    ]       )
							self.assertTrue(numpy.allclose(output[:, :3, :3].numpy()	,    expected_slice.numpy()	,    atol=1e-4       )       )












'''simple docstring'''







import os
import time

import pytest

from datasets.utils.filelock import FileLock, Timeout





def 		_UpperCAmelCase (				__A				: int   ):
				a_		:	Optional[Any]									=   FileLock(str(tmpdir / '''foo.lock'''   )   )
				a_		:	Union[str, Any]									=   FileLock(str(tmpdir / '''foo.lock'''   )   )
				a_		:	List[Any]									=   0.01
				with locka.acquire():
								with pytest.raises(__A   ):
												a_		:	List[Any]									=   time.time()
												locka.acquire(__A   )
												assert time.time() - _start > timeout





def 		_UpperCAmelCase (				__A				: Tuple   ):
				a_		:	Tuple									=   '''a''' * 10_00 + '''.lock'''
				a_		:	List[Any]									=   FileLock(str(tmpdir / filename   )   )
				assert locka._lock_file.endswith('''.lock'''   )
				assert not locka._lock_file.endswith(__A   )
				assert len(os.path.basename(locka._lock_file   )   ) <= 2_55
				a_		:	Optional[int]									=   FileLock(tmpdir / filename   )
				with locka.acquire():
								with pytest.raises(__A   ):
												locka.acquire(0   )




