



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_12, 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     SCREAMING_SNAKE_CASE(      __UpperCamelCase	,				__UpperCamelCase = 0.0	,				__UpperCamelCase = False )	->						Dict:
				if drop_prob == 0.0 or not training:
								return input
				a__	:							Tuple      						= 1 - drop_prob
				a__	:							Any      						= (input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
				a__	:							Dict      						= keep_prob + torch.rand(__UpperCamelCase	,				dtype=input.dtype	,				device=input.device )
				random_tensor.floor_()  # binarize
				a__	:							Optional[int]      						= input.div(__UpperCamelCase ) * random_tensor
				return output


class 				_a						(       nn.Module			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase = None				):







											"""simple docstring"""



											super().__init__()
											a__	:							Optional[Any]      						= drop_prob



							def 		_A	(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



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




							def 		_A	(							self				):







											"""simple docstring"""



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





class 				_a						(       nn.Module			):





							'''simple docstring'''

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







											"""simple docstring"""



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

											a__	:							int      						= nn.Convad(__UpperCAmelCase ,					__UpperCAmelCase ,					kernel_size=__UpperCAmelCase ,					stride=__UpperCAmelCase ,					padding=__UpperCAmelCase				)
											a__	:							Optional[Any]      						= norm_layer(__UpperCAmelCase				) if norm_layer else nn.Identity()




							def 		_A	(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



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





class 				_a						(       nn.GroupNorm			):





							'''simple docstring'''

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







											"""simple docstring"""



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





class 				_a						(       nn.Module			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											super().__init__()
											a__	:							List[str]      						= nn.AvgPoolad(__UpperCAmelCase ,					stride=1 ,					padding=pool_size // 2 ,					count_include_pad=__UpperCAmelCase				)




							def 		_A	(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											return self.pool(__UpperCAmelCase				) - hidden_states





class 				_a						(       nn.Module			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase				):







											"""simple docstring"""



											super().__init__()
											a__	:							str      						= nn.Convad(__UpperCAmelCase ,					__UpperCAmelCase ,					1				)
											a__	:							int      						= nn.Convad(__UpperCAmelCase ,					__UpperCAmelCase ,					1				)
											a__	:							Any      						= PoolFormerDropPath(__UpperCAmelCase				)
											if isinstance(config.hidden_act ,					__UpperCAmelCase				):
															a__	:							List[Any]      						= ACTaFN[config.hidden_act]
											else:
															a__	:							Union[str, Any]      						= config.hidden_act




							def 		_A	(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											a__	:							Union[str, Any]      						= self.conva(__UpperCAmelCase				)
											a__	:							Union[str, Any]      						= self.act_fn(__UpperCAmelCase				)
											a__	:							Dict      						= self.drop(__UpperCAmelCase				)
											a__	:							Union[str, Any]      						= self.conva(__UpperCAmelCase				)
											a__	:							int      						= self.drop(__UpperCAmelCase				)

											return hidden_states





class 				_a						(       nn.Module			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase				):







											"""simple docstring"""



											super().__init__()
											a__	:							Any      						= PoolFormerPooling(__UpperCAmelCase				)
											a__	:							Any      						= PoolFormerOutput(__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase				)
											a__	:							Any      						= PoolFormerGroupNorm(__UpperCAmelCase				)
											a__	:							Dict      						= PoolFormerGroupNorm(__UpperCAmelCase				)

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




							def 		_A	(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											if self.use_layer_scale:
															a__	:							Any      						= self.pooling(self.before_norm(__UpperCAmelCase				)				)
															a__	:							int      						= self.layer_scale_a.unsqueeze(-1				).unsqueeze(-1				) * pooling_output
															# First residual connection
															a__	:							Optional[int]      						= hidden_states + self.drop_path(__UpperCAmelCase				)
															a__	:							Dict      						= ()

															a__	:							List[Any]      						= self.output(self.after_norm(__UpperCAmelCase				)				)
															a__	:							Union[str, Any]      						= self.layer_scale_a.unsqueeze(-1				).unsqueeze(-1				) * layer_output
															# Second residual connection
															a__	:							Optional[Any]      						= hidden_states + self.drop_path(__UpperCAmelCase				)

															a__	:							Optional[int]      						= (output,) + outputs
															return outputs

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

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

															a__	:							Optional[Any]      						= (output,) + outputs
															return outputs





class 				_a						(       nn.Module			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



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

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

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




							def 		_A	(							self ,					__UpperCAmelCase ,					__UpperCAmelCase=False ,					__UpperCAmelCase=True				):







											"""simple docstring"""



											a__	:							int      						= () if output_hidden_states else None

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

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





class 				_a						(       SCREAMING_SNAKE_CASE			):





							'''simple docstring'''

							A				:Optional[int] =       PoolFormerConfig
							A				:List[str] =       "poolformer"
							A				:Tuple =       "pixel_values"
							A				:List[str] =       True



							def 		_A	(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											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 		_A	(							self ,					__UpperCAmelCase ,					__UpperCAmelCase=False				):







											"""simple docstring"""



											if isinstance(__UpperCAmelCase ,					__UpperCAmelCase				):
															a__	:							Dict      						= 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."       ,  SCREAMING_SNAKE_CASE       ,  )
class 				_a						(       SCREAMING_SNAKE_CASE			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											super().__init__(__UpperCAmelCase				)
											a__	:							Optional[Any]      						= config

											a__	:							int      						= PoolFormerEncoder(__UpperCAmelCase				)

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



							def 		_A	(							self				):







											"""simple docstring"""



											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 		_A	(							self ,					__UpperCAmelCase = None ,					__UpperCAmelCase = None ,					__UpperCAmelCase = None ,					):







											"""simple docstring"""



											a__	:							int      						= (
											    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
											)
											a__	:							str      						= 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__	:							List[Any]      						= self.encoder(
											    __UpperCAmelCase ,					output_hidden_states=__UpperCAmelCase ,					return_dict=__UpperCAmelCase ,					)
											a__	:							Any      						= 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 				_a						(       nn.Module			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											super().__init__()
											a__	:							Any      						= nn.Linear(config.hidden_size ,					config.hidden_size				)




							def 		_A	(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											a__	:							Any      						= self.dense(__UpperCAmelCase				)
											return output





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





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											super().__init__(__UpperCAmelCase				)
											a__	:							Optional[Any]      						= config.num_labels
											a__	:							int      						= PoolFormerModel(__UpperCAmelCase				)

											# Final norm
											a__	:							Any      						= PoolFormerGroupNorm(config.hidden_sizes[-1]				)
											# Classifier head
											a__	:							Dict      						= (
											    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 		_A	(							self ,					__UpperCAmelCase = None ,					__UpperCAmelCase = None ,					__UpperCAmelCase = None ,					__UpperCAmelCase = None ,					):







											"""simple docstring"""



											a__	:							str      						= return_dict if return_dict is not None else self.config.use_return_dict

											a__	:							str      						= self.poolformer(
											    __UpperCAmelCase ,					output_hidden_states=__UpperCAmelCase ,					return_dict=__UpperCAmelCase ,					)

											a__	:							Optional[Any]      						= outputs[0]

											a__	:							Union[str, Any]      						= self.classifier(self.norm(__UpperCAmelCase				).mean([-2, -1]				)				)

											a__	:							List[str]      						= 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[str]      						= "single_label_classification"
																			else:
																							a__	:							str      						= "multi_label_classification"

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

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

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







import argparse
import os

import transformers

from .convert_slow_tokenizer import SLOW_TO_FAST_CONVERTERS
from .utils import logging


logging.set_verbosity_info()

lowerCamelCase		       =      logging.get_logger(__name__)


lowerCamelCase		       =      {name: getattr(transformers, name + """Fast""") for name in SLOW_TO_FAST_CONVERTERS}






def     SCREAMING_SNAKE_CASE(      __UpperCamelCase	,				__UpperCamelCase	,				__UpperCamelCase	,				__UpperCamelCase )	->						Dict:
				if tokenizer_name is not None and tokenizer_name not in TOKENIZER_CLASSES:
								raise ValueError(F'Unrecognized tokenizer name, should be one of {list(TOKENIZER_CLASSES.keys() )}.' )

				if tokenizer_name is None:
								a__	:							str      						= TOKENIZER_CLASSES
				else:
								a__	:							int      						= {tokenizer_name: getattr(__UpperCamelCase	,				tokenizer_name + "Fast" )}

				logger.info(F'Loading tokenizer classes: {tokenizer_names}' )

				for tokenizer_name in tokenizer_names:
								a__	:							Any      						= TOKENIZER_CLASSES[tokenizer_name]

								a__	:							Dict      						= True
								if checkpoint_name is None:
												a__	:							Tuple      						= list(tokenizer_class.max_model_input_sizes.keys() )
								else:
												a__	:							List[str]      						= [checkpoint_name]

								logger.info(F'For tokenizer {tokenizer_class.__class__.__name__} loading checkpoints: {checkpoint_names}' )

								for checkpoint in checkpoint_names:
												logger.info(F'Loading {tokenizer_class.__class__.__name__} {checkpoint}' )

												# Load tokenizer
												a__	:							List[str]      						= tokenizer_class.from_pretrained(__UpperCamelCase	,				force_download=__UpperCamelCase )

												# Save fast tokenizer
												logger.info(F'Save fast tokenizer to {dump_path} with prefix {checkpoint} add_prefix {add_prefix}' )

												# For organization names we create sub-directories
												if "/" in checkpoint:
																a__						,   a__	:							Dict      						= checkpoint.split("/" )
																a__	:							Any      						= os.path.join(__UpperCamelCase	,				__UpperCamelCase )
												elif add_prefix:
																a__	:							int      						= checkpoint
																a__	:							Optional[Any]      						= dump_path
												else:
																a__	:							Dict      						= None
																a__	:							Optional[Any]      						= dump_path

												logger.info(F'=> {dump_path_full} with prefix {checkpoint_prefix_name}, add_prefix {add_prefix}' )

												if checkpoint in list(tokenizer.pretrained_vocab_files_map.values() )[0]:
																a__	:							Dict      						= list(tokenizer.pretrained_vocab_files_map.values() )[0][checkpoint]
																a__	:							Optional[Any]      						= file_path.split(__UpperCamelCase )[-1][0]
																if next_char == "/":
																				a__	:							Union[str, Any]      						= os.path.join(__UpperCamelCase	,				__UpperCamelCase )
																				a__	:							Any      						= None

																logger.info(F'=> {dump_path_full} with prefix {checkpoint_prefix_name}, add_prefix {add_prefix}' )

												a__	:							Any      						= tokenizer.save_pretrained(
												    __UpperCamelCase	,				legacy_format=__UpperCamelCase	,				filename_prefix=__UpperCamelCase )
												logger.info(F'=> File names {file_names}' )

												for file_name in file_names:
																if not file_name.endswith("tokenizer.json" ):
																				os.remove(__UpperCamelCase )
																				logger.info(F'=> removing {file_name}' )


if __name__ == "__main__":
						lowerCamelCase		       =      argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    """--dump_path""", default=None, type=str, required=True, help="""Path to output generated fast tokenizer files."""
						)
						parser.add_argument(
						    """--tokenizer_name""",
						    default=None,
						    type=str,
						    help=(
						        F'Optional tokenizer type selected in the list of {list(TOKENIZER_CLASSES.keys())}. If not given, will '
						        """download and convert all the checkpoints from AWS."""
						    ),
						)
						parser.add_argument(
						    """--checkpoint_name""",
						    default=None,
						    type=str,
						    help="""Optional checkpoint name. If not given, will download and convert the canonical checkpoints from AWS.""",
						)
						parser.add_argument(
						    """--force_download""",
						    action="""store_true""",
						    help="""Re-download checkpoints.""",
						)
						lowerCamelCase		       =      parser.parse_args()

						convert_slow_checkpoint_to_fast(args.tokenizer_name, args.checkpoint_name, args.dump_path, args.force_download)


