


from typing import List, Optional, Union

import numpy as np
import tensorflow as tf

from .utils import logging


__A	:	List[Any]     =     logging.get_logger(__name__)





def      UpperCAmelCase      (		lowerCamelCase_  :Any ):







		'''simple docstring'''
		if isinstance(__A				,      np.ndarray ):
				return list(tensor.shape )

		snake_case_		:							Optional[Any]			       =						tf.shape(__A )

		if tensor.shape == tf.TensorShape(__A ):
				return dynamic

		snake_case_		:							List[str]			       =						tensor.shape.as_list()

		return [dynamic[i] if s is None else s for i, s in enumerate(__A )]





def      UpperCAmelCase      (		lowerCamelCase_  :int				,      lowerCamelCase_  :str = None				,      lowerCamelCase_  :Dict = None ):







		'''simple docstring'''
		return tf.nn.softmax(logits=logits + 1E-9				,      axis=__A				,      name=__A )





def      UpperCAmelCase      (		lowerCamelCase_  :Dict				,      lowerCamelCase_  :List[str]				,      lowerCamelCase_  :Tuple				,      lowerCamelCase_  :Union[str, Any]=1E-5				,      lowerCamelCase_  :Dict=-1 ):







		'''simple docstring'''
		# This is a very simplified functional layernorm, designed to duplicate
		# the functionality of PyTorch nn.functional.layer_norm when this is needed to port
		# models in Transformers.

		if weight.shape.rank != 1 or bias.shape.rank != 1 or not isinstance(__A				,      __A ):
				raise NotImplementedError("""Only 1D weight and bias tensors are supported for now, with only a single axis.""" )

		# Get mean and variance on the axis to be normalized
		snake_case_       ,      snake_case_		:							List[str]			       =						tf.nn.moments(__A				,      axes=[axis]				,      keepdims=__A )

		if axis != -1:
				# Reshape scale and weight to have the same rank as inputs, but with 1 dimensions
				# on every dimension except axis
				snake_case_		:							Union[str, Any]			       =						[1] * inputs.shape.rank
				snake_case_		:							List[str]			       =						shape_list(__A )[axis]
				snake_case_		:							Union[str, Any]			       =						tf.reshape(__A				,      __A )
				snake_case_		:							Dict			       =						tf.reshape(__A				,      __A )

		# Compute layer normalization using the batch_normalization
		# function.
		snake_case_		:							Tuple			       =						tf.nn.batch_normalization(
		    __A				,      __A				,      __A				,      offset=__A				,      scale=__A				,      variance_epsilon=__A				,      )
		return outputs





def      UpperCAmelCase      (		lowerCamelCase_  :Dict				,      lowerCamelCase_  :Optional[Any]=0				,      lowerCamelCase_  :Optional[Any]=-1 ):







		'''simple docstring'''
		# Replicates the behavior of torch.flatten in TF

		# If end_dim or start_dim is negative, count them from the end
		if end_dim < 0:
				end_dim += input.shape.rank
		if start_dim < 0:
				start_dim += input.shape.rank

		if start_dim == end_dim:
				return input

		snake_case_		:							int			       =						tf.shape(__A )
		snake_case_		:							List[Any]			       =						tf.math.reduce_prod(in_shape[start_dim : end_dim + 1] )
		snake_case_		:							List[Any]			       =						tf.concat([in_shape[:start_dim], [flattened_dim], in_shape[end_dim + 1 :]]				,      axis=0 )
		return tf.reshape(__A				,      __A )





def      UpperCAmelCase      (		lowerCamelCase_  :Tuple ):







		'''simple docstring'''
		if not isinstance(__A				,      tf.Tensor ):
				snake_case_		:							Any			       =						tf.convert_to_tensor(__A )  # Catches stray NumPy inputs
		if encoder_attention_mask.shape.rank == 3:
				snake_case_		:							Dict			       =						encoder_attention_mask[:, None, :, :]
		if encoder_attention_mask.shape.rank == 2:
				snake_case_		:							Optional[int]			       =						encoder_attention_mask[:, None, None, :]
		# T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
		# Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow
		# /transformer/transformer_layers.py#L270
		# encoder_extended_attention_mask = (encoder_extended_attention_mask ==
		# encoder_extended_attention_mask.transpose(-1, -2))
		snake_case_		:							Dict			       =						(
		    tf.cast(1				,      encoder_attention_mask.dtype ) - encoder_extended_attention_mask
		) * encoder_extended_attention_mask.dtype.min

		return encoder_extended_attention_mask





