





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

# General docstring
lowercase__   :				int      =   '''PoolFormerConfig'''

# Base docstring
lowercase__   :				Tuple      =   '''sail/poolformer_s12'''
lowercase__   :				Tuple      =   [1, 5_12, 7, 7]

# Image classification docstring
lowercase__   :				int      =   '''sail/poolformer_s12'''
lowercase__   :				str      =   '''tabby, tabby cat'''

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






def 	_lowerCAmelCase      (  __snake_case  :			Optional[Any]			,							__snake_case  :			float = 0.0			,							__snake_case  :			bool = False      )					->					Any:





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


class 			SCREAMING_SNAKE_CASE   (nn.Module							):


     def __init__(						self						,     _UpperCAmelCase = None):

          '''simple docstring'''






          super().__init__()
          __A						:			int								=	drop_prob


     def 							SCREAMING_SNAKE_CASE  (						self						,     _UpperCAmelCase):

          '''simple docstring'''






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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






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




class 			SCREAMING_SNAKE_CASE   (nn.Module							):


     def __init__(						self						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase=None):

          '''simple docstring'''






          super().__init__()
          __A						:			List[str]								=	patch_size if isinstance(__A						,     collections.abc.Iterable) else (patch_size, patch_size)
          __A						:			Any								=	stride if isinstance(__A						,     collections.abc.Iterable) else (stride, stride)
          __A						:			Optional[int]								=	padding if isinstance(__A						,     collections.abc.Iterable) else (padding, padding)

          __A						:			Dict								=	nn.Convad(__A						,     __A						,     kernel_size=__A						,     stride=__A						,     padding=__A)
          __A						:			Any								=	norm_layer(__A) if norm_layer else nn.Identity()


     def 							SCREAMING_SNAKE_CASE  (						self						,     _UpperCAmelCase):

          '''simple docstring'''






          __A						:			Any								=	self.projection(__A)
          __A						:			Union[str, Any]								=	self.norm(__A)
          return embeddings




class 			SCREAMING_SNAKE_CASE   (nn.GroupNorm							):


     def __init__(						self						,     _UpperCAmelCase						,     **_UpperCAmelCase):

          '''simple docstring'''






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




class 			SCREAMING_SNAKE_CASE   (nn.Module							):


     def __init__(						self						,     _UpperCAmelCase):

          '''simple docstring'''






          super().__init__()
          __A						:			Union[str, Any]								=	nn.AvgPoolad(__A						,     stride=1						,     padding=pool_size // 2						,     count_include_pad=__A)


     def 							SCREAMING_SNAKE_CASE  (						self						,     _UpperCAmelCase):

          '''simple docstring'''






          return self.pool(__A) - hidden_states




class 			SCREAMING_SNAKE_CASE   (nn.Module							):


     def __init__(						self						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase):

          '''simple docstring'''






          super().__init__()
          __A						:			str								=	nn.Convad(__A						,     __A						,     1)
          __A						:			List[str]								=	nn.Convad(__A						,     __A						,     1)
          __A						:			Tuple								=	PoolFormerDropPath(__A)
          if isinstance(config.hidden_act						,     __A):
               __A						:			Optional[int]								=	ACTaFN[config.hidden_act]
          else:
               __A						:			Tuple								=	config.hidden_act


     def 							SCREAMING_SNAKE_CASE  (						self						,     _UpperCAmelCase):

          '''simple docstring'''






          __A						:			Tuple								=	self.conva(__A)
          __A						:			Union[str, Any]								=	self.act_fn(__A)
          __A						:			Union[str, Any]								=	self.drop(__A)
          __A						:			Dict								=	self.conva(__A)
          __A						:			Any								=	self.drop(__A)

          return hidden_states




class 			SCREAMING_SNAKE_CASE   (nn.Module							):


     def __init__(						self						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase						,     _UpperCAmelCase):

          '''simple docstring'''






          super().__init__()
          __A						:			Union[str, Any]								=	PoolFormerPooling(__A)
          __A						:			List[Any]								=	PoolFormerOutput(__A						,     __A						,     __A						,     __A)
          __A						:			Union[str, Any]								=	PoolFormerGroupNorm(__A)
          __A						:			Optional[int]								=	PoolFormerGroupNorm(__A)

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


     def 							SCREAMING_SNAKE_CASE  (						self						,     _UpperCAmelCase):

          '''simple docstring'''






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

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

               __A						:			Tuple								=	(output,) + outputs
               return outputs

          else:
               __A						:			Any								=	self.drop_path(self.pooling(self.before_norm(__A)))
               # First residual connection
               __A						:			Any								=	pooling_output + hidden_states
               __A						:			Optional[int]								=	()

               # Second residual connection inside the PoolFormerOutput block
               __A						:			Optional[int]								=	self.drop_path(self.output(self.after_norm(__A)))
               __A						:			Optional[int]								=	hidden_states + layer_output

               __A						:			List[Any]								=	(output,) + outputs
               return outputs




class 			SCREAMING_SNAKE_CASE   (nn.Module							):


     def __init__(						self						,     _UpperCAmelCase):

          '''simple docstring'''






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

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

          # Transformer blocks
          __A						:			Union[str, Any]								=	[]
          __A						:			int								=	0
          for i in range(config.num_encoder_blocks):
               # each block consists of layers
               __A						:			Dict								=	[]
               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))

          __A						:			Tuple								=	nn.ModuleList(__A)


     def 							SCREAMING_SNAKE_CASE  (						self						,     _UpperCAmelCase						,     _UpperCAmelCase=False						,     _UpperCAmelCase=True):

          '''simple docstring'''






          __A						:			Union[str, Any]								=	() if output_hidden_states else None

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

               if output_hidden_states:
                    __A						:			Dict								=	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 			SCREAMING_SNAKE_CASE   (A__							):
     lowerCAmelCase		      =	PoolFormerConfig
     lowerCAmelCase		      =	"poolformer"
     lowerCAmelCase		      =	"pixel_values"
     lowerCAmelCase		      =	True


     def 							SCREAMING_SNAKE_CASE  (						self						,     _UpperCAmelCase):

          '''simple docstring'''






          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						,     _UpperCAmelCase						,     _UpperCAmelCase=False):

          '''simple docstring'''






          if isinstance(__A						,     __A):
               __A						:			List[Any]								=	value




