





from typing import Optional

import torch
import torch.utils.checkpoint
from torch import Tensor, nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import (
    BackboneOutput,
    BaseModelOutputWithNoAttention,
    BaseModelOutputWithPoolingAndNoAttention,
    ImageClassifierOutputWithNoAttention,
)
from ...modeling_utils import PreTrainedModel
from ...utils import (
    add_code_sample_docstrings,
    add_start_docstrings,
    add_start_docstrings_to_model_forward,
    logging,
    replace_return_docstrings,
)
from ...utils.backbone_utils import BackboneMixin
from .configuration_resnet import ResNetConfig


_a						=   logging.get_logger(__name__)

# General docstring
_a						=   "ResNetConfig"

# Base docstring
_a						=   "microsoft/resnet-50"
_a						=   [1, 2_048, 7, 7]

# Image classification docstring
_a						=   "microsoft/resnet-50"
_a						=   "tiger cat"

_a						=   [
    "microsoft/resnet-50",
    # See all resnet models at https://huggingface.co/models?filter=resnet
]






class       __A					(					nn.Module				):

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase = 3 ,						__lowerCAmelCase = 1 ,						__lowerCAmelCase = "relu"							):



				'''simple docstring'''







				super().__init__()
				lowerCamelCase__              =	nn.Convad(
				    __lowerCAmelCase ,						__lowerCAmelCase ,						kernel_size=__lowerCAmelCase ,						stride=__lowerCAmelCase ,						padding=kernel_size // 2 ,						bias=__lowerCAmelCase							)
				lowerCamelCase__              =	nn.BatchNormad(__lowerCAmelCase							)
				lowerCamelCase__              =	ACTaFN[activation] if activation is not None else nn.Identity()





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	self.convolution(__lowerCAmelCase							)
				lowerCamelCase__              =	self.normalization(__lowerCAmelCase							)
				lowerCamelCase__              =	self.activation(__lowerCAmelCase							)
				return hidden_state






class       __A					(					nn.Module				):

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				super().__init__()
				lowerCamelCase__              =	ResNetConvLayer(
				    config.num_channels ,						config.embedding_size ,						kernel_size=7 ,						stride=2 ,						activation=config.hidden_act							)
				lowerCamelCase__              =	nn.MaxPoolad(kernel_size=3 ,						stride=2 ,						padding=1							)
				lowerCamelCase__              =	config.num_channels





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	pixel_values.shape[1]
				if num_channels != self.num_channels:
					raise ValueError(
					    '''Make sure that the channel dimension of the pixel values match with the one set in the configuration.'''							)
				lowerCamelCase__              =	self.embedder(__lowerCAmelCase							)
				lowerCamelCase__              =	self.pooler(__lowerCAmelCase							)
				return embedding






class       __A					(					nn.Module				):

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase = 2							):



				'''simple docstring'''







				super().__init__()
				lowerCamelCase__              =	nn.Convad(__lowerCAmelCase ,						__lowerCAmelCase ,						kernel_size=1 ,						stride=__lowerCAmelCase ,						bias=__lowerCAmelCase							)
				lowerCamelCase__              =	nn.BatchNormad(__lowerCAmelCase							)





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	self.convolution(__lowerCAmelCase							)
				lowerCamelCase__              =	self.normalization(__lowerCAmelCase							)
				return hidden_state






class       __A					(					nn.Module				):

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase = 1 ,						__lowerCAmelCase = "relu"							):



				'''simple docstring'''







				super().__init__()
				lowerCamelCase__              =	in_channels != out_channels or stride != 1
				lowerCamelCase__              =	(
				    ResNetShortCut(__lowerCAmelCase ,						__lowerCAmelCase ,						stride=__lowerCAmelCase							) if should_apply_shortcut else nn.Identity()
				)
				lowerCamelCase__              =	nn.Sequential(
				    ResNetConvLayer(__lowerCAmelCase ,						__lowerCAmelCase ,						stride=__lowerCAmelCase							) ,						ResNetConvLayer(__lowerCAmelCase ,						__lowerCAmelCase ,						activation=__lowerCAmelCase							) ,						)
				lowerCamelCase__              =	ACTaFN[activation]





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	hidden_state
				lowerCamelCase__              =	self.layer(__lowerCAmelCase							)
				lowerCamelCase__              =	self.shortcut(__lowerCAmelCase							)
				hidden_state += residual
				lowerCamelCase__              =	self.activation(__lowerCAmelCase							)
				return hidden_state