def      UpperCAmelCase      (		lowerCamelCase_  :Dict				,      lowerCamelCase_  :Any				,      lowerCamelCase_  :List[str] = "input_ids" ):







		'''simple docstring'''
		tf.debugging.assert_less(
		    __A				,      tf.cast(__A				,      dtype=tensor.dtype )				,      message=(
		        F'''The maximum value of {tensor_name} ({tf.math.reduce_max(__A )}) must be smaller than the embedding '''
		        F'''layer\'s input dimension ({embed_dim}). The likely cause is some problem at tokenization time.'''
		    )				,      )





def      UpperCAmelCase      (		lowerCamelCase_  :List[str]				,      lowerCamelCase_  :int				,      lowerCamelCase_  :Dict ):







		'''simple docstring'''
		snake_case_		:							Optional[int]			       =						6_45_12
		# Check that no item in `data` is larger than `HDF5_OBJECT_HEADER_LIMIT`
		# because in that case even chunking the array would not make the saving
		# possible.
		snake_case_		:							Optional[Any]			       =						[x for x in data if len(__A ) > HDF5_OBJECT_HEADER_LIMIT]

		# Expecting this to never be true.
		if bad_attributes:
				raise RuntimeError(
				    """The following attributes cannot be saved to HDF5 file because """
				    F'''they are larger than {HDF5_OBJECT_HEADER_LIMIT} '''
				    F'''bytes: {bad_attributes}''' )

		snake_case_		:							Tuple			       =						np.asarray(__A )

		snake_case_		:							Tuple			       =						1
		snake_case_		:							int			       =						np.array_split(__A				,      __A )

		# This will never loop forever thanks to the test above.
		while any(x.nbytes > HDF5_OBJECT_HEADER_LIMIT for x in chunked_data ):
				num_chunks += 1
				snake_case_		:							List[str]			       =						np.array_split(__A				,      __A )

		if num_chunks > 1:
				for chunk_id, chunk_data in enumerate(__A ):
						snake_case_		:							Optional[int]			       =						chunk_data
		else:
				snake_case_		:							Any			       =						data





def      UpperCAmelCase      (		lowerCamelCase_  :List[str]				,      lowerCamelCase_  :Optional[int] ):







		'''simple docstring'''
		if name in group.attrs:
				snake_case_		:							List[Any]			       =						[n.decode("""utf8""" ) if hasattr(__A				,      """decode""" ) else n for n in group.attrs[name]]
		else:
				snake_case_		:							List[str]			       =						[]
				snake_case_		:							Optional[int]			       =						0
				while "%s%d" % (name, chunk_id) in group.attrs:
						data.extend(
						    [n.decode("""utf8""" ) if hasattr(__A				,      """decode""" ) else n for n in group.attrs["""%s%d""" % (name, chunk_id)]] )
						chunk_id += 1
		return data






