



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


lowercase__		:				Union[str, Any]          =	logging.get_logger(__name__)

# General docstring
lowercase__		:				Any          =	"""PoolFormerConfig"""

# Base docstring
lowercase__		:				Optional[Any]          =	"""sail/poolformer_s12"""
lowercase__		:				int          =	[1, 5_1_2, 7, 7]

# Image classification docstring
lowercase__		:				Dict          =	"""sail/poolformer_s12"""
lowercase__		:				str          =	"""tabby, tabby cat"""

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

def 				UpperCamelCase_    (    lowerCAmelCase__      :				Optional[int]    ,   lowerCAmelCase__      :				Optional[Any] = 0.0    ,   lowerCAmelCase__      :				Tuple = False						)						->   Optional[Any]:

  """simple docstring"""




  if drop_prob == 0.0 or not training:
    return input
  lowerCAmelCase_   :				int								= 1 - drop_prob
  lowerCAmelCase_   :				List[str]								= (input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
  lowerCAmelCase_   :				Union[str, Any]								= keep_prob + torch.rand(__lowerCamelCase    ,   dtype=input.dtype    ,   device=input.device						)
  random_tensor.floor_()  # binarize
  lowerCAmelCase_   :				int								= input.div(__lowerCamelCase						) * random_tensor
  return output







class       UpperCamelCase__   (	nn.Module  ):




      """simple docstring"""




      def __init__(      self				:      str    ,			SCREAMING_SNAKE_CASE_				:      Optional[float] = None       ):
        super().__init__()
        lowerCAmelCase_   :				Dict								= drop_prob



      def      SCREAMING_SNAKE_CASE__ (      self				:      int    ,			SCREAMING_SNAKE_CASE_				:      torch.Tensor       ):
        return drop_path(_snake_case    ,			self.drop_prob    ,			self.training       )





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







class       UpperCamelCase__   (	nn.Module  ):




      """simple docstring"""




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

        lowerCAmelCase_   :				int								= nn.Convad(_snake_case    ,			_snake_case    ,			kernel_size=_snake_case    ,			stride=_snake_case    ,			padding=_snake_case       )
        lowerCAmelCase_   :				int								= norm_layer(_snake_case       ) if norm_layer else nn.Identity()





      def      SCREAMING_SNAKE_CASE__ (      self				:      Any    ,			SCREAMING_SNAKE_CASE_				:      int       ):
        lowerCAmelCase_   :				List[Any]								= self.projection(_snake_case       )
        lowerCAmelCase_   :				Tuple								= self.norm(_snake_case       )
        return embeddings







class       UpperCamelCase__   (	nn.GroupNorm  ):




      """simple docstring"""




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







class       UpperCamelCase__   (	nn.Module  ):




      """simple docstring"""




      def __init__(      self				:      Optional[Any]    ,			SCREAMING_SNAKE_CASE_				:      Optional[Any]       ):
        super().__init__()
        lowerCAmelCase_   :				Union[str, Any]								= nn.AvgPoolad(_snake_case    ,			stride=1    ,			padding=pool_size // 2    ,			count_include_pad=_snake_case       )





      def      SCREAMING_SNAKE_CASE__ (      self				:      Optional[int]    ,			SCREAMING_SNAKE_CASE_				:      Optional[int]       ):
        return self.pool(_snake_case       ) - hidden_states







class       UpperCamelCase__   (	nn.Module  ):




      """simple docstring"""




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





      def      SCREAMING_SNAKE_CASE__ (      self				:      Dict    ,			SCREAMING_SNAKE_CASE_				:      Optional[int]       ):
        lowerCAmelCase_   :				Optional[int]								= self.conva(_snake_case       )
        lowerCAmelCase_   :				Any								= self.act_fn(_snake_case       )
        lowerCAmelCase_   :				Union[str, Any]								= self.drop(_snake_case       )
        lowerCAmelCase_   :				int								= self.conva(_snake_case       )
        lowerCAmelCase_   :				str								= self.drop(_snake_case       )

        return hidden_states







class       UpperCamelCase__   (	nn.Module  ):




      """simple docstring"""




      def __init__(      self				:      Any    ,			SCREAMING_SNAKE_CASE_				:      List[str]    ,			SCREAMING_SNAKE_CASE_				:      str    ,			SCREAMING_SNAKE_CASE_				:      Any    ,			SCREAMING_SNAKE_CASE_				:      Dict    ,			SCREAMING_SNAKE_CASE_				:      List[str]    ,			SCREAMING_SNAKE_CASE_				:      Union[str, Any]       ):
        super().__init__()
        lowerCAmelCase_   :				List[Any]								= PoolFormerPooling(_snake_case       )
        lowerCAmelCase_   :				int								= PoolFormerOutput(_snake_case    ,			_snake_case    ,			_snake_case    ,			_snake_case       )
        lowerCAmelCase_   :				str								= PoolFormerGroupNorm(_snake_case       )
        lowerCAmelCase_   :				Optional[Any]								= PoolFormerGroupNorm(_snake_case       )

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





      def      SCREAMING_SNAKE_CASE__ (      self				:      Union[str, Any]    ,			SCREAMING_SNAKE_CASE_				:      List[str]       ):



        if self.use_layer_scale:
          lowerCAmelCase_   :				List[str]								= self.pooling(self.before_norm(_snake_case       )       )
          lowerCAmelCase_   :				Tuple								= self.layer_scale_a.unsqueeze(-1       ).unsqueeze(-1       ) * pooling_output
          # First residual connection
          lowerCAmelCase_   :				Any								= hidden_states + self.drop_path(_snake_case       )
          lowerCAmelCase_   :				int								= ()

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

          lowerCAmelCase_   :				Optional[int]								= (output,) + outputs
          return outputs

        else:
          lowerCAmelCase_   :				Any								= self.drop_path(self.pooling(self.before_norm(_snake_case       )       )       )
          # First residual connection
          lowerCAmelCase_   :				Dict								= pooling_output + hidden_states
          lowerCAmelCase_   :				Tuple								= ()

          # Second residual connection inside the PoolFormerOutput block
          lowerCAmelCase_   :				Union[str, Any]								= self.drop_path(self.output(self.after_norm(_snake_case       )       )       )
          lowerCAmelCase_   :				Optional[Any]								= hidden_states + layer_output

          lowerCAmelCase_   :				Dict								= (output,) + outputs
          return outputs







class       UpperCamelCase__   (	nn.Module  ):




      """simple docstring"""




      def __init__(      self				:      List[Any]    ,			SCREAMING_SNAKE_CASE_				:      List[Any]       ):
        super().__init__()
        lowerCAmelCase_   :				Union[str, Any]								= config
        # stochastic depth decay rule
        lowerCAmelCase_   :				List[Any]								= [x.item() for x in torch.linspace(0    ,			config.drop_path_rate    ,			sum(config.depths       )       )]

        # patch embeddings
        lowerCAmelCase_   :				Any								= []
        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_   :				Optional[Any]								= nn.ModuleList(_snake_case       )

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





      def      SCREAMING_SNAKE_CASE__ (      self				:      str    ,			SCREAMING_SNAKE_CASE_				:      Tuple    ,			SCREAMING_SNAKE_CASE_				:      List[Any]=False    ,			SCREAMING_SNAKE_CASE_				:      Union[str, Any]=True       ):
        lowerCAmelCase_   :				List[str]								= () if output_hidden_states else None

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

          if output_hidden_states:
            lowerCAmelCase_   :				Tuple								= 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       UpperCamelCase__   (	A_  ):




      """simple docstring"""




      _SCREAMING_SNAKE_CASE  				=     PoolFormerConfig
      _SCREAMING_SNAKE_CASE  				=     "poolformer"
      _SCREAMING_SNAKE_CASE  				=     "pixel_values"
      _SCREAMING_SNAKE_CASE  				=     True



      def      SCREAMING_SNAKE_CASE__ (      self				:      Union[str, Any]    ,			SCREAMING_SNAKE_CASE_				:      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      SCREAMING_SNAKE_CASE__ (      self				:      Optional[int]    ,			SCREAMING_SNAKE_CASE_				:      Optional[int]    ,			SCREAMING_SNAKE_CASE_				:      Optional[Any]=False       ):



        if isinstance(_snake_case    ,			_snake_case       ):
          lowerCAmelCase_   :				Optional[int]								= value




lowercase__		:				str          =	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"""

lowercase__		:				Tuple          =	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.""",					A_,					)
class       UpperCamelCase__   (	A_  ):




      """simple docstring"""




      def __init__(      self				:      Union[str, Any]    ,			SCREAMING_SNAKE_CASE_				:      Any       ):
        super().__init__(_snake_case       )
        lowerCAmelCase_   :				List[Any]								= config

        lowerCAmelCase_   :				Optional[int]								= PoolFormerEncoder(_snake_case       )

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



      def      SCREAMING_SNAKE_CASE__ (      self				:      str       ):
        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      SCREAMING_SNAKE_CASE__ (      self				:      List[str]    ,			SCREAMING_SNAKE_CASE_				:      Optional[torch.FloatTensor] = None    ,			SCREAMING_SNAKE_CASE_				:      Optional[bool] = None    ,			SCREAMING_SNAKE_CASE_				:      Optional[bool] = None    ,			):
        lowerCAmelCase_   :				Optional[int]								= (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        lowerCAmelCase_   :				Optional[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'       )

        lowerCAmelCase_   :				Tuple								= self.encoder(
            _snake_case    ,			output_hidden_states=_snake_case    ,			return_dict=_snake_case    ,			)
        lowerCAmelCase_   :				Optional[Any]								= 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       UpperCamelCase__   (	nn.Module  ):




      """simple docstring"""




      def __init__(      self				:      int    ,			SCREAMING_SNAKE_CASE_				:      Any       ):
        super().__init__()
        lowerCAmelCase_   :				Dict								= nn.Linear(config.hidden_size    ,			config.hidden_size       )





      def      SCREAMING_SNAKE_CASE__ (      self				:      str    ,			SCREAMING_SNAKE_CASE_				:      Dict       ):
        lowerCAmelCase_   :				List[Any]								= self.dense(_snake_case       )
        return output







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




      """simple docstring"""




      def __init__(      self				:      Tuple    ,			SCREAMING_SNAKE_CASE_				:      int       ):
        super().__init__(_snake_case       )
        lowerCAmelCase_   :				Optional[Any]								= config.num_labels
        lowerCAmelCase_   :				Optional[Any]								= PoolFormerModel(_snake_case       )

        # Final norm
        lowerCAmelCase_   :				Tuple								= PoolFormerGroupNorm(config.hidden_sizes[-1]       )
        # Classifier head
        lowerCAmelCase_   :				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      SCREAMING_SNAKE_CASE__ (      self				:      List[Any]    ,			SCREAMING_SNAKE_CASE_				:      Optional[torch.FloatTensor] = None    ,			SCREAMING_SNAKE_CASE_				:      Optional[torch.LongTensor] = None    ,			SCREAMING_SNAKE_CASE_				:      Optional[bool] = None    ,			SCREAMING_SNAKE_CASE_				:      Optional[bool] = None    ,			):
        lowerCAmelCase_   :				Optional[Any]								= return_dict if return_dict is not None else self.config.use_return_dict

        lowerCAmelCase_   :				Dict								= self.poolformer(
            _snake_case    ,			output_hidden_states=_snake_case    ,			return_dict=_snake_case    ,			)

        lowerCAmelCase_   :				Dict								= outputs[0]

        lowerCAmelCase_   :				Optional[int]								= self.classifier(self.norm(_snake_case       ).mean([-2, -1]       )       )

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

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

        if not return_dict:
          lowerCAmelCase_   :				List[str]								= (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"""







import unittest

from accelerate import debug_launcher
from accelerate.test_utils import require_cpu, test_ops, test_script
@require_cpu
class        __A       (	unittest.TestCase    ):







		'''simple docstring'''





		def 					UpperCAmelCase    (	self						:  Optional[int]						)					->		str:






				"""simple docstring"""



				debug_launcher(test_script.main						)







		def 					UpperCAmelCase    (	self						:  Dict						)					->		Tuple:






				"""simple docstring"""



				debug_launcher(test_ops.main						)