class       __A					(					nn.Module				):

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase = 1 ,						__lowerCAmelCase = "relu" ,						__lowerCAmelCase = 4							):



				'''simple docstring'''







				super().__init__()
				lowerCamelCase__              =	in_channels != out_channels or stride != 1
				lowerCamelCase__              =	out_channels // reduction
				lowerCamelCase__              =	(
				    ResNetShortCut(__lowerCAmelCase ,						__lowerCAmelCase ,						stride=__lowerCAmelCase							) if should_apply_shortcut else nn.Identity()
				)
				lowerCamelCase__              =	nn.Sequential(
				    ResNetConvLayer(__lowerCAmelCase ,						__lowerCAmelCase ,						kernel_size=1							) ,						ResNetConvLayer(__lowerCAmelCase ,						__lowerCAmelCase ,						stride=__lowerCAmelCase							) ,						ResNetConvLayer(__lowerCAmelCase ,						__lowerCAmelCase ,						kernel_size=1 ,						activation=__lowerCAmelCase							) ,						)
				lowerCamelCase__              =	ACTaFN[activation]





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	hidden_state
				lowerCamelCase__              =	self.layer(__lowerCAmelCase							)
				lowerCamelCase__              =	self.shortcut(__lowerCAmelCase							)
				hidden_state += residual
				lowerCamelCase__              =	self.activation(__lowerCAmelCase							)
				return hidden_state






class       __A					(					nn.Module				):

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase = 2 ,						__lowerCAmelCase = 2 ,						):



				'''simple docstring'''







				super().__init__()

				lowerCamelCase__              =	ResNetBottleNeckLayer if config.layer_type == '''bottleneck''' else ResNetBasicLayer

				lowerCamelCase__              =	nn.Sequential(
				    # downsampling is done in the first layer with stride of 2
				    layer(__lowerCAmelCase ,						__lowerCAmelCase ,						stride=__lowerCAmelCase ,						activation=config.hidden_act							) ,						*[layer(__lowerCAmelCase ,						__lowerCAmelCase ,						activation=config.hidden_act							) for _ in range(depth - 1							)] ,						)





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	input
				for layer in self.layers:
					lowerCamelCase__              =	layer(__lowerCAmelCase							)
				return hidden_state






class       __A					(					nn.Module				):

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				super().__init__()
				lowerCamelCase__              =	nn.ModuleList([]							)
				# based on `downsample_in_first_stage` the first layer of the first stage may or may not downsample the input
				self.stages.append(
				    ResNetStage(
				        __lowerCAmelCase ,						config.embedding_size ,						config.hidden_sizes[0] ,						stride=2 if config.downsample_in_first_stage else 1 ,						depth=config.depths[0] ,						)							)
				lowerCamelCase__              =	zip(config.hidden_sizes ,						config.hidden_sizes[1:]							)
				for (in_channels, out_channels), depth in zip(__lowerCAmelCase ,						config.depths[1:]							):
					self.stages.append(ResNetStage(__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						depth=__lowerCAmelCase							)							)





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase = False ,						__lowerCAmelCase = True							):



				'''simple docstring'''







				lowerCamelCase__              =	() if output_hidden_states else None

				for stage_module in self.stages:
					if output_hidden_states:
						lowerCamelCase__              =	hidden_states + (hidden_state,)

					lowerCamelCase__              =	stage_module(__lowerCAmelCase							)

				if output_hidden_states:
					lowerCamelCase__              =	hidden_states + (hidden_state,)

				if not return_dict:
					return tuple(v for v in [hidden_state, hidden_states] if v is not None							)

				return BaseModelOutputWithNoAttention(
				    last_hidden_state=__lowerCAmelCase ,						hidden_states=__lowerCAmelCase ,						)






