





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


_snake_case           =			logging.get_logger(__name__)

# General docstring
_snake_case           =			"""PoolFormerConfig"""

# Base docstring
_snake_case           =			"""sail/poolformer_s12"""
_snake_case           =			[1, 512, 7, 7]

# Image classification docstring
_snake_case           =			"""sail/poolformer_s12"""
_snake_case           =			"""tabby, tabby cat"""

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


def _A       (   __magic_name__  , __magic_name__ = 0.0  , __magic_name__ = False	):
   if drop_prob == 0.0 or not training:
      return input
   lowercase__			       =  1 - drop_prob
   lowercase__			       =  (input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
   lowercase__			       =  keep_prob + torch.rand(__magic_name__  , dtype=input.dtype  , device=input.device	)
   random_tensor.floor_()  # binarize
   lowercase__			       =  input.div(__magic_name__	) * random_tensor
   return output





class  lowerCAmelCase							(       nn.Module ):







  def __init__(      self  :Union[str, Any]   ,							_lowercase  :Optional[float] = None ):






     '''simple docstring'''





     super().__init__()
     lowercase__			       =  drop_prob







  def   UpperCAmelCase					(      self  :Optional[int]   ,							_lowercase  :torch.Tensor ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :Dict ):






     '''simple docstring'''





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





class  lowerCAmelCase							(       nn.Module ):







  def __init__(      self  :Tuple   ,							_lowercase  :Any   ,							_lowercase  :List[Any]   ,							_lowercase  :Dict   ,							_lowercase  :List[str]   ,							_lowercase  :Union[str, Any]   ,							_lowercase  :str=None ):






     '''simple docstring'''





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

     lowercase__			       =  nn.Convad(_lowercase   ,							_lowercase   ,							kernel_size=_lowercase   ,							stride=_lowercase   ,							padding=_lowercase )
     lowercase__			       =  norm_layer(_lowercase ) if norm_layer else nn.Identity()







  def   UpperCAmelCase					(      self  :Optional[int]   ,							_lowercase  :Any ):






     '''simple docstring'''





     lowercase__			       =  self.projection(_lowercase )
     lowercase__			       =  self.norm(_lowercase )
     return embeddings





class  lowerCAmelCase							(       nn.GroupNorm ):







  def __init__(      self  :Union[str, Any]   ,							_lowercase  :Dict   ,							**_lowercase  :int ):






     '''simple docstring'''





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





class  lowerCAmelCase							(       nn.Module ):







  def __init__(      self  :Any   ,							_lowercase  :int ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :Tuple   ,							_lowercase  :Tuple ):






     '''simple docstring'''





     return self.pool(_lowercase ) - hidden_states





class  lowerCAmelCase							(       nn.Module ):







  def __init__(      self  :Any   ,							_lowercase  :Tuple   ,							_lowercase  :Optional[Any]   ,							_lowercase  :Dict   ,							_lowercase  :int ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :List[Any]   ,							_lowercase  :Tuple ):






     '''simple docstring'''





     lowercase__			       =  self.conva(_lowercase )
     lowercase__			       =  self.act_fn(_lowercase )
     lowercase__			       =  self.drop(_lowercase )
     lowercase__			       =  self.conva(_lowercase )
     lowercase__			       =  self.drop(_lowercase )

     return hidden_states





class  lowerCAmelCase							(       nn.Module ):







  def __init__(      self  :List[str]   ,							_lowercase  :Tuple   ,							_lowercase  :Optional[Any]   ,							_lowercase  :Any   ,							_lowercase  :Tuple   ,							_lowercase  :List[Any]   ,							_lowercase  :List[Any] ):






     '''simple docstring'''





     super().__init__()
     lowercase__			       =  PoolFormerPooling(_lowercase )
     lowercase__			       =  PoolFormerOutput(_lowercase   ,							_lowercase   ,							_lowercase   ,							_lowercase )
     lowercase__			       =  PoolFormerGroupNorm(_lowercase )
     lowercase__			       =  PoolFormerGroupNorm(_lowercase )

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







  def   UpperCAmelCase					(      self  :Union[str, Any]   ,							_lowercase  :List[Any] ):






     '''simple docstring'''





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

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

        lowercase__			       =  (output,) + outputs
        return outputs

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

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

        lowercase__			       =  (output,) + outputs
        return outputs





class  lowerCAmelCase							(       nn.Module ):







  def __init__(      self  :str   ,							_lowercase  :Optional[Any] ):






     '''simple docstring'''





     super().__init__()
     lowercase__			       =  config
     # stochastic depth decay rule
     lowercase__			       =  [x.item() for x in torch.linspace(0   ,							config.drop_path_rate   ,							sum(config.depths ) )]

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

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

     lowercase__			       =  nn.ModuleList(_lowercase )







  def   UpperCAmelCase					(      self  :Optional[int]   ,							_lowercase  :Tuple   ,							_lowercase  :str=False   ,							_lowercase  :Any=True ):






     '''simple docstring'''





     lowercase__			       =  () if output_hidden_states else None

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

        if output_hidden_states:
           lowercase__			       =  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=_lowercase   ,							hidden_states=_lowercase )





class  lowerCAmelCase							(       lowercase_ ):
  __lowerCamelCase		=      PoolFormerConfig
  __lowerCamelCase		=      'poolformer'
  __lowerCamelCase		=      'pixel_values'
  __lowerCamelCase		=      True







  def   UpperCAmelCase					(      self  :Union[str, Any]   ,							_lowercase  :Tuple ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :Union[str, Any]   ,							_lowercase  :Dict   ,							_lowercase  :Optional[Any]=False ):






     '''simple docstring'''





     if isinstance(_lowercase   ,							_lowercase ):
        lowercase__			       =  value
