

'''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, 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  a__       (    _SCREAMING_SNAKE_CASE   :       Optional[Any]	,				_SCREAMING_SNAKE_CASE   :       float = 0.0	,				_SCREAMING_SNAKE_CASE   :       bool = False				)     ->					Tuple:


	"""simple docstring"""





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






			def __init__(  self			,_snake_case = None     ):
				super().__init__()
				UpperCAmelCase_     :      List[Any]						=  drop_prob






			def 		UpperCamelCase__							(  self			,_snake_case     ):
				return drop_path(_snake_case			,self.drop_prob			,self.training     )






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




class 					_snake_case		(nn.Module):






			def __init__(  self			,_snake_case			,_snake_case			,_snake_case			,_snake_case			,_snake_case			,_snake_case=None     ):
				super().__init__()
				UpperCAmelCase_     :      int						=  patch_size if isinstance(_snake_case			,collections.abc.Iterable     ) else (patch_size, patch_size)
				UpperCAmelCase_     :      Optional[int]						=  stride if isinstance(_snake_case			,collections.abc.Iterable     ) else (stride, stride)
				UpperCAmelCase_     :      Union[str, Any]						=  padding if isinstance(_snake_case			,collections.abc.Iterable     ) else (padding, padding)

				UpperCAmelCase_     :      Optional[int]						=  nn.Convad(_snake_case			,_snake_case			,kernel_size=_snake_case			,stride=_snake_case			,padding=_snake_case     )
				UpperCAmelCase_     :      Optional[Any]						=  norm_layer(_snake_case     ) if norm_layer else nn.Identity()






			def 		UpperCamelCase__							(  self			,_snake_case     ):
				UpperCAmelCase_     :      Union[str, Any]						=  self.projection(_snake_case     )
				UpperCAmelCase_     :      Any						=  self.norm(_snake_case     )
				return embeddings




class 					_snake_case		(nn.GroupNorm):






			def __init__(  self			,_snake_case			,**_snake_case     ):
				super().__init__(1			,_snake_case			,**_snake_case     )




class 					_snake_case		(nn.Module):






			def __init__(  self			,_snake_case     ):
				super().__init__()
				UpperCAmelCase_     :      List[Any]						=  nn.AvgPoolad(_snake_case			,stride=1			,padding=pool_size // 2			,count_include_pad=_snake_case     )






			def 		UpperCamelCase__							(  self			,_snake_case     ):
				return self.pool(_snake_case     ) - hidden_states




class 					_snake_case		(nn.Module):






			def __init__(  self			,_snake_case			,_snake_case			,_snake_case			,_snake_case     ):
				super().__init__()
				UpperCAmelCase_     :      Optional[Any]						=  nn.Convad(_snake_case			,_snake_case			,1     )
				UpperCAmelCase_     :      int						=  nn.Convad(_snake_case			,_snake_case			,1     )
				UpperCAmelCase_     :      Tuple						=  PoolFormerDropPath(_snake_case     )
				if isinstance(config.hidden_act			,_snake_case     ):
					UpperCAmelCase_     :      Dict						=  ACTaFN[config.hidden_act]
				else:
					UpperCAmelCase_     :      List[Any]						=  config.hidden_act






			def 		UpperCamelCase__							(  self			,_snake_case     ):
				UpperCAmelCase_     :      Optional[Any]						=  self.conva(_snake_case     )
				UpperCAmelCase_     :      Tuple						=  self.act_fn(_snake_case     )
				UpperCAmelCase_     :      Tuple						=  self.drop(_snake_case     )
				UpperCAmelCase_     :      Union[str, Any]						=  self.conva(_snake_case     )
				UpperCAmelCase_     :      str						=  self.drop(_snake_case     )

				return hidden_states




class 					_snake_case		(nn.Module):






			def __init__(  self			,_snake_case			,_snake_case			,_snake_case			,_snake_case			,_snake_case			,_snake_case     ):
				super().__init__()
				UpperCAmelCase_     :      int						=  PoolFormerPooling(_snake_case     )
				UpperCAmelCase_     :      Any						=  PoolFormerOutput(_snake_case			,_snake_case			,_snake_case			,_snake_case     )
				UpperCAmelCase_     :      List[str]						=  PoolFormerGroupNorm(_snake_case     )
				UpperCAmelCase_     :      Tuple						=  PoolFormerGroupNorm(_snake_case     )

				# Useful for training neural nets
				UpperCAmelCase_     :      Optional[Any]						=  PoolFormerDropPath(_snake_case     ) if drop_path > 0.0 else nn.Identity()
				UpperCAmelCase_     :      Union[str, Any]						=  config.use_layer_scale
				if config.use_layer_scale:
					UpperCAmelCase_     :      str						=  nn.Parameter(
					    config.layer_scale_init_value * torch.ones((_snake_case)     )			,requires_grad=_snake_case     )
					UpperCAmelCase_     :      List[Any]						=  nn.Parameter(
					    config.layer_scale_init_value * torch.ones((_snake_case)     )			,requires_grad=_snake_case     )






			def 		UpperCamelCase__							(  self			,_snake_case     ):
				if self.use_layer_scale:
					UpperCAmelCase_     :      Dict						=  self.pooling(self.before_norm(_snake_case     )     )
					UpperCAmelCase_     :      Tuple						=  self.layer_scale_a.unsqueeze(-1     ).unsqueeze(-1     ) * pooling_output
					# First residual connection
					UpperCAmelCase_     :      Optional[int]						=  hidden_states + self.drop_path(_snake_case     )
					UpperCAmelCase_     :      Optional[Any]						=  ()

					UpperCAmelCase_     :      List[str]						=  self.output(self.after_norm(_snake_case     )     )
					UpperCAmelCase_     :      Any						=  self.layer_scale_a.unsqueeze(-1     ).unsqueeze(-1     ) * layer_output
					# Second residual connection
					UpperCAmelCase_     :      List[str]						=  hidden_states + self.drop_path(_snake_case     )

					UpperCAmelCase_     :      Any						=  (output,) + outputs
					return outputs

				else:
					UpperCAmelCase_     :      Tuple						=  self.drop_path(self.pooling(self.before_norm(_snake_case     )     )     )
					# First residual connection
					UpperCAmelCase_     :      List[Any]						=  pooling_output + hidden_states
					UpperCAmelCase_     :      Optional[int]						=  ()

					# Second residual connection inside the PoolFormerOutput block
					UpperCAmelCase_     :      Optional[int]						=  self.drop_path(self.output(self.after_norm(_snake_case     )     )     )
					UpperCAmelCase_     :      List[Any]						=  hidden_states + layer_output

					UpperCAmelCase_     :      Optional[Any]						=  (output,) + outputs
					return outputs




class 					_snake_case		(nn.Module):






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

				# patch embeddings
				UpperCAmelCase_     :      str						=  []
				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]			,)     )
				UpperCAmelCase_     :      List[Any]						=  nn.ModuleList(_snake_case     )

				# Transformer blocks
				UpperCAmelCase_     :      int						=  []
				UpperCAmelCase_     :      List[Any]						=  0
				for i in range(config.num_encoder_blocks     ):
					# each block consists of layers
					UpperCAmelCase_     :      str						=  []
					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     )     )

				UpperCAmelCase_     :      List[Any]						=  nn.ModuleList(_snake_case     )






			def 		UpperCamelCase__							(  self			,_snake_case			,_snake_case=False			,_snake_case=True     ):
				UpperCAmelCase_     :      Tuple						=  () if output_hidden_states else None

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

					if output_hidden_states:
						UpperCAmelCase_     :      List[Any]						=  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 					_snake_case		(__SCREAMING_SNAKE_CASE):
			__A  :	List[str]        =PoolFormerConfig
			__A  :	Dict        ="poolformer"
			__A  :	Any        ="pixel_values"
			__A  :	Optional[Any]        =True






			def 		UpperCamelCase__							(  self			,_snake_case     ):
				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 		UpperCamelCase__							(  self			,_snake_case			,_snake_case=False     ):
				if isinstance(_snake_case			,_snake_case     ):
					UpperCAmelCase_     :      Any						=  value




_lowerCamelCase      					=     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.
"""

_lowerCamelCase      					=     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."			,			__SCREAMING_SNAKE_CASE			,			)
class 					_snake_case		(__SCREAMING_SNAKE_CASE):






			def __init__(  self			,_snake_case     ):
				super().__init__(_snake_case     )
				UpperCAmelCase_     :      str						=  config

				UpperCAmelCase_     :      str						=  PoolFormerEncoder(_snake_case     )

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






			def 		UpperCamelCase__							(  self     ):
				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 		UpperCamelCase__							(  self			,_snake_case = None			,_snake_case = None			,_snake_case = None			,):
				UpperCAmelCase_     :      int						=  (
				    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
				)
				UpperCAmelCase_     :      List[Any]						=  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"     )

				UpperCAmelCase_     :      str						=  self.encoder(
				    _snake_case			,output_hidden_states=_snake_case			,return_dict=_snake_case			,)
				UpperCAmelCase_     :      Tuple						=  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 					_snake_case		(nn.Module):






			def __init__(  self			,_snake_case     ):
				super().__init__()
				UpperCAmelCase_     :      List[Any]						=  nn.Linear(config.hidden_size			,config.hidden_size     )






			def 		UpperCamelCase__							(  self			,_snake_case     ):
				UpperCAmelCase_     :      Dict						=  self.dense(_snake_case     )
				return output




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






			def __init__(  self			,_snake_case     ):
				super().__init__(_snake_case     )
				UpperCAmelCase_     :      List[Any]						=  config.num_labels
				UpperCAmelCase_     :      Any						=  PoolFormerModel(_snake_case     )

				# Final norm
				UpperCAmelCase_     :      Optional[int]						=  PoolFormerGroupNorm(config.hidden_sizes[-1]     )
				# Classifier head
				UpperCAmelCase_     :      Optional[Any]						=  (
				    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 		UpperCamelCase__							(  self			,_snake_case = None			,_snake_case = None			,_snake_case = None			,_snake_case = None			,):
				UpperCAmelCase_     :      Optional[int]						=  return_dict if return_dict is not None else self.config.use_return_dict

				UpperCAmelCase_     :      Union[str, Any]						=  self.poolformer(
				    _snake_case			,output_hidden_states=_snake_case			,return_dict=_snake_case			,)

				UpperCAmelCase_     :      Any						=  outputs[0]

				UpperCAmelCase_     :      str						=  self.classifier(self.norm(_snake_case     ).mean([-2, -1]     )     )

				UpperCAmelCase_     :      Tuple						=  None
				if labels is not None:
					if self.config.problem_type is None:
						if self.num_labels == 1:
							UpperCAmelCase_     :      List[str]						=  "regression"
						elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
							UpperCAmelCase_     :      Optional[int]						=  "single_label_classification"
						else:
							UpperCAmelCase_     :      Tuple						=  "multi_label_classification"

					if self.config.problem_type == "regression":
						UpperCAmelCase_     :      str						=  MSELoss()
						if self.num_labels == 1:
							UpperCAmelCase_     :      int						=  loss_fct(logits.squeeze()			,labels.squeeze()     )
						else:
							UpperCAmelCase_     :      List[str]						=  loss_fct(_snake_case			,_snake_case     )
					elif self.config.problem_type == "single_label_classification":
						UpperCAmelCase_     :      Any						=  CrossEntropyLoss()
						UpperCAmelCase_     :      List[str]						=  loss_fct(logits.view(-1			,self.num_labels     )			,labels.view(-1     )     )
					elif self.config.problem_type == "multi_label_classification":
						UpperCAmelCase_     :      Optional[int]						=  BCEWithLogitsLoss()
						UpperCAmelCase_     :      Optional[int]						=  loss_fct(_snake_case			,_snake_case     )

				if not return_dict:
					UpperCAmelCase_     :      Any						=  (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     )








from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


UpperCAmelCase   :    Tuple     =		{
    "configuration_blenderbot_small": [
        "BLENDERBOT_SMALL_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "BlenderbotSmallConfig",
        "BlenderbotSmallOnnxConfig",
    ],
    "tokenization_blenderbot_small": ["BlenderbotSmallTokenizer"],
}

try:
		if not is_tokenizers_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		UpperCAmelCase   :    List[Any]     =		["BlenderbotSmallTokenizerFast"]

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		UpperCAmelCase   :    List[str]     =		[
		    "BLENDERBOT_SMALL_PRETRAINED_MODEL_ARCHIVE_LIST",
		    "BlenderbotSmallForCausalLM",
		    "BlenderbotSmallForConditionalGeneration",
		    "BlenderbotSmallModel",
		    "BlenderbotSmallPreTrainedModel",
		]

try:
		if not is_tf_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		UpperCAmelCase   :    Any     =		[
		    "TFBlenderbotSmallForConditionalGeneration",
		    "TFBlenderbotSmallModel",
		    "TFBlenderbotSmallPreTrainedModel",
		]

try:
		if not is_flax_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		UpperCAmelCase   :    List[Any]     =		[
		    "FlaxBlenderbotSmallForConditionalGeneration",
		    "FlaxBlenderbotSmallModel",
		    "FlaxBlenderbotSmallPreTrainedModel",
		]

if TYPE_CHECKING:
		from .configuration_blenderbot_small import (
		    BLENDERBOT_SMALL_PRETRAINED_CONFIG_ARCHIVE_MAP,
		    BlenderbotSmallConfig,
		    BlenderbotSmallOnnxConfig,
		)
		from .tokenization_blenderbot_small import BlenderbotSmallTokenizer

		try:
				if not is_tokenizers_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .tokenization_blenderbot_small_fast import BlenderbotSmallTokenizerFast

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_blenderbot_small import (
				    BLENDERBOT_SMALL_PRETRAINED_MODEL_ARCHIVE_LIST,
				    BlenderbotSmallForCausalLM,
				    BlenderbotSmallForConditionalGeneration,
				    BlenderbotSmallModel,
				    BlenderbotSmallPreTrainedModel,
				)

		try:
				if not is_tf_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_tf_blenderbot_small import (
				    TFBlenderbotSmallForConditionalGeneration,
				    TFBlenderbotSmallModel,
				    TFBlenderbotSmallPreTrainedModel,
				)

		try:
				if not is_flax_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_flax_blenderbot_small import (
				    FlaxBlenderbotSmallForConditionalGeneration,
				    FlaxBlenderbotSmallModel,
				    FlaxBlenderbotSmallPreTrainedModel,
				)

else:
		import sys

		UpperCAmelCase   :    List[Any]     =		_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)

