





import warnings

from ...utils import logging
from .image_processing_dpt import DPTImageProcessor


a__    :		Optional[Any]      						=     logging.get_logger(__name__)






class  lowercase    (      UpperCamelCase__       ):






					"""simple docstring"""



					def __init__(    self						:     Optional[int] ,					*a_						:     Optional[Any] ,					**a_						:     Optional[Any]							):

						"""simple docstring"""



						warnings.warn(
						    """The class DPTFeatureExtractor is deprecated and will be removed in version 5 of Transformers. Please"""
						    """ use DPTImageProcessor instead.""" ,					__A ,					)
						super().__init__(*__A ,					**__A							)







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, 512, 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 					SCREAMING_SNAKE_CASE__				(							UpperCamelCase__:				Optional[Any]       , UpperCamelCase__:				float = 0.0       , UpperCamelCase__:				bool = False     ):
		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(UpperCamelCase__       , dtype=input.dtype       , device=input.device     )
		random_tensor.floor_()  # binarize
		SCREAMING_SNAKE_CASE__								=   input.div(UpperCamelCase__     ) * random_tensor
		return output




class 	UpperCamelCase_							(  nn.Module     ):


							def __init__(				self     :Optional[Any]				,					__A     :Optional[float] = None		) -> None:



									"""simple docstring"""



									super().__init__()
									SCREAMING_SNAKE_CASE__								=   drop_prob


							def   _snake_case	(				self     :Any				,					__A     :torch.Tensor		) -> torch.Tensor:



									"""simple docstring"""



									return drop_path(__A				,					self.drop_prob				,					self.training		)


							def   _snake_case	(				self     :Dict		) -> str:



									"""simple docstring"""



									return "p={}".format(self.drop_prob		)




class 	UpperCamelCase_							(  nn.Module     ):


							def __init__(				self     :Dict				,					__A     :Optional[Any]				,					__A     :Dict				,					__A     :List[str]				,					__A     :Optional[Any]				,					__A     :Tuple				,					__A     :Optional[Any]=None		) -> Union[str, Any]:



									"""simple docstring"""



									super().__init__()
									SCREAMING_SNAKE_CASE__								=   patch_size if isinstance(__A				,					collections.abc.Iterable		) else (patch_size, patch_size)
									SCREAMING_SNAKE_CASE__								=   stride if isinstance(__A				,					collections.abc.Iterable		) else (stride, stride)
									SCREAMING_SNAKE_CASE__								=   padding if isinstance(__A				,					collections.abc.Iterable		) else (padding, padding)

									SCREAMING_SNAKE_CASE__								=   nn.Convad(__A				,					__A				,					kernel_size=__A				,					stride=__A				,					padding=__A		)
									SCREAMING_SNAKE_CASE__								=   norm_layer(__A		) if norm_layer else nn.Identity()


							def   _snake_case	(				self     :Dict				,					__A     :Union[str, Any]		) -> Optional[Any]:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self.projection(__A		)
									SCREAMING_SNAKE_CASE__								=   self.norm(__A		)
									return embeddings




class 	UpperCamelCase_							(  nn.GroupNorm     ):


							def __init__(				self     :Dict				,					__A     :Tuple				,					**__A     :Union[str, Any]		) -> Dict:



									"""simple docstring"""



									super().__init__(1				,					__A				,					**__A		)




