
"""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


SCREAMING_SNAKE_CASE_					     =	logging.get_logger(__name__)

# General docstring
SCREAMING_SNAKE_CASE_					     =	'''PoolFormerConfig'''

# Base docstring
SCREAMING_SNAKE_CASE_					     =	'''sail/poolformer_s12'''
SCREAMING_SNAKE_CASE_					     =	[1, 512, 7, 7]

# Image classification docstring
SCREAMING_SNAKE_CASE_					     =	'''sail/poolformer_s12'''
SCREAMING_SNAKE_CASE_					     =	'''tabby, tabby cat'''

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

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



							'''simple docstring'''





							def __init__( self    ,				snake_case_ = None				)       ->  None:
													super().__init__()
													__lowerCAmelCase     =  drop_prob




							def    A__						( self    ,				snake_case_				)       ->  torch.Tensor:
													return drop_path(snake_case_    ,				self.drop_prob    ,				self.training				)




							def    A__						( self				)       ->  str:
													return "p={}".format(self.drop_prob				)
class       lowerCAmelCase_		(				nn.Module ):



							'''simple docstring'''





							def __init__( self    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_=None				)       ->  List[Any]:
													super().__init__()
													__lowerCAmelCase     =  patch_size if isinstance(snake_case_    ,				collections.abc.Iterable				) else (patch_size, patch_size)
													__lowerCAmelCase     =  stride if isinstance(snake_case_    ,				collections.abc.Iterable				) else (stride, stride)
													__lowerCAmelCase     =  padding if isinstance(snake_case_    ,				collections.abc.Iterable				) else (padding, padding)

													__lowerCAmelCase     =  nn.Convad(snake_case_    ,				snake_case_    ,				kernel_size=snake_case_    ,				stride=snake_case_    ,				padding=snake_case_				)
													__lowerCAmelCase     =  norm_layer(snake_case_				) if norm_layer else nn.Identity()




							def    A__						( self    ,				snake_case_				)       ->  Optional[int]:
													__lowerCAmelCase     =  self.projection(snake_case_				)
													__lowerCAmelCase     =  self.norm(snake_case_				)
													return embeddings
class       lowerCAmelCase_		(				nn.GroupNorm ):



							'''simple docstring'''





							def __init__( self    ,				snake_case_    ,				**snake_case_				)       ->  Optional[Any]:
													super().__init__(1    ,				snake_case_    ,				**snake_case_				)
class       lowerCAmelCase_		(				nn.Module ):



							'''simple docstring'''





							def __init__( self    ,				snake_case_				)       ->  int:
													super().__init__()
													__lowerCAmelCase     =  nn.AvgPoolad(snake_case_    ,				stride=1    ,				padding=pool_size // 2    ,				count_include_pad=snake_case_				)




							def    A__						( self    ,				snake_case_				)       ->  Optional[int]:
													return self.pool(snake_case_				) - hidden_states
class       lowerCAmelCase_		(				nn.Module ):



							'''simple docstring'''





							def __init__( self    ,				snake_case_    ,				snake_case_    ,				snake_case_    ,				snake_case_				)       ->  Tuple:
													super().__init__()
													__lowerCAmelCase     =  nn.Convad(snake_case_    ,				snake_case_    ,				1				)
													__lowerCAmelCase     =  nn.Convad(snake_case_    ,				snake_case_    ,				1				)
													__lowerCAmelCase     =  PoolFormerDropPath(snake_case_				)
													if isinstance(config.hidden_act    ,				snake_case_				):
																			__lowerCAmelCase     =  ACTaFN[config.hidden_act]
													else:
																			__lowerCAmelCase     =  config.hidden_act




							def    A__						( self    ,				snake_case_				)       ->  Dict:
													__lowerCAmelCase     =  self.conva(snake_case_				)
													__lowerCAmelCase     =  self.act_fn(snake_case_				)
													__lowerCAmelCase     =  self.drop(snake_case_				)
													__lowerCAmelCase     =  self.conva(snake_case_				)
													__lowerCAmelCase     =  self.drop(snake_case_				)

													return hidden_states
class       lowerCAmelCase_		(				nn.Module ):



							'''simple docstring'''





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

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




							def    A__						( self    ,				snake_case_				)       ->  List[str]:
													if self.use_layer_scale:
																			__lowerCAmelCase     =  self.pooling(self.before_norm(snake_case_				)				)
																			__lowerCAmelCase     =  self.layer_scale_a.unsqueeze(-1				).unsqueeze(-1				) * pooling_output
																			# First residual connection
																			__lowerCAmelCase     =  hidden_states + self.drop_path(snake_case_				)
																			__lowerCAmelCase     =  ()

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

																			__lowerCAmelCase     =  (output,) + outputs
																			return outputs

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

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

																			__lowerCAmelCase     =  (output,) + outputs
																			return outputs
class       lowerCAmelCase_		(				nn.Module ):



							'''simple docstring'''





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

													# patch embeddings
													__lowerCAmelCase     =  []
													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]    ,				)				)
													__lowerCAmelCase     =  nn.ModuleList(snake_case_				)

													# Transformer blocks
													__lowerCAmelCase     =  []
													__lowerCAmelCase     =  0
													for i in range(config.num_encoder_blocks				):
																			# each block consists of layers
																			__lowerCAmelCase     =  []
																			if i != 0:
																									cur += config.depths[i - 1]
																			for j in range(config.depths[i]				):
																									layers.append(
																									    PoolFormerLayer(
																									        snake_case_    ,				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(snake_case_				)				)

													__lowerCAmelCase     =  nn.ModuleList(snake_case_				)




							def    A__						( self    ,				snake_case_    ,				snake_case_=False    ,				snake_case_=True				)       ->  Dict:
													__lowerCAmelCase     =  () if output_hidden_states else None

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

																			if output_hidden_states:
																									__lowerCAmelCase     =  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=snake_case_    ,				hidden_states=snake_case_				)
class       lowerCAmelCase_		(				A__ ):



							'''simple docstring'''





							_snake_case						      =		PoolFormerConfig
							_snake_case						      =		'''poolformer'''
							_snake_case						      =		'''pixel_values'''
							_snake_case						      =		True




							def    A__						( self    ,				snake_case_				)       ->  List[str]:
													if isinstance(snake_case_    ,				(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(snake_case_    ,				nn.LayerNorm				):
																			module.bias.data.zero_()
																			module.weight.data.fill_(1.0				)




							def    A__						( self    ,				snake_case_    ,				snake_case_=False				)       ->  List[str]:
													if isinstance(snake_case_    ,				snake_case_				):
																			__lowerCAmelCase     =  value
SCREAMING_SNAKE_CASE_					     =	R'''
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`PoolFormerConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
'''

SCREAMING_SNAKE_CASE_					     =	R'''
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`PoolFormerImageProcessor.__call__`] for details.
'''
@add_start_docstrings(
    '''The bare PoolFormer Model transformer outputting raw hidden-states without any specific head on top.'''				,				A__				,				)
class       lowerCAmelCase_		(				A__ ):



							'''simple docstring'''





							def __init__( self    ,				snake_case_				)       ->  Optional[int]:
													super().__init__(snake_case_				)
													__lowerCAmelCase     =  config

													__lowerCAmelCase     =  PoolFormerEncoder(snake_case_				)

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




							def    A__						( self				)       ->  Dict:
													return self.embeddings.patch_embeddings




							@add_start_docstrings_to_model_forward(snake_case_				)
							@add_code_sample_docstrings(
							    checkpoint=_CHECKPOINT_FOR_DOC    ,				output_type=snake_case_    ,				config_class=_CONFIG_FOR_DOC    ,				modality="""vision"""    ,				expected_output=_EXPECTED_OUTPUT_SHAPE    ,				)
							def    A__						( self    ,				snake_case_ = None    ,				snake_case_ = None    ,				snake_case_ = None    ,				)       ->  Union[Tuple, BaseModelOutputWithNoAttention]:
													__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

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

													__lowerCAmelCase     =  self.encoder(
													    snake_case_    ,				output_hidden_states=snake_case_    ,				return_dict=snake_case_    ,				)
													__lowerCAmelCase     =  encoder_outputs[0]

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

													return BaseModelOutputWithNoAttention(
													    last_hidden_state=snake_case_    ,				hidden_states=encoder_outputs.hidden_states    ,				)
class       lowerCAmelCase_		(				nn.Module ):



							'''simple docstring'''





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




							def    A__						( self    ,				snake_case_				)       ->  Any:
													__lowerCAmelCase     =  self.dense(snake_case_				)
													return output
@add_start_docstrings(
    '''
    PoolFormer Model transformer with an image classification head on top
    '''				,				A__				,				)
class       lowerCAmelCase_		(				A__ ):



							'''simple docstring'''





							def __init__( self    ,				snake_case_				)       ->  List[Any]:
													super().__init__(snake_case_				)
													__lowerCAmelCase     =  config.num_labels
													__lowerCAmelCase     =  PoolFormerModel(snake_case_				)

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

													__lowerCAmelCase     =  self.poolformer(
													    snake_case_    ,				output_hidden_states=snake_case_    ,				return_dict=snake_case_    ,				)

													__lowerCAmelCase     =  outputs[0]

													__lowerCAmelCase     =  self.classifier(self.norm(snake_case_				).mean([-2, -1]				)				)

													__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(snake_case_    ,				snake_case_				)
																			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(snake_case_    ,				snake_case_				)

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

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


"""simple docstring"""
from typing import List

from ...configuration_utils import PretrainedConfig
from ...utils import logging


SCREAMING_SNAKE_CASE_					     =	logging.get_logger(__name__)

SCREAMING_SNAKE_CASE_					     =	{
    '''snap-research/efficientformer-l1-300''': (
        '''https://huggingface.co/snap-research/efficientformer-l1-300/resolve/main/config.json'''
    ),
}
class       lowerCAmelCase_		(				A__ ):



							'''simple docstring'''





							_snake_case						      =		'''efficientformer'''




							def __init__( self    ,				snake_case_ = [3, 2, 6, 4]    ,				snake_case_ = [48, 96, 224, 448]    ,				snake_case_ = [True, True, True, True]    ,				snake_case_ = 448    ,				snake_case_ = 32    ,				snake_case_ = 4    ,				snake_case_ = 7    ,				snake_case_ = 5    ,				snake_case_ = 8    ,				snake_case_ = 4    ,				snake_case_ = 0.0    ,				snake_case_ = 16    ,				snake_case_ = 3    ,				snake_case_ = 3    ,				snake_case_ = 3    ,				snake_case_ = 2    ,				snake_case_ = 1    ,				snake_case_ = 0.0    ,				snake_case_ = 1    ,				snake_case_ = True    ,				snake_case_ = True    ,				snake_case_ = 1e-5    ,				snake_case_ = "gelu"    ,				snake_case_ = 0.02    ,				snake_case_ = 1e-1_2    ,				snake_case_ = 224    ,				snake_case_ = 1e-0_5    ,				**snake_case_    ,				)       ->  None:
													super().__init__(**snake_case_				)

													__lowerCAmelCase     =  hidden_act
													__lowerCAmelCase     =  hidden_dropout_prob
													__lowerCAmelCase     =  hidden_sizes
													__lowerCAmelCase     =  num_hidden_layers
													__lowerCAmelCase     =  num_attention_heads
													__lowerCAmelCase     =  initializer_range
													__lowerCAmelCase     =  layer_norm_eps
													__lowerCAmelCase     =  patch_size
													__lowerCAmelCase     =  num_channels
													__lowerCAmelCase     =  depths
													__lowerCAmelCase     =  mlp_expansion_ratio
													__lowerCAmelCase     =  downsamples
													__lowerCAmelCase     =  dim
													__lowerCAmelCase     =  key_dim
													__lowerCAmelCase     =  attention_ratio
													__lowerCAmelCase     =  resolution
													__lowerCAmelCase     =  pool_size
													__lowerCAmelCase     =  downsample_patch_size
													__lowerCAmelCase     =  downsample_stride
													__lowerCAmelCase     =  downsample_pad
													__lowerCAmelCase     =  drop_path_rate
													__lowerCAmelCase     =  num_metaad_blocks
													__lowerCAmelCase     =  distillation
													__lowerCAmelCase     =  use_layer_scale
													__lowerCAmelCase     =  layer_scale_init_value
													__lowerCAmelCase     =  image_size
													__lowerCAmelCase     =  batch_norm_eps