class       __A					(					lowerCAmelCase				):

			'''simple docstring'''







			lowerCAmelCase_							  =		ResNetConfig
			lowerCAmelCase_							  =		"""resnet"""
			lowerCAmelCase_							  =		"""pixel_values"""
			lowerCAmelCase_							  =		True





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				if isinstance(__lowerCAmelCase ,						nn.Convad							):
					nn.init.kaiming_normal_(module.weight ,						mode='''fan_out''' ,						nonlinearity='''relu'''							)
				elif isinstance(__lowerCAmelCase ,						(nn.BatchNormad, nn.GroupNorm)							):
					nn.init.constant_(module.weight ,						1							)
					nn.init.constant_(module.bias ,						0							)





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase=False							):



				'''simple docstring'''







				if isinstance(__lowerCAmelCase ,						__lowerCAmelCase							):
					lowerCamelCase__              =	value



_a						=   r"\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it\n    as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`ResNetConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n"

_a						=   r"\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`ConvNextImageProcessor.__call__`] for details.\n\n        output_hidden_states (`bool`, *optional*):\n            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for\n            more detail.\n        return_dict (`bool`, *optional*):\n            Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.\n"






@add_start_docstrings(
    """The bare ResNet model outputting raw features without any specific head on top."""    ,       lowerCAmelCase    ,       )
class       __A					(					lowerCAmelCase				):

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				super().__init__(__lowerCAmelCase							)
				lowerCamelCase__              =	config
				lowerCamelCase__              =	ResNetEmbeddings(__lowerCAmelCase							)
				lowerCamelCase__              =	ResNetEncoder(__lowerCAmelCase							)
				lowerCamelCase__              =	nn.AdaptiveAvgPoolad((1, 1)							)
				# Initialize weights and apply final processing
				self.post_init()





			@add_start_docstrings_to_model_forward(__lowerCAmelCase							)
			@add_code_sample_docstrings(
			    checkpoint=_CHECKPOINT_FOR_DOC ,						output_type=__lowerCAmelCase ,						config_class=_CONFIG_FOR_DOC ,						modality='''vision''' ,						expected_output=_EXPECTED_OUTPUT_SHAPE ,						)
			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase = None ,						__lowerCAmelCase = None							):



				'''simple docstring'''







				lowerCamelCase__              =	(
				    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
				)
				lowerCamelCase__              =	return_dict if return_dict is not None else self.config.use_return_dict

				lowerCamelCase__              =	self.embedder(__lowerCAmelCase							)

				lowerCamelCase__              =	self.encoder(
				    __lowerCAmelCase ,						output_hidden_states=__lowerCAmelCase ,						return_dict=__lowerCAmelCase							)

				lowerCamelCase__              =	encoder_outputs[0]

				lowerCamelCase__              =	self.pooler(__lowerCAmelCase							)

				if not return_dict:
					return (last_hidden_state, pooled_output) + encoder_outputs[1:]

				return BaseModelOutputWithPoolingAndNoAttention(
				    last_hidden_state=__lowerCAmelCase ,						pooler_output=__lowerCAmelCase ,						hidden_states=encoder_outputs.hidden_states ,						)






@add_start_docstrings(
    """
    ResNet Model with an image classification head on top (a linear layer on top of the pooled features), e.g. for
    ImageNet.
    """    ,       lowerCAmelCase    ,       )