class 	UpperCamelCase_							(  nn.Module     ):


							def __init__(				self     :List[str]				,					__A     :Optional[int]		) -> Any:



									"""simple docstring"""



									super().__init__()
									SCREAMING_SNAKE_CASE__								=   nn.AvgPoolad(__A				,					stride=1				,					padding=pool_size // 2				,					count_include_pad=__A		)


							def   _snake_case	(				self     :Any				,					__A     :Optional[Any]		) -> Optional[Any]:



									"""simple docstring"""



									return self.pool(__A		) - hidden_states




class 	UpperCamelCase_							(  nn.Module     ):


							def __init__(				self     :Optional[Any]				,					__A     :Tuple				,					__A     :Dict				,					__A     :int				,					__A     :Any		) -> str:



									"""simple docstring"""



									super().__init__()
									SCREAMING_SNAKE_CASE__								=   nn.Convad(__A				,					__A				,					1		)
									SCREAMING_SNAKE_CASE__								=   nn.Convad(__A				,					__A				,					1		)
									SCREAMING_SNAKE_CASE__								=   PoolFormerDropPath(__A		)
									if isinstance(config.hidden_act				,					__A		):
											SCREAMING_SNAKE_CASE__								=   ACTaFN[config.hidden_act]
									else:
											SCREAMING_SNAKE_CASE__								=   config.hidden_act


							def   _snake_case	(				self     :Union[str, Any]				,					__A     :Optional[int]		) -> Dict:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self.conva(__A		)
									SCREAMING_SNAKE_CASE__								=   self.act_fn(__A		)
									SCREAMING_SNAKE_CASE__								=   self.drop(__A		)
									SCREAMING_SNAKE_CASE__								=   self.conva(__A		)
									SCREAMING_SNAKE_CASE__								=   self.drop(__A		)

									return hidden_states




class 	UpperCamelCase_							(  nn.Module     ):


							def __init__(				self     :Any				,					__A     :str				,					__A     :List[str]				,					__A     :Tuple				,					__A     :Dict				,					__A     :Union[str, Any]				,					__A     :int		) -> Optional[int]:



									"""simple docstring"""



									super().__init__()
									SCREAMING_SNAKE_CASE__								=   PoolFormerPooling(__A		)
									SCREAMING_SNAKE_CASE__								=   PoolFormerOutput(__A				,					__A				,					__A				,					__A		)
									SCREAMING_SNAKE_CASE__								=   PoolFormerGroupNorm(__A		)
									SCREAMING_SNAKE_CASE__								=   PoolFormerGroupNorm(__A		)

									# Useful for training neural nets
									SCREAMING_SNAKE_CASE__								=   PoolFormerDropPath(__A		) 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((__A)		)				,					requires_grad=__A		)
											SCREAMING_SNAKE_CASE__								=   nn.Parameter(
											    config.layer_scale_init_value * torch.ones((__A)		)				,					requires_grad=__A		)


							def   _snake_case	(				self     :Optional[Any]				,					__A     :Optional[int]		) -> str:



									"""simple docstring"""



									if self.use_layer_scale:
											SCREAMING_SNAKE_CASE__								=   self.pooling(self.before_norm(__A		)		)
											SCREAMING_SNAKE_CASE__								=   self.layer_scale_a.unsqueeze(-1		).unsqueeze(-1		) * pooling_output
											# First residual connection
											SCREAMING_SNAKE_CASE__								=   hidden_states + self.drop_path(__A		)
											SCREAMING_SNAKE_CASE__								=   ()

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

											SCREAMING_SNAKE_CASE__								=   (output,) + outputs
											return outputs

									else:
											SCREAMING_SNAKE_CASE__								=   self.drop_path(self.pooling(self.before_norm(__A		)		)		)
											# 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(__A		)		)		)
											SCREAMING_SNAKE_CASE__								=   hidden_states + layer_output

											SCREAMING_SNAKE_CASE__								=   (output,) + outputs
											return outputs




class 	UpperCamelCase_							(  nn.Module     ):


							def __init__(				self     :Union[str, Any]				,					__A     :List[Any]		) -> Union[str, Any]:



									"""simple docstring"""



									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(__A		)

									# 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(
													        __A				,					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(__A		)		)

									SCREAMING_SNAKE_CASE__								=   nn.ModuleList(__A		)


							def   _snake_case	(				self     :str				,					__A     :Tuple				,					__A     :Dict=False				,					__A     :Tuple=True		) -> Optional[int]:



									"""simple docstring"""



									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(__A		)
											# Send the embeddings through the blocks
											for _, blk in enumerate(__A		):
													SCREAMING_SNAKE_CASE__								=   blk(__A		)
													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=__A				,					hidden_states=__A		)