def      UpperCAmelCase      (		lowerCamelCase_  :Dict ):







		'''simple docstring'''

		def _expand_single_ad_tensor(lowerCamelCase_  :List[str] ):
				if isinstance(__A				,      tf.Tensor ) and t.shape.rank == 1:
						return tf.expand_dims(__A				,      axis=-1 )
				return t

		return tf.nest.map_structure(_expand_single_ad_tensor				,      __A )



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





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





							def 		a__   (  self		:str			):
									snake_case_		:							Dict			       =						ids_tensor([self.batch_size, self.seq_length]  ,self.vocab_size			)

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

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

									snake_case_		:							Dict			       =						self.get_config()

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





							def 		a__   (  self		:Optional[Any]			):
									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 		a__   (  self		:str  ,_UpperCamelCase		:List[Any]  ,_UpperCamelCase		:List[str]  ,_UpperCamelCase		:Any  ,_UpperCamelCase		:Optional[int]  ,_UpperCamelCase		:List[Any]  ,_UpperCamelCase		:Optional[int]			):
									snake_case_		:							Union[str, Any]			       =						EsmModel(config=_UpperCamelCase			)
									model.to(_UpperCamelCase			)
									model.eval()
									snake_case_		:							int			       =						model(_UpperCamelCase  ,attention_mask=_UpperCamelCase			)
									snake_case_		:							str			       =						model(_UpperCamelCase			)
									snake_case_		:							Optional[int]			       =						model(_UpperCamelCase			)

									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 		a__   (  self		:Any  ,_UpperCamelCase		:Optional[Any]  ,_UpperCamelCase		:Optional[Any]  ,_UpperCamelCase		:Any  ,_UpperCamelCase		:Dict  ,_UpperCamelCase		:Union[str, Any]  ,_UpperCamelCase		:Optional[int]			):
									snake_case_		:							str			       =						EsmForMaskedLM(config=_UpperCamelCase			)
									model.to(_UpperCamelCase			)
									model.eval()
									snake_case_		:							Dict			       =						model(_UpperCamelCase  ,attention_mask=_UpperCamelCase  ,labels=_UpperCamelCase			)
									self.parent.assertEqual(result.logits.shape  ,(self.batch_size, self.seq_length, self.vocab_size)			)





							def 		a__   (  self		:Optional[int]  ,_UpperCamelCase		:Optional[Any]  ,_UpperCamelCase		:List[str]  ,_UpperCamelCase		:Union[str, Any]  ,_UpperCamelCase		:Union[str, Any]  ,_UpperCamelCase		:Union[str, Any]  ,_UpperCamelCase		:Any			):
									snake_case_		:							List[Any]			       =						self.num_labels
									snake_case_		:							int			       =						EsmForTokenClassification(config=_UpperCamelCase			)
									model.to(_UpperCamelCase			)
									model.eval()
									snake_case_		:							int			       =						model(_UpperCamelCase  ,attention_mask=_UpperCamelCase  ,labels=_UpperCamelCase			)
									self.parent.assertEqual(result.logits.shape  ,(self.batch_size, self.seq_length, self.num_labels)			)





							def 		a__   (  self		:Union[str, Any]			):
									snake_case_		:							Dict			       =						self.prepare_config_and_inputs()
									(
									    (
									    snake_case_
									)       ,      (
									    snake_case_
									)       ,      (
									    snake_case_
									)       ,      (
									    snake_case_
									)       ,      (
									    snake_case_
									)       ,      (
									    snake_case_
									)       ,      
									)		:							Optional[int]			       =						config_and_inputs
									snake_case_		:							List[str]			       =						{"""input_ids""": input_ids, """attention_mask""": input_mask}
									return config, inputs_dict





