


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__							:						Dict							     =   logging.get_logger(__name__)

# General docstring
lowerCamelCase__							:						Union[str, Any]							     =   """PoolFormerConfig"""

# Base docstring
lowerCamelCase__							:						List[str]							     =   """sail/poolformer_s12"""
lowerCamelCase__							:						Optional[Any]							     =   [1, 5_1_2, 7, 7]

# Image classification docstring
lowerCamelCase__							:						Tuple							     =   """sail/poolformer_s12"""
lowerCamelCase__							:						Optional[int]							     =   """tabby, tabby cat"""

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







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





     '''simple docstring'''



     def __init__(					self:Tuple     ,		_a:Optional[float] = None   ):
       super().__init__()
       snake_case__						= drop_prob

     def 	SCREAMING_SNAKE_CASE__						(					self:List[Any]     ,		_a:torch.Tensor   ):
       return drop_path(_a     ,		self.drop_prob     ,		self.training   )







     def 	SCREAMING_SNAKE_CASE__						(					self:Optional[int]   ):
       return "p={}".format(self.drop_prob   )




class 		__magic_name__						(nn.Module				):





     '''simple docstring'''



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

       snake_case__						= nn.Convad(_a     ,		_a     ,		kernel_size=_a     ,		stride=_a     ,		padding=_a   )
       snake_case__						= norm_layer(_a   ) if norm_layer else nn.Identity()







     def 	SCREAMING_SNAKE_CASE__						(					self:Tuple     ,		_a:Optional[Any]   ):
       snake_case__						= self.projection(_a   )
       snake_case__						= self.norm(_a   )
       return embeddings




class 		__magic_name__						(nn.GroupNorm				):





     '''simple docstring'''



     def __init__(					self:Dict     ,		_a:Tuple     ,		**_a:int   ):
       super().__init__(1     ,		_a     ,		**_a   )




class 		__magic_name__						(nn.Module				):





     '''simple docstring'''



     def __init__(					self:Union[str, Any]     ,		_a:List[Any]   ):
       super().__init__()
       snake_case__						= nn.AvgPoolad(_a     ,		stride=1     ,		padding=pool_size // 2     ,		count_include_pad=_a   )







     def 	SCREAMING_SNAKE_CASE__						(					self:Any     ,		_a:int   ):
       return self.pool(_a   ) - hidden_states




class 		__magic_name__						(nn.Module				):





     '''simple docstring'''



     def __init__(					self:str     ,		_a:Optional[Any]     ,		_a:Optional[Any]     ,		_a:str     ,		_a:List[Any]   ):
       super().__init__()
       snake_case__						= nn.Convad(_a     ,		_a     ,		1   )
       snake_case__						= nn.Convad(_a     ,		_a     ,		1   )
       snake_case__						= PoolFormerDropPath(_a   )
       if isinstance(config.hidden_act     ,		_a   ):
         snake_case__						= ACTaFN[config.hidden_act]
       else:
         snake_case__						= config.hidden_act







     def 	SCREAMING_SNAKE_CASE__						(					self:int     ,		_a:Any   ):
       snake_case__						= self.conva(_a   )
       snake_case__						= self.act_fn(_a   )
       snake_case__						= self.drop(_a   )
       snake_case__						= self.conva(_a   )
       snake_case__						= self.drop(_a   )

       return hidden_states




class 		__magic_name__						(nn.Module				):





     '''simple docstring'''



     def __init__(					self:Any     ,		_a:Union[str, Any]     ,		_a:List[Any]     ,		_a:int     ,		_a:str     ,		_a:Dict     ,		_a:List[Any]   ):
       super().__init__()
       snake_case__						= PoolFormerPooling(_a   )
       snake_case__						= PoolFormerOutput(_a     ,		_a     ,		_a     ,		_a   )
       snake_case__						= PoolFormerGroupNorm(_a   )
       snake_case__						= PoolFormerGroupNorm(_a   )

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







     def 	SCREAMING_SNAKE_CASE__						(					self:List[Any]     ,		_a:Optional[Any]   ):
       if self.use_layer_scale:
         snake_case__						= self.pooling(self.before_norm(_a   )   )
         snake_case__						= self.layer_scale_a.unsqueeze(-1   ).unsqueeze(-1   ) * pooling_output
         # First residual connection
         snake_case__						= hidden_states + self.drop_path(_a   )
         snake_case__						= ()

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

         snake_case__						= (output,) + outputs
         return outputs

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

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

         snake_case__						= (output,) + outputs
         return outputs




class 		__magic_name__						(nn.Module				):





     '''simple docstring'''



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

       # patch embeddings
       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]     ,		)   )
       snake_case__						= nn.ModuleList(_a   )

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

       snake_case__						= nn.ModuleList(_a   )







     def 	SCREAMING_SNAKE_CASE__						(					self:List[Any]     ,		_a:Any     ,		_a:List[str]=False     ,		_a:Optional[Any]=True   ):
       snake_case__						= () if output_hidden_states else None

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

         if output_hidden_states:
           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 		__magic_name__						(snake_case_				):





     '''simple docstring'''



     __lowercase				:							Tuple	   =       PoolFormerConfig
     __lowercase				:							str	   =       'poolformer'
     __lowercase				:							Optional[int]	   =       'pixel_values'
     __lowercase				:							int	   =       True

     def 	SCREAMING_SNAKE_CASE__						(					self:Union[str, Any]     ,		_a:Tuple   ):
       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 	SCREAMING_SNAKE_CASE__						(					self:List[Any]     ,		_a:Optional[Any]     ,		_a:int=False   ):
       if isinstance(_a     ,		_a   ):
         snake_case__						= value




