

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

# General docstring
lowerCAmelCase__	   = """PoolFormerConfig"""

# Base docstring
lowerCAmelCase__	   = """sail/poolformer_s12"""
lowerCAmelCase__	   = [1, 5_1_2, 7, 7]

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

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



def 					lowerCAmelCase__							(	SCREAMING_SNAKE_CASE_:							Optional[int]	,       SCREAMING_SNAKE_CASE_:							float = 0.0	,       SCREAMING_SNAKE_CASE_:							bool = False      )			->      List[Any]:




				'''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(SCREAMING_SNAKE_CASE_	,       dtype=input.dtype	,       device=input.device      )
				random_tensor.floor_()  # binarize
				A__        =			input.div(SCREAMING_SNAKE_CASE_      ) * random_tensor
				return output

class 					a__						(					nn.Module			):
			"""simple docstring"""







			def __init__(    self      ,      lowercase = None       )  ->						None:
							'''simple docstring'''






							super().__init__()
							A__        =			drop_prob


			def      UpperCamelCase		(    self      ,      lowercase       )  ->						torch.Tensor:
							'''simple docstring'''






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


			def      UpperCamelCase		(    self       )  ->						str:
							'''simple docstring'''






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

class 					a__						(					nn.Module			):
			"""simple docstring"""







			def __init__(    self      ,      lowercase      ,      lowercase      ,      lowercase      ,      lowercase      ,      lowercase      ,      lowercase=None       )  ->						int:
							'''simple docstring'''






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

							A__        =			nn.Convad(lowercase      ,      lowercase      ,      kernel_size=lowercase      ,      stride=lowercase      ,      padding=lowercase       )
							A__        =			norm_layer(lowercase       ) if norm_layer else nn.Identity()


			def      UpperCamelCase		(    self      ,      lowercase       )  ->						Optional[Any]:
							'''simple docstring'''






							A__        =			self.projection(lowercase       )
							A__        =			self.norm(lowercase       )
							return embeddings

class 					a__						(					nn.GroupNorm			):
			"""simple docstring"""







			def __init__(    self      ,      lowercase      ,      **lowercase       )  ->						Any:
							'''simple docstring'''






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

class 					a__						(					nn.Module			):
			"""simple docstring"""







			def __init__(    self      ,      lowercase       )  ->						Union[str, Any]:
							'''simple docstring'''






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


			def      UpperCamelCase		(    self      ,      lowercase       )  ->						Any:
							'''simple docstring'''






							return self.pool(lowercase       ) - hidden_states

class 					a__						(					nn.Module			):
			"""simple docstring"""







			def __init__(    self      ,      lowercase      ,      lowercase      ,      lowercase      ,      lowercase       )  ->						Any:
							'''simple docstring'''






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


			def      UpperCamelCase		(    self      ,      lowercase       )  ->						Union[str, Any]:
							'''simple docstring'''






							A__        =			self.conva(lowercase       )
							A__        =			self.act_fn(lowercase       )
							A__        =			self.drop(lowercase       )
							A__        =			self.conva(lowercase       )
							A__        =			self.drop(lowercase       )

							return hidden_states

class 					a__						(					nn.Module			):
			"""simple docstring"""







			def __init__(    self      ,      lowercase      ,      lowercase      ,      lowercase      ,      lowercase      ,      lowercase      ,      lowercase       )  ->						Optional[Any]:
							'''simple docstring'''






							super().__init__()
							A__        =			PoolFormerPooling(lowercase       )
							A__        =			PoolFormerOutput(lowercase      ,      lowercase      ,      lowercase      ,      lowercase       )
							A__        =			PoolFormerGroupNorm(lowercase       )
							A__        =			PoolFormerGroupNorm(lowercase       )

							# Useful for training neural nets
							A__        =			PoolFormerDropPath(lowercase       ) 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((lowercase)       )      ,      requires_grad=lowercase       )
											A__        =			nn.Parameter(
											    config.layer_scale_init_value * torch.ones((lowercase)       )      ,      requires_grad=lowercase       )


			def      UpperCamelCase		(    self      ,      lowercase       )  ->						List[Any]:
							'''simple docstring'''






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

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

											A__        =			(output,) + outputs
											return outputs

							else:
											A__        =			self.drop_path(self.pooling(self.before_norm(lowercase       )       )       )
											# 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(lowercase       )       )       )
											A__        =			hidden_states + layer_output

											A__        =			(output,) + outputs
											return outputs

class 					a__						(					nn.Module			):
			"""simple docstring"""







			def __init__(    self      ,      lowercase       )  ->						Tuple:
							'''simple docstring'''






							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(lowercase       )

							# 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(
															        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       )       )

							A__        =			nn.ModuleList(lowercase       )


			def      UpperCamelCase		(    self      ,      lowercase      ,      lowercase=False      ,      lowercase=True       )  ->						str:
							'''simple docstring'''






							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(lowercase       )
											# Send the embeddings through the blocks
											for _, blk in enumerate(lowercase       ):
															A__        =			blk(lowercase       )
															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=lowercase      ,      hidden_states=lowercase       )

class 					a__						(					snake_case			):
			"""simple docstring"""







			__lowerCamelCase        =       PoolFormerConfig
			__lowerCamelCase        =       'poolformer'
			__lowerCamelCase        =       'pixel_values'
			__lowerCamelCase        =       True


			def      UpperCamelCase		(    self      ,      lowercase       )  ->						str:
							'''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      ,      lowercase      ,      lowercase=False       )  ->						Optional[int]:
							'''simple docstring'''






							if isinstance(lowercase      ,      lowercase       ):
											A__        =			value