@require_torch
class 						__UpperCamelCase	(						lowercase__						,  lowercase__						,  unittest.TestCase							):
							lowercase			:		Optional[int]   		=  False

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





							def 		a__   (  self		:Any			):
									snake_case_		:							Any			       =						EsmModelTester(self			)
									snake_case_		:							str			       =						ConfigTester(self  ,config_class=_UpperCamelCase  ,hidden_size=3_7			)





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





							def 		a__   (  self		:Optional[Any]			):
									snake_case_		:							str			       =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_model(*_UpperCamelCase			)





							def 		a__   (  self		:Dict			):
									snake_case_		:							Dict			       =						self.model_tester.prepare_config_and_inputs()
									for type in ["absolute", "relative_key", "relative_key_query"]:
											snake_case_		:							int			       =						type
											self.model_tester.create_and_check_model(*_UpperCamelCase			)





							def 		a__   (  self		:Dict			):
									snake_case_		:							str			       =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_masked_lm(*_UpperCamelCase			)





							def 		a__   (  self		:List[str]			):
									snake_case_		:							List[str]			       =						self.model_tester.prepare_config_and_inputs()
									self.model_tester.create_and_check_for_token_classification(*_UpperCamelCase			)





							@slow
							def 		a__   (  self		:Union[str, Any]			):
									for model_name in ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
											snake_case_		:							int			       =						EsmModel.from_pretrained(_UpperCamelCase			)
											self.assertIsNotNone(_UpperCamelCase			)





							def 		a__   (  self		:Tuple			):
									snake_case_		:							Dict			       =						self.model_tester.prepare_config_and_inputs()[0]
									snake_case_		:							Optional[int]			       =						EsmEmbeddings(config=_UpperCamelCase			)

									snake_case_		:							List[Any]			       =						torch.as_tensor([[1_2, 3_1, 1_3, model.padding_idx]]			)
									snake_case_		:							Union[str, Any]			       =						torch.as_tensor(
									    [
									        [
									            0 + model.padding_idx + 1,
									            1 + model.padding_idx + 1,
									            2 + model.padding_idx + 1,
									            model.padding_idx,
									        ]
									    ]			)
									snake_case_		:							Optional[Any]			       =						create_position_ids_from_input_ids(_UpperCamelCase  ,model.padding_idx			)
									self.assertEqual(position_ids.shape  ,expected_positions.shape			)
									self.assertTrue(torch.all(torch.eq(_UpperCamelCase  ,_UpperCamelCase			)			)			)





							def 		a__   (  self		:List[Any]			):
									snake_case_		:							Union[str, Any]			       =						self.model_tester.prepare_config_and_inputs()[0]
									snake_case_		:							List[Any]			       =						EsmEmbeddings(config=_UpperCamelCase			)

									snake_case_		:							Dict			       =						torch.empty(2  ,4  ,3_0			)
									snake_case_		:							List[Any]			       =						[
									    0 + embeddings.padding_idx + 1,
									    1 + embeddings.padding_idx + 1,
									    2 + embeddings.padding_idx + 1,
									    3 + embeddings.padding_idx + 1,
									]
									snake_case_		:							Any			       =						torch.as_tensor([expected_single_positions, expected_single_positions]			)
									snake_case_		:							str			       =						embeddings.create_position_ids_from_inputs_embeds(_UpperCamelCase			)
									self.assertEqual(position_ids.shape  ,expected_positions.shape			)
									self.assertTrue(torch.all(torch.eq(_UpperCamelCase  ,_UpperCamelCase			)			)			)





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





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





							@unittest.skip("""Will be fixed soon by reducing the size of the model used for common tests."""			)
							def 		a__   (  self		:Optional[int]			):
									pass





@require_torch
class 						__UpperCamelCase	(						lowercase__							):





							@slow
							def 		a__   (  self		:Any			):
									with torch.no_grad():
											snake_case_		:							Optional[Any]			       =						EsmForMaskedLM.from_pretrained("""facebook/esm2_t6_8M_UR50D"""			)
											model.eval()
											snake_case_		:							List[Any]			       =						torch.tensor([[0, 1, 2, 3, 4, 5]]			)
											snake_case_		:							Dict			       =						model(_UpperCamelCase			)[0]

											snake_case_		:							Optional[int]			       =						3_3

											snake_case_		:							List[Any]			       =						torch.Size((1, 6, vocab_size)			)
											self.assertEqual(output.shape  ,_UpperCamelCase			)

											snake_case_		:							Union[str, Any]			       =						torch.tensor(
											    [[[8.92_15, -10.58_98, -6.46_71], [-6.39_67, -13.91_14, -1.12_12], [-7.78_12, -13.95_16, -3.74_06]]]			)
											self.assertTrue(torch.allclose(output[:, :3, :3]  ,_UpperCamelCase  ,atol=1E-4			)			)





							@slow
							def 		a__   (  self		:List[Any]			):
									with torch.no_grad():
											snake_case_		:							List[Any]			       =						EsmModel.from_pretrained("""facebook/esm2_t6_8M_UR50D"""			)
											model.eval()

											snake_case_		:							str			       =						torch.tensor([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]			)
											snake_case_		:							Optional[Any]			       =						model(_UpperCamelCase			)[0]
											# compare the actual values for a slice.
											snake_case_		:							List[str]			       =						torch.tensor(
											    [[[0.14_44, 0.54_13, 0.32_48], [0.30_34, 0.00_53, 0.31_08], [0.32_28, -0.24_99, 0.34_15]]]			)
											self.assertTrue(torch.allclose(output[:, :3, :3]  ,_UpperCamelCase  ,atol=1E-4			)			)