lowercase__   :				List[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__   :				Dict      =   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 			SCREAMING_SNAKE_CASE   (A__							):


     def __init__(						self						,     _UpperCAmelCase):

          '''simple docstring'''






          super().__init__(__A)
          __A						:			Dict								=	config

          __A						:			str								=	PoolFormerEncoder(__A)

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


     def 							SCREAMING_SNAKE_CASE  (						self):

          '''simple docstring'''






          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						,     _UpperCAmelCase = None						,     _UpperCAmelCase = None						,     _UpperCAmelCase = None						,     ):

          '''simple docstring'''






          __A						:			List[Any]								=	(
              output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
          )
          __A						:			Optional[int]								=	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						:			Optional[Any]								=	self.encoder(
              __A						,     output_hidden_states=__A						,     return_dict=__A						,     )
          __A						:			Optional[Any]								=	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 			SCREAMING_SNAKE_CASE   (nn.Module							):


     def __init__(						self						,     _UpperCAmelCase):

          '''simple docstring'''






          super().__init__()
          __A						:			str								=	nn.Linear(config.hidden_size						,     config.hidden_size)


     def 							SCREAMING_SNAKE_CASE  (						self						,     _UpperCAmelCase):

          '''simple docstring'''






          __A						:			List[Any]								=	self.dense(__A)
          return output




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


     def __init__(						self						,     _UpperCAmelCase):

          '''simple docstring'''






          super().__init__(__A)
          __A						:			List[str]								=	config.num_labels
          __A						:			Tuple								=	PoolFormerModel(__A)

          # Final norm
          __A						:			Optional[int]								=	PoolFormerGroupNorm(config.hidden_sizes[-1])
          # Classifier head
          __A						:			Optional[int]								=	(
              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						,     _UpperCAmelCase = None						,     _UpperCAmelCase = None						,     _UpperCAmelCase = None						,     _UpperCAmelCase = None						,     ):

          '''simple docstring'''






          __A						:			Any								=	return_dict if return_dict is not None else self.config.use_return_dict

          __A						:			List[str]								=	self.poolformer(
              __A						,     output_hidden_states=__A						,     return_dict=__A						,     )

          __A						:			Optional[int]								=	outputs[0]

          __A						:			Dict								=	self.classifier(self.norm(__A).mean([-2, -1]))

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

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

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

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






'''simple docstring'''




from arguments import InitializationArguments

from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, HfArgumentParser


# Configuration
lowercase__   :				List[Any]      =   HfArgumentParser(InitializationArguments)
lowercase__   :				List[Any]      =   parser.parse_args()

# Load codeparrot tokenizer trained for Python code tokenization
lowercase__   :				Union[str, Any]      =   AutoTokenizer.from_pretrained(args.tokenizer_name)

# Config: "scale_attn_by_layer_idx" and "reorder_and_upcast_attn" are Mistral stability tweaks
lowercase__   :				List[str]      =   {
    '''vocab_size''': len(tokenizer),
    '''scale_attn_by_inverse_layer_idx''': True,
    '''reorder_and_upcast_attn''': True,
}

# Load model config (GPT-2 large in this case)
lowercase__   :				Optional[int]      =   AutoConfig.from_pretrained(args.config_name, **config_kwargs)

# Initialize new model with config
lowercase__   :				Optional[int]      =   AutoModelForCausalLM.from_config(config)

# Save model to the hub
model.save_pretrained(args.model_name, push_to_hub=args.push_to_hub)