class       __A					(					lowerCAmelCase				):

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				super().__init__(__lowerCAmelCase							)
				lowerCamelCase__              =	config.num_labels
				lowerCamelCase__              =	ResNetModel(__lowerCAmelCase							)
				# classification head
				lowerCamelCase__              =	nn.Sequential(
				    nn.Flatten() ,						nn.Linear(config.hidden_sizes[-1] ,						config.num_labels							) if config.num_labels > 0 else nn.Identity() ,						)
				# initialize weights and apply final processing
				self.post_init()





			@add_start_docstrings_to_model_forward(__lowerCAmelCase							)
			@add_code_sample_docstrings(
			    checkpoint=_IMAGE_CLASS_CHECKPOINT ,						output_type=__lowerCAmelCase ,						config_class=_CONFIG_FOR_DOC ,						expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT ,						)
			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase = None ,						__lowerCAmelCase = None ,						__lowerCAmelCase = None ,						__lowerCAmelCase = None ,						):



				'''simple docstring'''







				lowerCamelCase__              =	return_dict if return_dict is not None else self.config.use_return_dict

				lowerCamelCase__              =	self.resnet(__lowerCAmelCase ,						output_hidden_states=__lowerCAmelCase ,						return_dict=__lowerCAmelCase							)

				lowerCamelCase__              =	outputs.pooler_output if return_dict else outputs[1]

				lowerCamelCase__              =	self.classifier(__lowerCAmelCase							)

				lowerCamelCase__              =	None

				if labels is not None:
					if self.config.problem_type is None:
						if self.num_labels == 1:
							lowerCamelCase__              =	'''regression'''
						elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
							lowerCamelCase__              =	'''single_label_classification'''
						else:
							lowerCamelCase__              =	'''multi_label_classification'''
					if self.config.problem_type == "regression":
						lowerCamelCase__              =	MSELoss()
						if self.num_labels == 1:
							lowerCamelCase__              =	loss_fct(logits.squeeze() ,						labels.squeeze()							)
						else:
							lowerCamelCase__              =	loss_fct(__lowerCAmelCase ,						__lowerCAmelCase							)
					elif self.config.problem_type == "single_label_classification":
						lowerCamelCase__              =	CrossEntropyLoss()
						lowerCamelCase__              =	loss_fct(logits.view(-1 ,						self.num_labels							) ,						labels.view(-1							)							)
					elif self.config.problem_type == "multi_label_classification":
						lowerCamelCase__              =	BCEWithLogitsLoss()
						lowerCamelCase__              =	loss_fct(__lowerCAmelCase ,						__lowerCAmelCase							)

				if not return_dict:
					lowerCamelCase__              =	(logits,) + outputs[2:]
					return (loss,) + output if loss is not None else output

				return ImageClassifierOutputWithNoAttention(loss=__lowerCAmelCase ,						logits=__lowerCAmelCase ,						hidden_states=outputs.hidden_states							)






@add_start_docstrings(
    """
    ResNet backbone, to be used with frameworks like DETR and MaskFormer.
    """    ,       lowerCAmelCase    ,       )