class 	UpperCamelCase_							(  UpperCamelCase__     ):
							lowerCamelCase_							     =       PoolFormerConfig
							lowerCamelCase_							     =       "poolformer"
							lowerCamelCase_							     =       "pixel_values"
							lowerCamelCase_							     =       True


							def   _snake_case	(				self     :Optional[Any]				,					__A     :Tuple		) -> Dict:



									"""simple docstring"""



									if isinstance(__A				,					(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(__A				,					nn.LayerNorm		):
											module.bias.data.zero_()
											module.weight.data.fill_(1.0		)


							def   _snake_case	(				self     :str				,					__A     :Optional[Any]				,					__A     :Union[str, Any]=False		) -> Any:



									"""simple docstring"""



									if isinstance(__A				,					__A		):
											SCREAMING_SNAKE_CASE__								=   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."					,	UpperCamelCase__					,	)
class 	UpperCamelCase_							(  UpperCamelCase__     ):


							def __init__(				self     :Union[str, Any]				,					__A     :Any		) -> int:



									"""simple docstring"""



									super().__init__(__A		)
									SCREAMING_SNAKE_CASE__								=   config

									SCREAMING_SNAKE_CASE__								=   PoolFormerEncoder(__A		)

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


							def   _snake_case	(				self     :Union[str, Any]		) -> Union[str, Any]:



									"""simple docstring"""



									return self.embeddings.patch_embeddings


							@add_start_docstrings_to_model_forward(__A		)
							@add_code_sample_docstrings(
							    checkpoint=_CHECKPOINT_FOR_DOC				,					output_type=__A				,					config_class=_CONFIG_FOR_DOC				,					modality="""vision"""				,					expected_output=_EXPECTED_OUTPUT_SHAPE				,					)
							def   _snake_case	(				self     :Dict				,					__A     :Optional[torch.FloatTensor] = None				,					__A     :Optional[bool] = None				,					__A     :Optional[bool] = None				,					) -> Union[Tuple, BaseModelOutputWithNoAttention]:



									"""simple docstring"""



									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(
									    __A				,					output_hidden_states=__A				,					return_dict=__A				,					)
									SCREAMING_SNAKE_CASE__								=   encoder_outputs[0]

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

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




class 	UpperCamelCase_							(  nn.Module     ):


							def __init__(				self     :int				,					__A     :Optional[int]		) -> Tuple:



									"""simple docstring"""



									super().__init__()
									SCREAMING_SNAKE_CASE__								=   nn.Linear(config.hidden_size				,					config.hidden_size		)


							def   _snake_case	(				self     :List[Any]				,					__A     :Dict		) -> int:



									"""simple docstring"""



									SCREAMING_SNAKE_CASE__								=   self.dense(__A		)
									return output




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


							def __init__(				self     :str				,					__A     :Union[str, Any]		) -> int:



									"""simple docstring"""



									super().__init__(__A		)
									SCREAMING_SNAKE_CASE__								=   config.num_labels
									SCREAMING_SNAKE_CASE__								=   PoolFormerModel(__A		)

									# 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(__A		)
							@add_code_sample_docstrings(
							    checkpoint=_IMAGE_CLASS_CHECKPOINT				,					output_type=__A				,					config_class=_CONFIG_FOR_DOC				,					expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT				,					)
							def   _snake_case	(				self     :int				,					__A     :Optional[torch.FloatTensor] = None				,					__A     :Optional[torch.LongTensor] = None				,					__A     :Optional[bool] = None				,					__A     :Optional[bool] = None				,					) -> Union[Tuple, ImageClassifierOutputWithNoAttention]:



									"""simple docstring"""



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

									SCREAMING_SNAKE_CASE__								=   self.poolformer(
									    __A				,					output_hidden_states=__A				,					return_dict=__A				,					)

									SCREAMING_SNAKE_CASE__								=   outputs[0]

									SCREAMING_SNAKE_CASE__								=   self.classifier(self.norm(__A		).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(__A				,					__A		)
											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(__A				,					__A		)

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

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