lowerCAmelCase__	   = 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__	   = 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 					a__						(					snake_case			):
			"""simple docstring"""







			def __init__(    self      ,      lowercase       )  ->						int:
							'''simple docstring'''






							super().__init__(lowercase       )
							A__        =			config

							A__        =			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      ,      lowercase = None      ,      lowercase = None      ,      lowercase = None      ,      )  ->						Union[Tuple, BaseModelOutputWithNoAttention]:
							'''simple docstring'''






							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(
							    lowercase      ,      output_hidden_states=lowercase      ,      return_dict=lowercase      ,      )
							A__        =			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 					a__						(					nn.Module			):
			"""simple docstring"""







			def __init__(    self      ,      lowercase       )  ->						int:
							'''simple docstring'''






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


			def      UpperCamelCase		(    self      ,      lowercase       )  ->						List[Any]:
							'''simple docstring'''






							A__        =			self.dense(lowercase       )
							return output

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







			def __init__(    self      ,      lowercase       )  ->						List[str]:
							'''simple docstring'''






							super().__init__(lowercase       )
							A__        =			config.num_labels
							A__        =			PoolFormerModel(lowercase       )

							# 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(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      ,      lowercase = None      ,      lowercase = None      ,      lowercase = None      ,      lowercase = None      ,      )  ->						Union[Tuple, ImageClassifierOutputWithNoAttention]:
							'''simple docstring'''






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

							A__        =			self.poolformer(
							    lowercase      ,      output_hidden_states=lowercase      ,      return_dict=lowercase      ,      )

							A__        =			outputs[0]

							A__        =			self.classifier(self.norm(lowercase       ).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(lowercase      ,      lowercase       )
											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(lowercase      ,      lowercase       )

							if not return_dict:
											A__        =			(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 dataclasses import dataclass, field
from typing import Optional

from transformers import AutoConfig, AutoImageProcessor, AutoTokenizer, FlaxVisionEncoderDecoderModel, HfArgumentParser

@dataclass
class 					a__						:
			"""simple docstring"""







			__lowerCamelCase        =       field(
			    metadata={'help': 'The output directory where the model will be written.'}							,	)
			__lowerCamelCase        =       field(
			    metadata={
			        'help': (
			            'The encoder model checkpoint for weights initialization.'
			            'Don\'t set if you want to train an encoder model from scratch.'
			        )
			    }							,	)
			__lowerCamelCase        =       field(
			    metadata={
			        'help': (
			            'The decoder model checkpoint for weights initialization.'
			            'Don\'t set if you want to train a decoder model from scratch.'
			        )
			    }							,	)
			__lowerCamelCase        =       field(
			    default=snake_case							,	metadata={'help': 'Pretrained encoder config name or path if not the same as encoder_model_name'}			)
			__lowerCamelCase        =       field(
			    default=snake_case							,	metadata={'help': 'Pretrained decoder config name or path if not the same as decoder_model_name'}			)



def 					lowerCAmelCase__							(	)			->      List[Any]:




				'''simple docstring'''





				A__        =			HfArgumentParser((ModelArguments,)      )
				((A__)    ,				)        =			parser.parse_args_into_dataclasses()

				# Load pretrained model and tokenizer

				# Use explicit specified encoder config
				if model_args.encoder_config_name:
								A__        =			AutoConfig.from_pretrained(model_args.encoder_config_name      )
				# Use pretrained encoder model's config
				else:
								A__        =			AutoConfig.from_pretrained(model_args.encoder_model_name_or_path      )

				# Use explicit specified decoder config
				if model_args.decoder_config_name:
								A__        =			AutoConfig.from_pretrained(model_args.decoder_config_name      )
				# Use pretrained decoder model's config
				else:
								A__        =			AutoConfig.from_pretrained(model_args.decoder_model_name_or_path      )

				# necessary for `from_encoder_decoder_pretrained` when `decoder_config` is passed
				A__        =			True
				A__        =			True

				A__        =			FlaxVisionEncoderDecoderModel.from_encoder_decoder_pretrained(
				    encoder_pretrained_model_name_or_path=model_args.encoder_model_name_or_path	,       decoder_pretrained_model_name_or_path=model_args.decoder_model_name_or_path	,       encoder_config=SCREAMING_SNAKE_CASE_	,       decoder_config=SCREAMING_SNAKE_CASE_	,       )

				# GPT2 only has bos/eos tokens but not decoder_start/pad tokens
				A__        =			decoder_config.decoder_start_token_id
				A__        =			decoder_config.pad_token_id
				if decoder_start_token_id is None:
								A__        =			decoder_config.bos_token_id
				if pad_token_id is None:
								A__        =			decoder_config.eos_token_id

				# This is necessary to make Flax's generate() work
				A__        =			decoder_config.eos_token_id
				A__        =			decoder_start_token_id
				A__        =			pad_token_id

				A__        =			AutoImageProcessor.from_pretrained(model_args.encoder_model_name_or_path      )

				A__        =			AutoTokenizer.from_pretrained(model_args.decoder_model_name_or_path      )
				A__        =			tokenizer.convert_ids_to_tokens(model.config.pad_token_id      )

				model.save_pretrained(model_args.output_dir      )
				image_processor.save_pretrained(model_args.output_dir      )
				tokenizer.save_pretrained(model_args.output_dir      )


if __name__ == "__main__":
	main()