class       __A					(					lowerCAmelCase    ,       lowerCAmelCase				):

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase							):



				'''simple docstring'''







				super().__init__(__lowerCAmelCase							)
				super()._init_backbone(__lowerCAmelCase							)

				lowerCamelCase__              =	[config.embedding_size] + config.hidden_sizes
				lowerCamelCase__              =	ResNetEmbeddings(__lowerCAmelCase							)
				lowerCamelCase__              =	ResNetEncoder(__lowerCAmelCase							)

				# initialize weights and apply final processing
				self.post_init()





			@add_start_docstrings_to_model_forward(__lowerCAmelCase							)
			@replace_return_docstrings(output_type=__lowerCAmelCase ,						config_class=_CONFIG_FOR_DOC							)
			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase = None ,						__lowerCAmelCase = None							):



				'''simple docstring'''







				lowerCamelCase__              =	return_dict if return_dict is not None else self.config.use_return_dict
				lowerCamelCase__              =	(
				    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
				)

				lowerCamelCase__              =	self.embedder(__lowerCAmelCase							)

				lowerCamelCase__              =	self.encoder(__lowerCAmelCase ,						output_hidden_states=__lowerCAmelCase ,						return_dict=__lowerCAmelCase							)

				lowerCamelCase__              =	outputs.hidden_states

				lowerCamelCase__              =	()
				for idx, stage in enumerate(self.stage_names							):
					if stage in self.out_features:
						feature_maps += (hidden_states[idx],)

				if not return_dict:
					lowerCamelCase__              =	(feature_maps,)
					if output_hidden_states:
						output += (outputs.hidden_states,)
					return output

				return BackboneOutput(
				    feature_maps=__lowerCAmelCase ,						hidden_states=outputs.hidden_states if output_hidden_states else None ,						attentions=__lowerCAmelCase ,						)










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

			'''simple docstring'''







			def __init__(					self ,						__lowerCAmelCase ,						):



				'''simple docstring'''







				lowerCamelCase__              =	parent
				lowerCamelCase__              =	1_3
				lowerCamelCase__              =	7
				lowerCamelCase__              =	True
				lowerCamelCase__              =	True
				lowerCamelCase__              =	True
				lowerCamelCase__              =	9_9
				lowerCamelCase__              =	3_2
				lowerCamelCase__              =	2
				lowerCamelCase__              =	4
				lowerCamelCase__              =	3_7
				lowerCamelCase__              =	'''gelu'''
				lowerCamelCase__              =	0.1
				lowerCamelCase__              =	0.1
				lowerCamelCase__              =	5_1_2
				lowerCamelCase__              =	1_6
				lowerCamelCase__              =	2
				lowerCamelCase__              =	0.02
				lowerCamelCase__              =	3
				lowerCamelCase__              =	4
				lowerCamelCase__              =	None





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	ids_tensor([self.batch_size, self.seq_length] ,						self.vocab_size							)

				lowerCamelCase__              =	None
				if self.use_input_mask:
					lowerCamelCase__              =	random_attention_mask([self.batch_size, self.seq_length]							)

				lowerCamelCase__              =	None
				lowerCamelCase__              =	None
				lowerCamelCase__              =	None
				if self.use_labels:
					lowerCamelCase__              =	ids_tensor([self.batch_size] ,						self.type_sequence_label_size							)
					lowerCamelCase__              =	ids_tensor([self.batch_size, self.seq_length] ,						self.num_labels							)
					lowerCamelCase__              =	ids_tensor([self.batch_size] ,						self.num_choices							)

				lowerCamelCase__              =	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 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				(
				    (
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			(
				    lowerCamelCase__
				)       ,			
				)              =	self.prepare_config_and_inputs()

				lowerCamelCase__              =	True
				lowerCamelCase__              =	floats_tensor([self.batch_size, self.seq_length, self.hidden_size]							)
				lowerCamelCase__              =	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 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	TFEsmModel(config=__lowerCAmelCase							)
				lowerCamelCase__              =	{'''input_ids''': input_ids, '''attention_mask''': input_mask}
				lowerCamelCase__              =	model(__lowerCAmelCase							)

				lowerCamelCase__              =	[input_ids, input_mask]
				lowerCamelCase__              =	model(__lowerCAmelCase							)

				lowerCamelCase__              =	model(__lowerCAmelCase							)

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





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						):



				'''simple docstring'''







				lowerCamelCase__              =	True

				lowerCamelCase__              =	TFEsmModel(config=__lowerCAmelCase							)
				lowerCamelCase__              =	{
				    '''input_ids''': input_ids,
				    '''attention_mask''': input_mask,
				    '''encoder_hidden_states''': encoder_hidden_states,
				    '''encoder_attention_mask''': encoder_attention_mask,
				}
				lowerCamelCase__              =	model(__lowerCAmelCase							)

				lowerCamelCase__              =	[input_ids, input_mask]
				lowerCamelCase__              =	model(__lowerCAmelCase ,						encoder_hidden_states=__lowerCAmelCase							)

				# Also check the case where encoder outputs are not passed
				lowerCamelCase__              =	model(__lowerCAmelCase ,						attention_mask=__lowerCAmelCase							)

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





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	TFEsmForMaskedLM(config=__lowerCAmelCase							)
				lowerCamelCase__              =	model([input_ids, input_mask]							)
				self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.seq_length, self.vocab_size)							)





			def 	__lowerCamelCase       (					self ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase							):



				'''simple docstring'''







				lowerCamelCase__              =	self.num_labels
				lowerCamelCase__              =	TFEsmForTokenClassification(config=__lowerCAmelCase							)
				lowerCamelCase__              =	{'''input_ids''': input_ids, '''attention_mask''': input_mask}
				lowerCamelCase__              =	model(__lowerCAmelCase							)
				self.parent.assertEqual(result.logits.shape ,						(self.batch_size, self.seq_length, self.num_labels)							)





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







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