_snake_case           =			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.
"""

_snake_case           =			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.'							,							lowercase_							,							)
class  lowerCAmelCase							(       lowercase_ ):







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






     '''simple docstring'''





     super().__init__(_lowercase )
     lowercase__			       =  config

     lowercase__			       =  PoolFormerEncoder(_lowercase )

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







  def   UpperCAmelCase					(      self  :Dict ):






     '''simple docstring'''





     return self.embeddings.patch_embeddings







  @add_start_docstrings_to_model_forward(_lowercase )
  @add_code_sample_docstrings(
      checkpoint=_CHECKPOINT_FOR_DOC   ,							output_type=_lowercase   ,							config_class=_CONFIG_FOR_DOC   ,							modality="vision"   ,							expected_output=_EXPECTED_OUTPUT_SHAPE   ,							)
  def   UpperCAmelCase					(      self  :Union[str, Any]   ,							_lowercase  :Optional[torch.FloatTensor] = None   ,							_lowercase  :Optional[bool] = None   ,							_lowercase  :Optional[bool] = None   ,							):






     '''simple docstring'''





     lowercase__			       =  (
         output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
     )
     lowercase__			       =  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" )

     lowercase__			       =  self.encoder(
         _lowercase   ,							output_hidden_states=_lowercase   ,							return_dict=_lowercase   ,							)
     lowercase__			       =  encoder_outputs[0]

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

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





class  lowerCAmelCase							(       nn.Module ):







  def __init__(      self  :int   ,							_lowercase  :Optional[int] ):






     '''simple docstring'''





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







  def   UpperCAmelCase					(      self  :int   ,							_lowercase  :Optional[int] ):






     '''simple docstring'''





     lowercase__			       =  self.dense(_lowercase )
     return output





@add_start_docstrings(
    '\n    PoolFormer Model transformer with an image classification head on top\n    '							,							lowercase_							,							)
class  lowerCAmelCase							(       lowercase_ ):







  def __init__(      self  :List[str]   ,							_lowercase  :Optional[Any] ):






     '''simple docstring'''





     super().__init__(_lowercase )
     lowercase__			       =  config.num_labels
     lowercase__			       =  PoolFormerModel(_lowercase )

     # Final norm
     lowercase__			       =  PoolFormerGroupNorm(config.hidden_sizes[-1] )
     # Classifier head
     lowercase__			       =  (
         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(_lowercase )
  @add_code_sample_docstrings(
      checkpoint=_IMAGE_CLASS_CHECKPOINT   ,							output_type=_lowercase   ,							config_class=_CONFIG_FOR_DOC   ,							expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT   ,							)
  def   UpperCAmelCase					(      self  :Optional[int]   ,							_lowercase  :Optional[torch.FloatTensor] = None   ,							_lowercase  :Optional[torch.LongTensor] = None   ,							_lowercase  :Optional[bool] = None   ,							_lowercase  :Optional[bool] = None   ,							):






     '''simple docstring'''





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

     lowercase__			       =  self.poolformer(
         _lowercase   ,							output_hidden_states=_lowercase   ,							return_dict=_lowercase   ,							)

     lowercase__			       =  outputs[0]

     lowercase__			       =  self.classifier(self.norm(_lowercase ).mean([-2, -1] ) )

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

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

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

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













from abc import ABC, abstractmethod
from typing import Optional, Union

from .. import Dataset, DatasetDict, Features, IterableDataset, IterableDatasetDict, NamedSplit
from ..utils.typing import NestedDataStructureLike, PathLike





class  lowerCAmelCase							(       lowercase_ ):







  def __init__(      self  :str   ,							_lowercase  :Optional[NestedDataStructureLike[PathLike]] = None   ,							_lowercase  :Optional[NamedSplit] = None   ,							_lowercase  :Optional[Features] = None   ,							_lowercase  :str = None   ,							_lowercase  :bool = False   ,							_lowercase  :bool = False   ,							_lowercase  :Optional[int] = None   ,							**_lowercase  :Tuple   ,							):






     '''simple docstring'''





     lowercase__			       =  path_or_paths
     lowercase__			       =  split if split or isinstance(_lowercase   ,							_lowercase ) else "train"
     lowercase__			       =  features
     lowercase__			       =  cache_dir
     lowercase__			       =  keep_in_memory
     lowercase__			       =  streaming
     lowercase__			       =  num_proc
     lowercase__			       =  kwargs







  @abstractmethod
  def   UpperCAmelCase					(      self  :Any ):






     '''simple docstring'''





     pass





class  lowerCAmelCase							(       lowercase_ ):







  def __init__(      self  :List[Any]   ,							_lowercase  :Optional[Features] = None   ,							_lowercase  :str = None   ,							_lowercase  :bool = False   ,							_lowercase  :bool = False   ,							_lowercase  :Optional[int] = None   ,							**_lowercase  :Optional[int]   ,							):






     '''simple docstring'''





     lowercase__			       =  features
     lowercase__			       =  cache_dir
     lowercase__			       =  keep_in_memory
     lowercase__			       =  streaming
     lowercase__			       =  num_proc
     lowercase__			       =  kwargs







  @abstractmethod
  def   UpperCAmelCase					(      self  :int ):






     '''simple docstring'''





     pass






