


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


__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_   (   _lowerCAmelCase     :					List[str],     _lowerCAmelCase     :					float = 0.0,     _lowerCAmelCase     :					bool = False       ):






     """simple docstring"""

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


class     __lowerCamelCase  (					nn.Module				):

     '''simple docstring'''


     def __init__(      self    ,				__UpperCAmelCase = None							)       ->							None:
          super().__init__()
          _a    	=				drop_prob


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							torch.Tensor:
          return drop_path(__UpperCAmelCase    ,				self.drop_prob    ,				self.training							)





     def 							_UpperCAmelCase      (      self							)       ->							str:
          return "p={}".format(self.drop_prob							)







class     __lowerCamelCase  (					nn.Module				):

     '''simple docstring'''


     def __init__(      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase=None							)       ->							str:
          super().__init__()
          _a    	=				patch_size if isinstance(__UpperCAmelCase    ,				collections.abc.Iterable							) else (patch_size, patch_size)
          _a    	=				stride if isinstance(__UpperCAmelCase    ,				collections.abc.Iterable							) else (stride, stride)
          _a    	=				padding if isinstance(__UpperCAmelCase    ,				collections.abc.Iterable							) else (padding, padding)

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





     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							Dict:
          _a    	=				self.projection(__UpperCAmelCase							)
          _a    	=				self.norm(__UpperCAmelCase							)
          return embeddings







class     __lowerCamelCase  (					nn.GroupNorm				):

     '''simple docstring'''


     def __init__(      self    ,				__UpperCAmelCase    ,				**__UpperCAmelCase							)       ->							int:
          super().__init__(1    ,				__UpperCAmelCase    ,				**__UpperCAmelCase							)







class     __lowerCamelCase  (					nn.Module				):

     '''simple docstring'''


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





     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							int:
          return self.pool(__UpperCAmelCase							) - hidden_states







class     __lowerCamelCase  (					nn.Module				):

     '''simple docstring'''


     def __init__(      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)       ->							Any:
          super().__init__()
          _a    	=				nn.Convad(__UpperCAmelCase    ,				__UpperCAmelCase    ,				1							)
          _a    	=				nn.Convad(__UpperCAmelCase    ,				__UpperCAmelCase    ,				1							)
          _a    	=				PoolFormerDropPath(__UpperCAmelCase							)
          if isinstance(config.hidden_act    ,				__UpperCAmelCase							):
               _a    	=				ACTaFN[config.hidden_act]
          else:
               _a    	=				config.hidden_act





     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							Any:
          _a    	=				self.conva(__UpperCAmelCase							)
          _a    	=				self.act_fn(__UpperCAmelCase							)
          _a    	=				self.drop(__UpperCAmelCase							)
          _a    	=				self.conva(__UpperCAmelCase							)
          _a    	=				self.drop(__UpperCAmelCase							)

          return hidden_states







class     __lowerCamelCase  (					nn.Module				):

     '''simple docstring'''


     def __init__(      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)       ->							Any:
          super().__init__()
          _a    	=				PoolFormerPooling(__UpperCAmelCase							)
          _a    	=				PoolFormerOutput(__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase    ,				__UpperCAmelCase							)
          _a    	=				PoolFormerGroupNorm(__UpperCAmelCase							)
          _a    	=				PoolFormerGroupNorm(__UpperCAmelCase							)

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





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

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

               _a    	=				(output,) + outputs
               return outputs

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

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

               _a    	=				(output,) + outputs
               return outputs







class     __lowerCamelCase  (					nn.Module				):

     '''simple docstring'''


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

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

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

          _a    	=				nn.ModuleList(__UpperCAmelCase							)





     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase=False    ,				__UpperCAmelCase=True							)       ->							Tuple:
          _a    	=				() if output_hidden_states else None

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

               if output_hidden_states:
                    _a    	=				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=__UpperCAmelCase    ,				hidden_states=__UpperCAmelCase							)







