


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


__UpperCamelCase						=      logging.get_logger(__name__)

# General docstring
__UpperCamelCase						=      '''PoolFormerConfig'''

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

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

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


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




		'''simple docstring'''







		def __init__(	self						,       lowerCAmelCase__ = None				)     ->							None:
								super().__init__()
								SCREAMING_SNAKE_CASE													=					drop_prob







		def 							__A						(	self						,       lowerCAmelCase__				)     ->							torch.Tensor:
								return drop_path(lowerCAmelCase__						,       self.drop_prob						,       self.training				)







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





class 					lowerCAmelCase	(					nn.Module    ):




		'''simple docstring'''







		def __init__(	self						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__=None				)     ->							Union[str, Any]:
								super().__init__()
								SCREAMING_SNAKE_CASE													=					patch_size if isinstance(lowerCAmelCase__						,       collections.abc.Iterable				) else (patch_size, patch_size)
								SCREAMING_SNAKE_CASE													=					stride if isinstance(lowerCAmelCase__						,       collections.abc.Iterable				) else (stride, stride)
								SCREAMING_SNAKE_CASE													=					padding if isinstance(lowerCAmelCase__						,       collections.abc.Iterable				) else (padding, padding)

								SCREAMING_SNAKE_CASE													=					nn.Convad(lowerCAmelCase__						,       lowerCAmelCase__						,       kernel_size=lowerCAmelCase__						,       stride=lowerCAmelCase__						,       padding=lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					norm_layer(lowerCAmelCase__				) if norm_layer else nn.Identity()







		def 							__A						(	self						,       lowerCAmelCase__				)     ->							Union[str, Any]:
								SCREAMING_SNAKE_CASE													=					self.projection(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					self.norm(lowerCAmelCase__				)
								return embeddings





class 					lowerCAmelCase	(					nn.GroupNorm    ):




		'''simple docstring'''







		def __init__(	self						,       lowerCAmelCase__						,       **lowerCAmelCase__				)     ->							str:
								super().__init__(1						,       lowerCAmelCase__						,       **lowerCAmelCase__				)





class 					lowerCAmelCase	(					nn.Module    ):




		'''simple docstring'''







		def __init__(	self						,       lowerCAmelCase__				)     ->							Any:
								super().__init__()
								SCREAMING_SNAKE_CASE													=					nn.AvgPoolad(lowerCAmelCase__						,       stride=1						,       padding=pool_size // 2						,       count_include_pad=lowerCAmelCase__				)







		def 							__A						(	self						,       lowerCAmelCase__				)     ->							int:
								return self.pool(lowerCAmelCase__				) - hidden_states





class 					lowerCAmelCase	(					nn.Module    ):




		'''simple docstring'''







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







		def 							__A						(	self						,       lowerCAmelCase__				)     ->							Optional[int]:
								SCREAMING_SNAKE_CASE													=					self.conva(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					self.act_fn(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					self.drop(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					self.conva(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					self.drop(lowerCAmelCase__				)

								return hidden_states





class 					lowerCAmelCase	(					nn.Module    ):




		'''simple docstring'''







		def __init__(	self						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__				)     ->							Any:
								super().__init__()
								SCREAMING_SNAKE_CASE													=					PoolFormerPooling(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					PoolFormerOutput(lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__						,       lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					PoolFormerGroupNorm(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					PoolFormerGroupNorm(lowerCAmelCase__				)

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







		def 							__A						(	self						,       lowerCAmelCase__				)     ->							Union[str, Any]:
								if self.use_layer_scale:
														SCREAMING_SNAKE_CASE													=					self.pooling(self.before_norm(lowerCAmelCase__				)				)
														SCREAMING_SNAKE_CASE													=					self.layer_scale_a.unsqueeze(-1				).unsqueeze(-1				) * pooling_output
														# First residual connection
														SCREAMING_SNAKE_CASE													=					hidden_states + self.drop_path(lowerCAmelCase__				)
														SCREAMING_SNAKE_CASE													=					()

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

														SCREAMING_SNAKE_CASE													=					(output,) + outputs
														return outputs

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

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

														SCREAMING_SNAKE_CASE													=					(output,) + outputs
														return outputs





class 					lowerCAmelCase	(					nn.Module    ):




		'''simple docstring'''







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

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

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

								SCREAMING_SNAKE_CASE													=					nn.ModuleList(lowerCAmelCase__				)







		def 							__A						(	self						,       lowerCAmelCase__						,       lowerCAmelCase__=False						,       lowerCAmelCase__=True				)     ->							int:
								SCREAMING_SNAKE_CASE													=					() if output_hidden_states else None

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

														if output_hidden_states:
																				SCREAMING_SNAKE_CASE													=					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=lowerCAmelCase__						,       hidden_states=lowerCAmelCase__				)





class 					lowerCAmelCase	(					lowerCamelCase_    ):




		'''simple docstring'''







		SCREAMING_SNAKE_CASE_						:     Optional[int]								=     PoolFormerConfig
		SCREAMING_SNAKE_CASE_						:     Union[str, Any]								=     """poolformer"""
		SCREAMING_SNAKE_CASE_						:     List[Any]								=     """pixel_values"""
		SCREAMING_SNAKE_CASE_						:     Optional[Any]								=     True







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







		def 							__A						(	self						,       lowerCAmelCase__						,       lowerCAmelCase__=False				)     ->							int:
								if isinstance(lowerCAmelCase__						,       lowerCAmelCase__				):
														SCREAMING_SNAKE_CASE													=					value





__UpperCamelCase						=      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.
'''

__UpperCamelCase						=      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."""	,    lowerCamelCase_	,    )
class 					lowerCAmelCase	(					lowerCamelCase_    ):




		'''simple docstring'''







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

								SCREAMING_SNAKE_CASE													=					PoolFormerEncoder(lowerCAmelCase__				)

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







		def 							__A						(	self				)     ->							Optional[Any]:
								return self.embeddings.patch_embeddings







		@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 							__A						(	self						,       lowerCAmelCase__ = None						,       lowerCAmelCase__ = None						,       lowerCAmelCase__ = None						,       )     ->							Union[Tuple, BaseModelOutputWithNoAttention]:
								SCREAMING_SNAKE_CASE													=					(
								    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
								)
								SCREAMING_SNAKE_CASE													=					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'				)

								SCREAMING_SNAKE_CASE													=					self.encoder(
								    lowerCAmelCase__						,       output_hidden_states=lowerCAmelCase__						,       return_dict=lowerCAmelCase__						,       )
								SCREAMING_SNAKE_CASE													=					encoder_outputs[0]

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

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





class 					lowerCAmelCase	(					nn.Module    ):




		'''simple docstring'''







		def __init__(	self						,       lowerCAmelCase__				)     ->							Union[str, Any]:
								super().__init__()
								SCREAMING_SNAKE_CASE													=					nn.Linear(config.hidden_size						,       config.hidden_size				)







		def 							__A						(	self						,       lowerCAmelCase__				)     ->							Any:
								SCREAMING_SNAKE_CASE													=					self.dense(lowerCAmelCase__				)
								return output





@add_start_docstrings(
    """
    PoolFormer Model transformer with an image classification head on top
    """	,    lowerCamelCase_	,    )
class 					lowerCAmelCase	(					lowerCamelCase_    ):




		'''simple docstring'''







		def __init__(	self						,       lowerCAmelCase__				)     ->							int:
								super().__init__(lowerCAmelCase__				)
								SCREAMING_SNAKE_CASE													=					config.num_labels
								SCREAMING_SNAKE_CASE													=					PoolFormerModel(lowerCAmelCase__				)

								# Final norm
								SCREAMING_SNAKE_CASE													=					PoolFormerGroupNorm(config.hidden_sizes[-1]				)
								# Classifier head
								SCREAMING_SNAKE_CASE													=					(
								    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 							__A						(	self						,       lowerCAmelCase__ = None						,       lowerCAmelCase__ = None						,       lowerCAmelCase__ = None						,       lowerCAmelCase__ = None						,       )     ->							Union[Tuple, ImageClassifierOutputWithNoAttention]:
								SCREAMING_SNAKE_CASE													=					return_dict if return_dict is not None else self.config.use_return_dict

								SCREAMING_SNAKE_CASE													=					self.poolformer(
								    lowerCAmelCase__						,       output_hidden_states=lowerCAmelCase__						,       return_dict=lowerCAmelCase__						,       )

								SCREAMING_SNAKE_CASE													=					outputs[0]

								SCREAMING_SNAKE_CASE													=					self.classifier(self.norm(lowerCAmelCase__				).mean([-2, -1]				)				)

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

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

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

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






"""simple docstring"""




# This is the module that test_patching.py uses to test patch_submodule()

import os  # noqa: this is just for tests
import os as renamed_os  # noqa: this is just for tests
from os import path  # noqa: this is just for tests
from os import path as renamed_path  # noqa: this is just for tests
from os.path import join  # noqa: this is just for tests
from os.path import join as renamed_join  # noqa: this is just for tests


__UpperCamelCase						=      open  # noqa: we just need to have a builtin inside this module to test it properly


