"""simple docstring"""




import collections.abc
from typing import Optional, Tuple, Union

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

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithNoAttention, ImageClassifierOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_poolformer import PoolFormerConfig


__lowerCamelCase											=  logging.get_logger(__name__)

# General docstring
__lowerCamelCase											=  "PoolFormerConfig"

# Base docstring
__lowerCamelCase											=  "sail/poolformer_s12"
__lowerCamelCase											=  [1, 5_12, 7, 7]

# Image classification docstring
__lowerCamelCase											=  "sail/poolformer_s12"
__lowerCamelCase											=  "tabby, tabby cat"

__lowerCamelCase											=  [
    "sail/poolformer_s12",
    # See all PoolFormer models at https://huggingface.co/models?filter=poolformer
]






def   UpperCAmelCase					(	UpperCamelCase__       ,    UpperCamelCase__ = 0.0       ,    UpperCamelCase__ = False				):






							"""simple docstring"""





							if drop_prob == 0.0 or not training:
														return input
							A__	     =							1 - drop_prob
							A__	     =							(input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
							A__	     =							keep_prob + torch.rand(UpperCamelCase__       ,    dtype=input.dtype       ,    device=input.device				)
							random_tensor.floor_()  # binarize
							A__	     =							input.div(UpperCamelCase__				) * random_tensor
							return output
class     UpperCamelCase__(			nn.Module ):





						def __init__(       self		,__UpperCAmelCase = None     )					->				None:
													super().__init__()
													A__	     =							drop_prob





						def 				snake_case__	(       self		,__UpperCAmelCase     )					->				torch.Tensor:
													return drop_path(__UpperCAmelCase		,self.drop_prob		,self.training     )





						def 				snake_case__	(       self     )					->				str:
													return "p={}".format(self.drop_prob     )






class     UpperCamelCase__(			nn.Module ):





						def __init__(       self		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase=None     )					->				str:
													super().__init__()
													A__	     =							patch_size if isinstance(__UpperCAmelCase		,collections.abc.Iterable     ) else (patch_size, patch_size)
													A__	     =							stride if isinstance(__UpperCAmelCase		,collections.abc.Iterable     ) else (stride, stride)
													A__	     =							padding if isinstance(__UpperCAmelCase		,collections.abc.Iterable     ) else (padding, padding)

													A__	     =							nn.Convad(__UpperCAmelCase		,__UpperCAmelCase		,kernel_size=__UpperCAmelCase		,stride=__UpperCAmelCase		,padding=__UpperCAmelCase     )
													A__	     =							norm_layer(__UpperCAmelCase     ) if norm_layer else nn.Identity()





						def 				snake_case__	(       self		,__UpperCAmelCase     )					->				List[Any]:
													A__	     =							self.projection(__UpperCAmelCase     )
													A__	     =							self.norm(__UpperCAmelCase     )
													return embeddings






class     UpperCamelCase__(			nn.GroupNorm ):





						def __init__(       self		,__UpperCAmelCase		,**__UpperCAmelCase     )					->				Union[str, Any]:
													super().__init__(1		,__UpperCAmelCase		,**__UpperCAmelCase     )






class     UpperCamelCase__(			nn.Module ):





						def __init__(       self		,__UpperCAmelCase     )					->				Optional[int]:
													super().__init__()
													A__	     =							nn.AvgPoolad(__UpperCAmelCase		,stride=1		,padding=pool_size // 2		,count_include_pad=__UpperCAmelCase     )





						def 				snake_case__	(       self		,__UpperCAmelCase     )					->				Optional[int]:
													return self.pool(__UpperCAmelCase     ) - hidden_states






class     UpperCamelCase__(			nn.Module ):





						def __init__(       self		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase     )					->				List[Any]:
													super().__init__()
													A__	     =							nn.Convad(__UpperCAmelCase		,__UpperCAmelCase		,1     )
													A__	     =							nn.Convad(__UpperCAmelCase		,__UpperCAmelCase		,1     )
													A__	     =							PoolFormerDropPath(__UpperCAmelCase     )
													if isinstance(config.hidden_act		,__UpperCAmelCase     ):
																				A__	     =							ACTaFN[config.hidden_act]
													else:
																				A__	     =							config.hidden_act





						def 				snake_case__	(       self		,__UpperCAmelCase     )					->				Tuple:
													A__	     =							self.conva(__UpperCAmelCase     )
													A__	     =							self.act_fn(__UpperCAmelCase     )
													A__	     =							self.drop(__UpperCAmelCase     )
													A__	     =							self.conva(__UpperCAmelCase     )
													A__	     =							self.drop(__UpperCAmelCase     )

													return hidden_states






class     UpperCamelCase__(			nn.Module ):





						def __init__(       self		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase     )					->				List[Any]:
													super().__init__()
													A__	     =							PoolFormerPooling(__UpperCAmelCase     )
													A__	     =							PoolFormerOutput(__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase		,__UpperCAmelCase     )
													A__	     =							PoolFormerGroupNorm(__UpperCAmelCase     )
													A__	     =							PoolFormerGroupNorm(__UpperCAmelCase     )

													# Useful for training neural nets
													A__	     =							PoolFormerDropPath(__UpperCAmelCase     ) if drop_path > 0.0 else nn.Identity()
													A__	     =							config.use_layer_scale
													if config.use_layer_scale:
																				A__	     =							nn.Parameter(
																				    config.layer_scale_init_value * torch.ones((__UpperCAmelCase)     )		,requires_grad=__UpperCAmelCase     )
																				A__	     =							nn.Parameter(
																				    config.layer_scale_init_value * torch.ones((__UpperCAmelCase)     )		,requires_grad=__UpperCAmelCase     )





						def 				snake_case__	(       self		,__UpperCAmelCase     )					->				int:
													if self.use_layer_scale:
																				A__	     =							self.pooling(self.before_norm(__UpperCAmelCase     )     )
																				A__	     =							self.layer_scale_a.unsqueeze(-1     ).unsqueeze(-1     ) * pooling_output
																				# First residual connection
																				A__	     =							hidden_states + self.drop_path(__UpperCAmelCase     )
																				A__	     =							()

																				A__	     =							self.output(self.after_norm(__UpperCAmelCase     )     )
																				A__	     =							self.layer_scale_a.unsqueeze(-1     ).unsqueeze(-1     ) * layer_output
																				# Second residual connection
																				A__	     =							hidden_states + self.drop_path(__UpperCAmelCase     )

																				A__	     =							(output,) + outputs
																				return outputs

													else:
																				A__	     =							self.drop_path(self.pooling(self.before_norm(__UpperCAmelCase     )     )     )
																				# First residual connection
																				A__	     =							pooling_output + hidden_states
																				A__	     =							()

																				# Second residual connection inside the PoolFormerOutput block
																				A__	     =							self.drop_path(self.output(self.after_norm(__UpperCAmelCase     )     )     )
																				A__	     =							hidden_states + layer_output

																				A__	     =							(output,) + outputs
																				return outputs






class     UpperCamelCase__(			nn.Module ):





						def __init__(       self		,__UpperCAmelCase     )					->				List[Any]:
													super().__init__()
													A__	     =							config
													# stochastic depth decay rule
													A__	     =							[x.item() for x in torch.linspace(0		,config.drop_path_rate		,sum(config.depths     )     )]

													# patch embeddings
													A__	     =							[]
													for i in range(config.num_encoder_blocks     ):
																				embeddings.append(
																				    PoolFormerEmbeddings(
																				        patch_size=config.patch_sizes[i]		,stride=config.strides[i]		,padding=config.padding[i]		,num_channels=config.num_channels if i == 0 else config.hidden_sizes[i - 1]		,hidden_size=config.hidden_sizes[i]		,)     )
													A__	     =							nn.ModuleList(__UpperCAmelCase     )

													# Transformer blocks
													A__	     =							[]
													A__	     =							0
													for i in range(config.num_encoder_blocks     ):
																				# each block consists of layers
																				A__	     =							[]
																				if i != 0:
																											cur += config.depths[i - 1]
																				for j in range(config.depths[i]     ):
																											layers.append(
																											    PoolFormerLayer(
																											        __UpperCAmelCase		,num_channels=config.hidden_sizes[i]		,pool_size=config.pool_size		,hidden_size=config.hidden_sizes[i]		,intermediate_size=int(config.hidden_sizes[i] * config.mlp_ratio     )		,drop_path=dpr[cur + j]		,)     )
																				blocks.append(nn.ModuleList(__UpperCAmelCase     )     )

													A__	     =							nn.ModuleList(__UpperCAmelCase     )





						def 				snake_case__	(       self		,__UpperCAmelCase		,__UpperCAmelCase=False		,__UpperCAmelCase=True     )					->				Any:
													A__	     =							() if output_hidden_states else None

													A__	     =							pixel_values
													for idx, layers in enumerate(zip(self.patch_embeddings		,self.block     )     ):
																				A__     ,  A__	     =							layers
																				# Get patch embeddings from hidden_states
																				A__	     =							embedding_layer(__UpperCAmelCase     )
																				# Send the embeddings through the blocks
																				for _, blk in enumerate(__UpperCAmelCase     ):
																											A__	     =							blk(__UpperCAmelCase     )
																											A__	     =							layer_outputs[0]

																				if output_hidden_states:
																											A__	     =							all_hidden_states + (hidden_states,)

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

													return BaseModelOutputWithNoAttention(last_hidden_state=__UpperCAmelCase		,hidden_states=__UpperCAmelCase     )






class     UpperCamelCase__(			__A ):
						lowerCAmelCase__  :						Union[str, Any]													=		PoolFormerConfig
						lowerCAmelCase__  :						Tuple													=		'poolformer'
						lowerCAmelCase__  :						List[Any]													=		'pixel_values'
						lowerCAmelCase__  :						Union[str, Any]													=		True





						def 				snake_case__	(       self		,__UpperCAmelCase     )					->				Any:
													if isinstance(__UpperCAmelCase		,(nn.Linear, nn.Convad)     ):
																				module.weight.data.normal_(mean=0.0		,std=self.config.initializer_range     )
																				if module.bias is not None:
																											module.bias.data.zero_()
													elif isinstance(__UpperCAmelCase		,nn.LayerNorm     ):
																				module.bias.data.zero_()
																				module.weight.data.fill_(1.0     )





						def 				snake_case__	(       self		,__UpperCAmelCase		,__UpperCAmelCase=False     )					->				Tuple:
													if isinstance(__UpperCAmelCase		,__UpperCAmelCase     ):
																				A__	     =							value






__lowerCamelCase											=  R"\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use\n    it 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 ([`PoolFormerConfig`]): 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"

__lowerCamelCase											=  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            [`PoolFormerImageProcessor.__call__`] for details.\n"
@add_start_docstrings(
    'The bare PoolFormer Model transformer outputting raw hidden-states without any specific head on top.'		,						__A		,						)
class     UpperCamelCase__(			__A ):





						def __init__(       self		,__UpperCAmelCase     )					->				Optional[Any]:
													super().__init__(__UpperCAmelCase     )
													A__	     =							config

													A__	     =							PoolFormerEncoder(__UpperCAmelCase     )

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





						def 				snake_case__	(       self     )					->				int:
													return self.embeddings.patch_embeddings





						@add_start_docstrings_to_model_forward(__UpperCAmelCase     )
						@add_code_sample_docstrings(
						    checkpoint=_CHECKPOINT_FOR_DOC		,output_type=__UpperCAmelCase		,config_class=_CONFIG_FOR_DOC		,modality='vision'		,expected_output=_EXPECTED_OUTPUT_SHAPE		,)
						def 				snake_case__	(       self		,__UpperCAmelCase = None		,__UpperCAmelCase = None		,__UpperCAmelCase = None		,)					->				Union[Tuple, BaseModelOutputWithNoAttention]:
													A__	     =							(
													    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
													)
													A__	     =							return_dict if return_dict is not None else self.config.use_return_dict

													if pixel_values is None:
																				raise ValueError('You have to specify pixel_values'     )

													A__	     =							self.encoder(
													    __UpperCAmelCase		,output_hidden_states=__UpperCAmelCase		,return_dict=__UpperCAmelCase		,)
													A__	     =							encoder_outputs[0]

													if not return_dict:
																				return (sequence_output, None) + encoder_outputs[1:]

													return BaseModelOutputWithNoAttention(
													    last_hidden_state=__UpperCAmelCase		,hidden_states=encoder_outputs.hidden_states		,)






class     UpperCamelCase__(			nn.Module ):





						def __init__(       self		,__UpperCAmelCase     )					->				str:
													super().__init__()
													A__	     =							nn.Linear(config.hidden_size		,config.hidden_size     )





						def 				snake_case__	(       self		,__UpperCAmelCase     )					->				Optional[int]:
													A__	     =							self.dense(__UpperCAmelCase     )
													return output






@add_start_docstrings(
    '\n    PoolFormer Model transformer with an image classification head on top\n    '		,						__A		,						)
class     UpperCamelCase__(			__A ):





						def __init__(       self		,__UpperCAmelCase     )					->				Union[str, Any]:
													super().__init__(__UpperCAmelCase     )
													A__	     =							config.num_labels
													A__	     =							PoolFormerModel(__UpperCAmelCase     )

													# Final norm
													A__	     =							PoolFormerGroupNorm(config.hidden_sizes[-1]     )
													# Classifier head
													A__	     =							(
													    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(__UpperCAmelCase     )
						@add_code_sample_docstrings(
						    checkpoint=_IMAGE_CLASS_CHECKPOINT		,output_type=__UpperCAmelCase		,config_class=_CONFIG_FOR_DOC		,expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT		,)
						def 				snake_case__	(       self		,__UpperCAmelCase = None		,__UpperCAmelCase = None		,__UpperCAmelCase = None		,__UpperCAmelCase = None		,)					->				Union[Tuple, ImageClassifierOutputWithNoAttention]:
													A__	     =							return_dict if return_dict is not None else self.config.use_return_dict

													A__	     =							self.poolformer(
													    __UpperCAmelCase		,output_hidden_states=__UpperCAmelCase		,return_dict=__UpperCAmelCase		,)

													A__	     =							outputs[0]

													A__	     =							self.classifier(self.norm(__UpperCAmelCase     ).mean([-2, -1]     )     )

													A__	     =							None
													if labels is not None:
																				if self.config.problem_type is None:
																											if self.num_labels == 1:
																																		A__	     =							'regression'
																											elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
																																		A__	     =							'single_label_classification'
																											else:
																																		A__	     =							'multi_label_classification'

																				if self.config.problem_type == "regression":
																											A__	     =							MSELoss()
																											if self.num_labels == 1:
																																		A__	     =							loss_fct(logits.squeeze()		,labels.squeeze()     )
																											else:
																																		A__	     =							loss_fct(__UpperCAmelCase		,__UpperCAmelCase     )
																				elif self.config.problem_type == "single_label_classification":
																											A__	     =							CrossEntropyLoss()
																											A__	     =							loss_fct(logits.view(-1		,self.num_labels     )		,labels.view(-1     )     )
																				elif self.config.problem_type == "multi_label_classification":
																											A__	     =							BCEWithLogitsLoss()
																											A__	     =							loss_fct(__UpperCAmelCase		,__UpperCAmelCase     )

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

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






"""simple docstring"""




import os
from typing import List, Optional, Union

from ...tokenization_utils import PreTrainedTokenizer
from ...tokenization_utils_base import AddedToken
from ...utils import logging


__lowerCamelCase											=  logging.get_logger(__name__)

__lowerCamelCase											=  {"vocab_file": "vocab.txt"}

__lowerCamelCase											=  {
    "vocab_file": {
        "facebook/esm2_t6_8M_UR50D": "https://huggingface.co/facebook/esm2_t6_8M_UR50D/resolve/main/vocab.txt",
        "facebook/esm2_t12_35M_UR50D": "https://huggingface.co/facebook/esm2_t12_35M_UR50D/resolve/main/vocab.txt",
    },
}

__lowerCamelCase											=  {
    "facebook/esm2_t6_8M_UR50D": 10_24,
    "facebook/esm2_t12_35M_UR50D": 10_24,
}






def   UpperCAmelCase					(	UpperCamelCase__				):






							"""simple docstring"""





							with open(UpperCamelCase__       ,    'r'				) as f:
														A__	     =							f.read().splitlines()
														return [l.strip() for l in lines]


class     UpperCamelCase__(			__A ):
						lowerCAmelCase__  :						Any													=		VOCAB_FILES_NAMES
						lowerCAmelCase__  :						List[str]													=		PRETRAINED_VOCAB_FILES_MAP
						lowerCAmelCase__  :						List[Any]													=		PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
						lowerCAmelCase__  :						Dict													=		['input_ids', 'attention_mask']





						def __init__(       self		,__UpperCAmelCase		,__UpperCAmelCase="<unk>"		,__UpperCAmelCase="<cls>"		,__UpperCAmelCase="<pad>"		,__UpperCAmelCase="<mask>"		,__UpperCAmelCase="<eos>"		,**__UpperCAmelCase		,)					->				Optional[int]:
													super().__init__(**__UpperCAmelCase     )
													A__	     =							load_vocab_file(__UpperCAmelCase     )
													A__	     =							dict(enumerate(self.all_tokens     )     )
													A__	     =							{tok: ind for ind, tok in enumerate(self.all_tokens     )}
													A__	     =							unk_token
													A__	     =							cls_token
													A__	     =							pad_token
													A__	     =							mask_token
													A__	     =							eos_token
													A__	     =							self.all_tokens
													self._create_trie(self.unique_no_split_tokens     )





						def 				snake_case__	(       self		,__UpperCAmelCase     )					->				str:
													return self._id_to_token.get(__UpperCAmelCase		,self.unk_token     )





						def 				snake_case__	(       self		,__UpperCAmelCase     )					->				int:
													return self._token_to_id.get(__UpperCAmelCase		,self._token_to_id.get(self.unk_token     )     )





						def 				snake_case__	(       self		,__UpperCAmelCase		,**__UpperCAmelCase     )					->				Optional[int]:
													return text.split()





						def 				snake_case__	(       self		,__UpperCAmelCase=False     )					->				int:
													return len(self._id_to_token     )





						def 				snake_case__	(       self     )					->				List[str]:
													return {token: i for i, token in enumerate(self.all_tokens     )}





						def 				snake_case__	(       self		,__UpperCAmelCase     )					->				int:
													return self._token_to_id.get(__UpperCAmelCase		,self._token_to_id.get(self.unk_token     )     )





						def 				snake_case__	(       self		,__UpperCAmelCase     )					->				str:
													return self._id_to_token.get(__UpperCAmelCase		,self.unk_token     )





						def 				snake_case__	(       self		,__UpperCAmelCase		,__UpperCAmelCase = None     )					->				List[int]:
													A__	     =							[self.cls_token_id]
													A__	     =							[self.eos_token_id]  # No sep token in ESM vocabulary
													if token_ids_a is None:
																				if self.eos_token_id is None:
																											return cls + token_ids_a
																				else:
																											return cls + token_ids_a + sep
													elif self.eos_token_id is None:
																				raise ValueError('Cannot tokenize multiple sequences when EOS token is not set!'     )
													return cls + token_ids_a + sep + token_ids_a + sep  # Multiple inputs always have an EOS token





						def 				snake_case__	(       self		,__UpperCAmelCase		,__UpperCAmelCase = None		,__UpperCAmelCase = False     )					->				List[int]:
													if already_has_special_tokens:
																				if token_ids_a is not None:
																											raise ValueError(
																											    'You should not supply a second sequence if the provided sequence of '
																											    'ids is already formatted with special tokens for the model.'     )

																				return [1 if token in self.all_special_ids else 0 for token in token_ids_a]
													A__	     =							[1] + ([0] * len(__UpperCAmelCase     )) + [1]
													if token_ids_a is not None:
																				mask += [0] * len(__UpperCAmelCase     ) + [1]
													return mask





						def 				snake_case__	(       self		,__UpperCAmelCase		,__UpperCAmelCase     )					->				List[str]:
													A__	     =							os.path.join(__UpperCAmelCase		,(filename_prefix + '-' if filename_prefix else '') + 'vocab.txt'     )
													with open(__UpperCAmelCase		,'w'     ) as f:
																				f.write('\n'.join(self.all_tokens     )     )
													return (vocab_file,)





						@property
						def 				snake_case__	(       self     )					->				int:
													return self.get_vocab_size(with_added_tokens=__UpperCAmelCase     )





						def 				snake_case__	(       self		,__UpperCAmelCase		,__UpperCAmelCase = False     )					->				int:
													return super()._add_tokens(__UpperCAmelCase		,special_tokens=__UpperCAmelCase     )