class     __lowerCamelCase  (					a__				):

     '''simple docstring'''
     A_       :	Optional[Any]    =		PoolFormerConfig
     A_       :	Tuple    =		'poolformer'
     A_       :	Any    =		'pixel_values'
     A_       :	List[str]    =		True


     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							Union[str, Any]:
          if isinstance(__UpperCAmelCase    ,				(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(__UpperCAmelCase    ,				nn.LayerNorm							):
               module.bias.data.zero_()
               module.weight.data.fill_(1.0							)





     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase    ,				__UpperCAmelCase=False							)       ->							int:
          if isinstance(__UpperCAmelCase    ,				__UpperCAmelCase							):
               _a    	=				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.'  ,					a__  ,					)
class     __lowerCamelCase  (					a__				):

     '''simple docstring'''


     def __init__(      self    ,				__UpperCAmelCase							)       ->							Dict:
          super().__init__(__UpperCAmelCase							)
          _a    	=				config

          _a    	=				PoolFormerEncoder(__UpperCAmelCase							)

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


     def 							_UpperCAmelCase      (      self							)       ->							Optional[Any]:
          return self.embeddings.patch_embeddings





     @add_start_docstrings_to_model_forward(__UpperCAmelCase							)
     @add_code_sample_docstrings(
         checkpoint=_CHECKPOINT_FOR_DOC    ,				output_type=__UpperCAmelCase    ,				config_class=_CONFIG_FOR_DOC    ,				modality='''vision'''    ,				expected_output=_EXPECTED_OUTPUT_SHAPE    ,				)
     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase = None    ,				__UpperCAmelCase = None    ,				__UpperCAmelCase = None    ,				)       ->							Union[Tuple, BaseModelOutputWithNoAttention]:
          _a    	=				(
              output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
          )
          _a    	=				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'''							)

          _a    	=				self.encoder(
              __UpperCAmelCase    ,				output_hidden_states=__UpperCAmelCase    ,				return_dict=__UpperCAmelCase    ,				)
          _a    	=				encoder_outputs[0]

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

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







class     __lowerCamelCase  (					nn.Module				):

     '''simple docstring'''


     def __init__(      self    ,				__UpperCAmelCase							)       ->							List[Any]:
          super().__init__()
          _a    	=				nn.Linear(config.hidden_size    ,				config.hidden_size							)





     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase							)       ->							Dict:
          _a    	=				self.dense(__UpperCAmelCase							)
          return output







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

     '''simple docstring'''


     def __init__(      self    ,				__UpperCAmelCase							)       ->							Tuple:
          super().__init__(__UpperCAmelCase							)
          _a    	=				config.num_labels
          _a    	=				PoolFormerModel(__UpperCAmelCase							)

          # Final norm
          _a    	=				PoolFormerGroupNorm(config.hidden_sizes[-1]							)
          # Classifier head
          _a    	=				(
              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(__UpperCAmelCase							)
     @add_code_sample_docstrings(
         checkpoint=_IMAGE_CLASS_CHECKPOINT    ,				output_type=__UpperCAmelCase    ,				config_class=_CONFIG_FOR_DOC    ,				expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT    ,				)
     def 							_UpperCAmelCase      (      self    ,				__UpperCAmelCase = None    ,				__UpperCAmelCase = None    ,				__UpperCAmelCase = None    ,				__UpperCAmelCase = None    ,				)       ->							Union[Tuple, ImageClassifierOutputWithNoAttention]:
          _a    	=				return_dict if return_dict is not None else self.config.use_return_dict

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

          _a    	=				outputs[0]

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

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

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

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

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



"""simple docstring"""






import argparse

import torch

from transformers import BertConfig, BertForPreTraining, load_tf_weights_in_bert
from transformers.utils import logging


logging.set_verbosity_info()
def       A_   (   _lowerCAmelCase     :					Dict,     _lowerCAmelCase     :					List[Any],     _lowerCAmelCase     :					Union[str, Any]       ):






     """simple docstring"""

     _a    	=				BertConfig.from_json_file(_lowerCAmelCase       )
     print(f'Building PyTorch model from configuration: {config}'       )
     _a    	=				BertForPreTraining(_lowerCAmelCase       )

     # Load weights from tf checkpoint
     load_tf_weights_in_bert(_lowerCAmelCase,     _lowerCAmelCase,     _lowerCAmelCase       )

     # Save pytorch-model
     print(f'Save PyTorch model to {pytorch_dump_path}'       )
     torch.save(model.state_dict(),     _lowerCAmelCase       )


if __name__ == "__main__":
   __snake_case		=       argparse.ArgumentParser()
   # Required parameters
   parser.add_argument(
       '''--tf_checkpoint_path''', default=None, type=str, required=True, help='''Path to the TensorFlow checkpoint path.'''
   )
   parser.add_argument(
       '''--bert_config_file''',
       default=None,
       type=str,
       required=True,
       help=(
           '''The config json file corresponding to the pre-trained BERT model. \n'''
           '''This specifies the model architecture.'''
       ),
   )
   parser.add_argument(
       '''--pytorch_dump_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
   )
   __snake_case		=       parser.parse_args()
   convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.bert_config_file, args.pytorch_dump_path)