






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

# General docstring
_lowerCAmelCase	:	List[Any]				=      """PoolFormerConfig"""

# Base docstring
_lowerCAmelCase	:	List[str]				=      """sail/poolformer_s12"""
_lowerCAmelCase	:	Optional[int]				=      [1, 512, 7, 7]

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

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


def 	SCREAMING_SNAKE_CASE__       (  snake_case       : str			,					snake_case       : float = 0.0			,					snake_case       : bool = False				)->		Dict:






     '''simple docstring'''



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




class 			lowerCAmelCase__			(    nn.Module   ):
   def __init__(		self      :	List[str]							,			snake_case__      :	Optional[float] = None							):


        '''simple docstring'''
        super().__init__()
        UpperCAmelCase__						:    List[str]							=							drop_prob
   def       __a							(		self      :	Any							,			snake_case__      :	torch.Tensor							):


        '''simple docstring'''
        return drop_path(snake_case__							,			self.drop_prob							,			self.training							)




   def       __a							(		self      :	Optional[Any]							):


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




class 			lowerCAmelCase__			(    nn.Module   ):
   def __init__(		self      :	Union[str, Any]							,			snake_case__      :	List[Any]							,			snake_case__      :	Optional[int]							,			snake_case__      :	Optional[Any]							,			snake_case__      :	str							,			snake_case__      :	Optional[int]							,			snake_case__      :	Optional[Any]=None							):


        '''simple docstring'''
        super().__init__()
        UpperCAmelCase__						:    Tuple							=							patch_size if isinstance(snake_case__							,			collections.abc.Iterable							) else (patch_size, patch_size)
        UpperCAmelCase__						:    str							=							stride if isinstance(snake_case__							,			collections.abc.Iterable							) else (stride, stride)
        UpperCAmelCase__						:    Optional[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__						:    Dict							=							norm_layer(snake_case__							) if norm_layer else nn.Identity()




   def       __a							(		self      :	int							,			snake_case__      :	List[str]							):


        '''simple docstring'''
        UpperCAmelCase__						:    Optional[Any]							=							self.projection(snake_case__							)
        UpperCAmelCase__						:    Optional[Any]							=							self.norm(snake_case__							)
        return embeddings




class 			lowerCAmelCase__			(    nn.GroupNorm   ):
   def __init__(		self      :	Union[str, Any]							,			snake_case__      :	List[str]							,			**snake_case__      :	Union[str, Any]							):


        '''simple docstring'''
        super().__init__(1							,			snake_case__							,			**snake_case__							)




class 			lowerCAmelCase__			(    nn.Module   ):
   def __init__(		self      :	Dict							,			snake_case__      :	List[str]							):


        '''simple docstring'''
        super().__init__()
        UpperCAmelCase__						:    int							=							nn.AvgPoolad(snake_case__							,			stride=1							,			padding=pool_size // 2							,			count_include_pad=snake_case__							)




   def       __a							(		self      :	int							,			snake_case__      :	str							):


        '''simple docstring'''
        return self.pool(snake_case__							) - hidden_states




class 			lowerCAmelCase__			(    nn.Module   ):
   def __init__(		self      :	Any							,			snake_case__      :	Union[str, Any]							,			snake_case__      :	Any							,			snake_case__      :	str							,			snake_case__      :	Tuple							):


        '''simple docstring'''
        super().__init__()
        UpperCAmelCase__						:    int							=							nn.Convad(snake_case__							,			snake_case__							,			1							)
        UpperCAmelCase__						:    Dict							=							nn.Convad(snake_case__							,			snake_case__							,			1							)
        UpperCAmelCase__						:    List[Any]							=							PoolFormerDropPath(snake_case__							)
        if isinstance(config.hidden_act							,			snake_case__							):
             UpperCAmelCase__						:    List[Any]							=							ACTaFN[config.hidden_act]
        else:
             UpperCAmelCase__						:    int							=							config.hidden_act




   def       __a							(		self      :	List[Any]							,			snake_case__      :	Tuple							):


        '''simple docstring'''
        UpperCAmelCase__						:    Dict							=							self.conva(snake_case__							)
        UpperCAmelCase__						:    int							=							self.act_fn(snake_case__							)
        UpperCAmelCase__						:    Optional[int]							=							self.drop(snake_case__							)
        UpperCAmelCase__						:    int							=							self.conva(snake_case__							)
        UpperCAmelCase__						:    str							=							self.drop(snake_case__							)

        return hidden_states




class 			lowerCAmelCase__			(    nn.Module   ):
   def __init__(		self      :	Optional[int]							,			snake_case__      :	Dict							,			snake_case__      :	Optional[int]							,			snake_case__      :	int							,			snake_case__      :	Union[str, Any]							,			snake_case__      :	List[Any]							,			snake_case__      :	int							):


        '''simple docstring'''
        super().__init__()
        UpperCAmelCase__						:    Optional[int]							=							PoolFormerPooling(snake_case__							)
        UpperCAmelCase__						:    Dict							=							PoolFormerOutput(snake_case__							,			snake_case__							,			snake_case__							,			snake_case__							)
        UpperCAmelCase__						:    Tuple							=							PoolFormerGroupNorm(snake_case__							)
        UpperCAmelCase__						:    str							=							PoolFormerGroupNorm(snake_case__							)

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




   def       __a							(		self      :	Tuple							,			snake_case__      :	int							):


        '''simple docstring'''
        if self.use_layer_scale:
             UpperCAmelCase__						:    Dict							=							self.pooling(self.before_norm(snake_case__							)							)
             UpperCAmelCase__						:    str							=							self.layer_scale_a.unsqueeze(-1							).unsqueeze(-1							) * pooling_output
             # First residual connection
             UpperCAmelCase__						:    int							=							hidden_states + self.drop_path(snake_case__							)
             UpperCAmelCase__						:    Dict							=							()

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

             UpperCAmelCase__						:    int							=							(output,) + outputs
             return outputs

        else:
             UpperCAmelCase__						:    Tuple							=							self.drop_path(self.pooling(self.before_norm(snake_case__							)							)							)
             # First residual connection
             UpperCAmelCase__						:    Tuple							=							pooling_output + hidden_states
             UpperCAmelCase__						:    Union[str, Any]							=							()

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

             UpperCAmelCase__						:    int							=							(output,) + outputs
             return outputs




class 			lowerCAmelCase__			(    nn.Module   ):
   def __init__(		self      :	Any							,			snake_case__      :	Optional[int]							):


        '''simple docstring'''
        super().__init__()
        UpperCAmelCase__						:    List[Any]							=							config
        # stochastic depth decay rule
        UpperCAmelCase__						:    Optional[int]							=							[x.item() for x in torch.linspace(0							,			config.drop_path_rate							,			sum(config.depths							)							)]

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

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




   def       __a							(		self      :	List[str]							,			snake_case__      :	Tuple							,			snake_case__      :	Tuple=False							,			snake_case__      :	Any=True							):


        '''simple docstring'''
        UpperCAmelCase__						:    List[str]							=							() if output_hidden_states else None

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

             if output_hidden_states:
                  UpperCAmelCase__						:    Optional[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 			lowerCAmelCase__			(    __magic_name__   ):
   SCREAMING_SNAKE_CASE_    =PoolFormerConfig
   SCREAMING_SNAKE_CASE_    ='''poolformer'''
   SCREAMING_SNAKE_CASE_    ='''pixel_values'''
   SCREAMING_SNAKE_CASE_    =True
   def       __a							(		self      :	Dict							,			snake_case__      :	Tuple							):


        '''simple docstring'''
        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       __a							(		self      :	List[Any]							,			snake_case__      :	Tuple							,			snake_case__      :	Union[str, Any]=False							):


        '''simple docstring'''
        if isinstance(snake_case__							,			snake_case__							):
             UpperCAmelCase__						:    Any							=							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	:	Tuple				=      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.'''							,						__magic_name__							,						)
class 			lowerCAmelCase__			(    __magic_name__   ):
   def __init__(		self      :	Any							,			snake_case__      :	int							):


        '''simple docstring'''
        super().__init__(snake_case__							)
        UpperCAmelCase__						:    List[str]							=							config

        UpperCAmelCase__						:    List[Any]							=							PoolFormerEncoder(snake_case__							)

        # Initialize weights and apply final processing
        self.post_init()
   def       __a							(		self      :	Optional[Any]							):


        '''simple docstring'''
        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       __a							(		self      :	str							,			snake_case__      :	Optional[torch.FloatTensor] = None							,			snake_case__      :	Optional[bool] = None							,			snake_case__      :	Optional[bool] = None							,			):


        '''simple docstring'''
        UpperCAmelCase__						:    List[Any]							=							(
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        UpperCAmelCase__						:    Tuple							=							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__						:    List[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 			lowerCAmelCase__			(    nn.Module   ):
   def __init__(		self      :	List[str]							,			snake_case__      :	int							):


        '''simple docstring'''
        super().__init__()
        UpperCAmelCase__						:    List[str]							=							nn.Linear(config.hidden_size							,			config.hidden_size							)




   def       __a							(		self      :	Dict							,			snake_case__      :	str							):


        '''simple docstring'''
        UpperCAmelCase__						:    Tuple							=							self.dense(snake_case__							)
        return output




@add_start_docstrings(
    '''
    PoolFormer Model transformer with an image classification head on top
    '''							,						__magic_name__							,						)
class 			lowerCAmelCase__			(    __magic_name__   ):
   def __init__(		self      :	List[str]							,			snake_case__      :	Union[str, Any]							):


        '''simple docstring'''
        super().__init__(snake_case__							)
        UpperCAmelCase__						:    List[Any]							=							config.num_labels
        UpperCAmelCase__						:    Optional[Any]							=							PoolFormerModel(snake_case__							)

        # Final norm
        UpperCAmelCase__						:    Tuple							=							PoolFormerGroupNorm(config.hidden_sizes[-1]							)
        # Classifier head
        UpperCAmelCase__						:    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       __a							(		self      :	List[Any]							,			snake_case__      :	Optional[torch.FloatTensor] = None							,			snake_case__      :	Optional[torch.LongTensor] = None							,			snake_case__      :	Optional[bool] = None							,			snake_case__      :	Optional[bool] = None							,			):


        '''simple docstring'''
        UpperCAmelCase__						:    Dict							=							return_dict if return_dict is not None else self.config.use_return_dict

        UpperCAmelCase__						:    Dict							=							self.poolformer(
            snake_case__							,			output_hidden_states=snake_case__							,			return_dict=snake_case__							,			)

        UpperCAmelCase__						:    Any							=							outputs[0]

        UpperCAmelCase__						:    List[Any]							=							self.classifier(self.norm(snake_case__							).mean([-2, -1]							)							)

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

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

        if not return_dict:
             UpperCAmelCase__						:    Dict							=							(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"""



from ...configuration_utils import PretrainedConfig
from ...utils import logging


_lowerCAmelCase	:	Any				=      logging.get_logger(__name__)

_lowerCAmelCase	:	Union[str, Any]				=      {
    # See all MEGATRON_BERT models at https://huggingface.co/models?filter=bert
}




class 			lowerCAmelCase__			(    __magic_name__   ):
   SCREAMING_SNAKE_CASE_    ='''megatron-bert'''
   def __init__(		self      :	Tuple							,			snake_case__      :	List[Any]=2_9_0_5_6							,			snake_case__      :	Any=1_0_2_4							,			snake_case__      :	Tuple=2_4							,			snake_case__      :	Optional[Any]=1_6							,			snake_case__      :	List[str]=4_0_9_6							,			snake_case__      :	Union[str, Any]="gelu"							,			snake_case__      :	int=0.1							,			snake_case__      :	Dict=0.1							,			snake_case__      :	str=5_1_2							,			snake_case__      :	Optional[Any]=2							,			snake_case__      :	Any=0.02							,			snake_case__      :	Optional[Any]=1e-12							,			snake_case__      :	str=0							,			snake_case__      :	Optional[Any]="absolute"							,			snake_case__      :	Optional[Any]=True							,			**snake_case__      :	Any							,			):


        '''simple docstring'''
        super().__init__(pad_token_id=snake_case__							,			**snake_case__							)

        UpperCAmelCase__						:    Union[str, Any]							=							vocab_size
        UpperCAmelCase__						:    Union[str, Any]							=							hidden_size
        UpperCAmelCase__						:    Optional[Any]							=							num_hidden_layers
        UpperCAmelCase__						:    Tuple							=							num_attention_heads
        UpperCAmelCase__						:    List[Any]							=							hidden_act
        UpperCAmelCase__						:    Optional[int]							=							intermediate_size
        UpperCAmelCase__						:    Dict							=							hidden_dropout_prob
        UpperCAmelCase__						:    str							=							attention_probs_dropout_prob
        UpperCAmelCase__						:    str							=							max_position_embeddings
        UpperCAmelCase__						:    List[str]							=							type_vocab_size
        UpperCAmelCase__						:    Tuple							=							initializer_range
        UpperCAmelCase__						:    str							=							layer_norm_eps
        UpperCAmelCase__						:    List[Any]							=							position_embedding_type
        UpperCAmelCase__						:    Dict							=							use_cache