@require_tf
class       __A					(					lowerCAmelCase    ,       lowerCAmelCase    ,       unittest.TestCase				):

			'''simple docstring'''







			lowerCAmelCase_							  =		(
			    (
			        TFEsmModel,
			        TFEsmForMaskedLM,
			        TFEsmForSequenceClassification,
			        TFEsmForTokenClassification,
			    )
			    if is_tf_available()
			    else ()
			)
			lowerCAmelCase_							  =		(
			    {
			        """feature-extraction""": TFEsmModel,
			        """fill-mask""": TFEsmForMaskedLM,
			        """text-classification""": TFEsmForSequenceClassification,
			        """token-classification""": TFEsmForTokenClassification,
			        """zero-shot""": TFEsmForSequenceClassification,
			    }
			    if is_tf_available()
			    else {}
			)
			lowerCAmelCase_							  =		False
			lowerCAmelCase_							  =		False





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	TFEsmModelTester(self							)
				lowerCamelCase__              =	ConfigTester(self ,						config_class=__lowerCAmelCase ,						hidden_size=3_7							)





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				self.config_tester.run_common_tests()





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*__lowerCAmelCase							)





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	self.model_tester.prepare_config_and_inputs_for_decoder()
				self.model_tester.create_and_check_model_as_decoder(*__lowerCAmelCase							)





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_masked_lm(*__lowerCAmelCase							)





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_token_classification(*__lowerCAmelCase							)





			@slow
			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				for model_name in TF_ESM_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					lowerCamelCase__              =	TFEsmModel.from_pretrained(__lowerCAmelCase							)
					self.assertIsNotNone(__lowerCAmelCase							)





			@unittest.skip('''Protein models do not support embedding resizing.'''							)
			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				pass





			@unittest.skip('''Protein models do not support embedding resizing.'''							)
			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				pass





			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__       ,			lowerCamelCase__              =	self.model_tester.prepare_config_and_inputs_for_common()

				for model_class in self.all_model_classes:
					lowerCamelCase__              =	model_class(__lowerCAmelCase							)
					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
						lowerCamelCase__              =	model.get_bias()
						assert isinstance(__lowerCAmelCase ,						__lowerCAmelCase							)
						for k, v in name.items():
							assert isinstance(__lowerCAmelCase ,						tf.Variable							)
					else:
						lowerCamelCase__              =	model.get_output_embeddings()
						assert x is None
						lowerCamelCase__              =	model.get_bias()
						assert name is None






@require_tf
class       __A					(					unittest.TestCase				):

			'''simple docstring'''







			@slow
			def 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	TFEsmForMaskedLM.from_pretrained('''facebook/esm2_t6_8M_UR50D'''							)

				lowerCamelCase__              =	tf.constant([[0, 1, 2, 3, 4, 5]]							)
				lowerCamelCase__              =	model(__lowerCAmelCase							)[0]
				lowerCamelCase__              =	[1, 6, 3_3]
				self.assertEqual(list(output.numpy().shape							) ,						__lowerCAmelCase							)
				# compare the actual values for a slice.
				lowerCamelCase__              =	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 	__lowerCamelCase       (					self							):



				'''simple docstring'''







				lowerCamelCase__              =	TFEsmModel.from_pretrained('''facebook/esm2_t6_8M_UR50D'''							)

				lowerCamelCase__              =	tf.constant([[0, 6, 4, 1_3, 5, 4, 1_6, 1_2, 1_1, 7, 2]]							)
				lowerCamelCase__              =	model(__lowerCAmelCase							)[0]
				# compare the actual values for a slice.
				lowerCamelCase__              =	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							)							)



