




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							:						List[Any]					       =  logging.get_logger(__name__)

# General docstring
lowerCamelCase							:						int					       =  "PoolFormerConfig"

# Base docstring
lowerCamelCase							:						Dict					       =  "sail/poolformer_s12"
lowerCamelCase							:						str					       =  [1, 512, 7, 7]

# Image classification docstring
lowerCamelCase							:						Union[str, Any]					       =  "sail/poolformer_s12"
lowerCamelCase							:						int					       =  "tabby, tabby cat"

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

def      _SCREAMING_SNAKE_CASE     (				lowercase      :			Union[str, Any]      ,				lowercase      :			float = 0.0      ,				lowercase      :			bool = False	):


      '''simple docstring'''



      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(lowercase      ,				dtype=input.dtype      ,				device=input.device	)
      random_tensor.floor_()  # binarize
      lowerCamelCase_							=       input.div(lowercase	) * random_tensor
      return output







class     A(		nn.Module  ):



 '''simple docstring'''


 def __init__(							self   :    Tuple		,				A_   :    Optional[float] = None    )  ->      None:

       """simple docstring"""






       super().__init__()
       lowerCamelCase_							=       drop_prob

 def 							a__					(							self   :    List[Any]		,				A_   :    torch.Tensor    )  ->      torch.Tensor:

       """simple docstring"""






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





 def 							a__					(							self   :    Tuple    )  ->      str:

       """simple docstring"""






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







class     A(		nn.Module  ):



 '''simple docstring'''


 def __init__(							self   :    List[str]		,				A_   :    Optional[Any]		,				A_   :    List[Any]		,				A_   :    int		,				A_   :    List[str]		,				A_   :    Any		,				A_   :    Tuple=None    )  ->      Optional[Any]:

       """simple docstring"""






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

       lowerCamelCase_							=       nn.Convad(A_		,				A_		,				kernel_size=A_		,				stride=A_		,				padding=A_    )
       lowerCamelCase_							=       norm_layer(A_    ) if norm_layer else nn.Identity()





 def 							a__					(							self   :    Tuple		,				A_   :    str    )  ->      List[str]:

       """simple docstring"""






       lowerCamelCase_							=       self.projection(A_    )
       lowerCamelCase_							=       self.norm(A_    )
       return embeddings







class     A(		nn.GroupNorm  ):



 '''simple docstring'''


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

       """simple docstring"""






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







class     A(		nn.Module  ):



 '''simple docstring'''


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

       """simple docstring"""






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





 def 							a__					(							self   :    List[str]		,				A_   :    Union[str, Any]    )  ->      Dict:

       """simple docstring"""






       return self.pool(A_    ) - hidden_states







class     A(		nn.Module  ):



 '''simple docstring'''


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

       """simple docstring"""






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





 def 							a__					(							self   :    List[str]		,				A_   :    int    )  ->      List[Any]:

       """simple docstring"""






       lowerCamelCase_							=       self.conva(A_    )
       lowerCamelCase_							=       self.act_fn(A_    )
       lowerCamelCase_							=       self.drop(A_    )
       lowerCamelCase_							=       self.conva(A_    )
       lowerCamelCase_							=       self.drop(A_    )

       return hidden_states







class     A(		nn.Module  ):



 '''simple docstring'''


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

       """simple docstring"""






       super().__init__()
       lowerCamelCase_							=       PoolFormerPooling(A_    )
       lowerCamelCase_							=       PoolFormerOutput(A_		,				A_		,				A_		,				A_    )
       lowerCamelCase_							=       PoolFormerGroupNorm(A_    )
       lowerCamelCase_							=       PoolFormerGroupNorm(A_    )

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





 def 							a__					(							self   :    Dict		,				A_   :    str    )  ->      Any:

       """simple docstring"""






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

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

             lowerCamelCase_							=       (output,) + outputs
             return outputs

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

             lowerCamelCase_							=       (output,) + outputs
             return outputs