lowerCamelCase__							:						Optional[int]							     =   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__							:						List[str]							     =   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.'   ,snake_case_   ,)
class 		__magic_name__						(snake_case_				):





     '''simple docstring'''



     def __init__(					self:Tuple     ,		_a:List[str]   ):
       super().__init__(_a   )
       snake_case__						= config

       snake_case__						= PoolFormerEncoder(_a   )

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

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

       snake_case__						= self.encoder(
           _a     ,		output_hidden_states=_a     ,		return_dict=_a     ,		)
       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 		__magic_name__						(nn.Module				):





     '''simple docstring'''



     def __init__(					self:Optional[Any]     ,		_a:List[str]   ):
       super().__init__()
       snake_case__						= nn.Linear(config.hidden_size     ,		config.hidden_size   )







     def 	SCREAMING_SNAKE_CASE__						(					self:List[str]     ,		_a:Tuple   ):
       snake_case__						= self.dense(_a   )
       return output




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





     '''simple docstring'''



     def __init__(					self:Tuple     ,		_a:Dict   ):
       super().__init__(_a   )
       snake_case__						= config.num_labels
       snake_case__						= PoolFormerModel(_a   )

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

       snake_case__						= self.poolformer(
           _a     ,		output_hidden_states=_a     ,		return_dict=_a     ,		)

       snake_case__						= outputs[0]

       snake_case__						= self.classifier(self.norm(_a   ).mean([-2, -1]   )   )

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

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

       if not return_dict:
         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   )





from collections import OrderedDict
from typing import Mapping

from packaging import version

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import logging


lowerCamelCase__							:						List[Any]							     =   logging.get_logger(__name__)

lowerCamelCase__							:						Optional[int]							     =   {
    """facebook/data2vec-vision-base-ft""": (
        """https://huggingface.co/facebook/data2vec-vision-base-ft/resolve/main/config.json"""
    ),
}
class 		__magic_name__						(snake_case_				):





     '''simple docstring'''



     __lowercase				:							Optional[int]	   =       'data2vec-vision'

     def __init__(					self:int     ,		_a:Tuple=7_68     ,		_a:int=12     ,		_a:Any=12     ,		_a:Optional[int]=30_72     ,		_a:Optional[int]="gelu"     ,		_a:Any=0.0     ,		_a:Any=0.0     ,		_a:List[str]=0.02     ,		_a:Dict=1e-12     ,		_a:Tuple=2_24     ,		_a:Any=16     ,		_a:str=3     ,		_a:str=False     ,		_a:Union[str, Any]=False     ,		_a:Optional[int]=False     ,		_a:Any=False     ,		_a:Dict=0.1     ,		_a:Dict=0.1     ,		_a:str=True     ,		_a:str=[3, 5, 7, 11]     ,		_a:List[str]=[1, 2, 3, 6]     ,		_a:List[str]=True     ,		_a:Any=0.4     ,		_a:str=2_56     ,		_a:Union[str, Any]=1     ,		_a:int=False     ,		_a:Optional[int]=2_55     ,		**_a:Dict     ,		):
       super().__init__(**_a   )

       snake_case__						= hidden_size
       snake_case__						= num_hidden_layers
       snake_case__						= num_attention_heads
       snake_case__						= intermediate_size
       snake_case__						= hidden_act
       snake_case__						= hidden_dropout_prob
       snake_case__						= attention_probs_dropout_prob
       snake_case__						= initializer_range
       snake_case__						= layer_norm_eps

       snake_case__						= image_size
       snake_case__						= patch_size
       snake_case__						= num_channels
       snake_case__						= use_mask_token
       snake_case__						= use_absolute_position_embeddings
       snake_case__						= use_relative_position_bias
       snake_case__						= use_shared_relative_position_bias
       snake_case__						= layer_scale_init_value
       snake_case__						= drop_path_rate
       snake_case__						= use_mean_pooling
       # decode head attributes (semantic segmentation)
       snake_case__						= out_indices
       snake_case__						= pool_scales
       # auxiliary head attributes (semantic segmentation)
       snake_case__						= use_auxiliary_head
       snake_case__						= auxiliary_loss_weight
       snake_case__						= auxiliary_channels
       snake_case__						= auxiliary_num_convs
       snake_case__						= auxiliary_concat_input
       snake_case__						= semantic_loss_ignore_index




class 		__magic_name__						(snake_case_				):





     '''simple docstring'''



     __lowercase				:							Any	   =       version.parse('1.11'				)

     @property
     def 	SCREAMING_SNAKE_CASE__						(					self:List[str]   ):
       return OrderedDict(
           [
               ('''pixel_values''', {0: '''batch''', 1: '''num_channels''', 2: '''height''', 3: '''width'''}),
           ]   )







     @property
     def 	SCREAMING_SNAKE_CASE__						(					self:Tuple   ):
       return 1e-4