class     A(		nn.Module  ):



 '''simple docstring'''


 def __init__(							self   :    Tuple		,				A_   :    List[str]    )  ->      Any:

       """simple docstring"""






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

       # 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(
                           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_    )    )

       lowerCamelCase_							=       nn.ModuleList(A_    )





 def 							a__					(							self   :    str		,				A_   :    Union[str, Any]		,				A_   :    Dict=False		,				A_   :    Union[str, Any]=True    )  ->      str:

       """simple docstring"""






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







class     A(		UpperCamelCase  ):



 '''simple docstring'''


 UpperCamelCase		  =		PoolFormerConfig
 UpperCamelCase		  =		'''poolformer'''
 UpperCamelCase		  =		'''pixel_values'''
 UpperCamelCase		  =		True

 def 							a__					(							self   :    Optional[Any]		,				A_   :    Any    )  ->      str:

       """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 							a__					(							self   :    Dict		,				A_   :    Union[str, Any]		,				A_   :    Dict=False    )  ->      int:

       """simple docstring"""






       if isinstance(A_		,				A_    ):
             lowerCamelCase_							=       value

lowerCamelCase							:						Any					       =  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							:						Dict					       =  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     A(		UpperCamelCase  ):



 '''simple docstring'''


 def __init__(							self   :    Dict		,				A_   :    Any    )  ->      Optional[Any]:

       """simple docstring"""






       super().__init__(A_    )
       lowerCamelCase_							=       config

       lowerCamelCase_							=       PoolFormerEncoder(A_    )

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

 def 							a__					(							self   :    Optional[Any]    )  ->      List[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 							a__					(							self   :    Union[str, Any]		,				A_   :    Optional[torch.FloatTensor] = None		,				A_   :    Optional[bool] = None		,				A_   :    Optional[bool] = None		,				)  ->      Union[Tuple, BaseModelOutputWithNoAttention]:

       """simple docstring"""






       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(
           A_		,				output_hidden_states=A_		,				return_dict=A_		,				)
       lowerCamelCase_							=       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     A(		nn.Module  ):



 '''simple docstring'''


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

       """simple docstring"""






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





 def 							a__					(							self   :    Any		,				A_   :    str    )  ->      Optional[Any]:

       """simple docstring"""






       lowerCamelCase_							=       self.dense(A_    )
       return output







@add_start_docstrings(
    '''
    PoolFormer Model transformer with an image classification head on top
    '''				, UpperCamelCase				, )
class     A(		UpperCamelCase  ):



 '''simple docstring'''


 def __init__(							self   :    Dict		,				A_   :    Dict    )  ->      Any:

       """simple docstring"""






       super().__init__(A_    )
       lowerCamelCase_							=       config.num_labels
       lowerCamelCase_							=       PoolFormerModel(A_    )

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

       """simple docstring"""






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

       lowerCamelCase_							=       self.poolformer(
           A_		,				output_hidden_states=A_		,				return_dict=A_		,				)

       lowerCamelCase_							=       outputs[0]

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

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

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





'''simple docstring'''
import enum
import shutil
import sys


lowercase__				, lowercase__											=shutil.get_terminal_size()

lowercase__											={'UP': 'A', 'DOWN': 'B', 'RIGHT': 'C', 'LEFT': 'D'}







class   a_ (      enum.Enum					):
			lowerCamelCase__	:			int       =    0
			lowerCamelCase__	:			Union[str, Any]       =    1







def      UpperCamelCase_				(      A__  ,      A__=""						):
	sys.stdout.write(str(A__						) + end						)
	sys.stdout.flush()






def      UpperCamelCase_				(      A__  ,      A__  ,      A__=""						):
	forceWrite(F'''\u001b[{color}m{content}\u001b[0m'''  ,      A__						)






def      UpperCamelCase_				(      ):
	forceWrite("""\r"""						)






def      UpperCamelCase_				(      A__  ,      A__						):
	forceWrite(F'''\033[{num_lines}{CURSOR_TO_CHAR[direction.upper()]}'''						)






def      UpperCamelCase_				(      ):
	forceWrite(""" """ * TERMINAL_WIDTH						)
	reset_cursor()






def      UpperCamelCase_				(      ):
	reset_cursor()
	forceWrite("""-""" * TERMINAL_WIDTH						)
