
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


__a  			=					logging.get_logger(__name__)

# General docstring
__a  			=					'PoolFormerConfig'

# Base docstring
__a  			=					'sail/poolformer_s12'
__a  			=					[1, 5_1_2, 7, 7]

# Image classification docstring
__a  			=					'sail/poolformer_s12'
__a  			=					'tabby, tabby cat'

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




def 				a						(							snake_case__:					List[str]						,   snake_case__:					float = 0.0						,   snake_case__:					bool = False					):


				'''simple docstring'''
				if drop_prob == 0.0 or not training:
								return input
				lowercase_       =					1 - drop_prob
				lowercase_       =					(input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
				lowercase_       =					keep_prob + torch.rand(snake_case__						,   dtype=input.dtype						,   device=input.device					)
				random_tensor.floor_()  # binarize
				lowercase_       =					input.div(snake_case__					) * random_tensor
				return output
class  lowercase__(				nn.Module							):



	"""simple docstring"""




	def __init__(       self  :			Dict       ,							SCREAMING_SNAKE_CASE_  :			Optional[float] = None       )       ->     None:
					super().__init__()
					lowercase_       =					drop_prob
	def 						_lowercase  (       self  :			Tuple       ,							SCREAMING_SNAKE_CASE_  :			torch.Tensor       )       ->     torch.Tensor:
					return drop_path(SCREAMING_SNAKE_CASE_       ,							self.drop_prob       ,							self.training       )




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

class  lowercase__(				nn.Module							):



	"""simple docstring"""




	def __init__(       self  :			Tuple       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]       ,							SCREAMING_SNAKE_CASE_  :			Any       ,							SCREAMING_SNAKE_CASE_  :			Dict       ,							SCREAMING_SNAKE_CASE_  :			Dict       ,							SCREAMING_SNAKE_CASE_  :			Any=None       )       ->     int:
					super().__init__()
					lowercase_       =					patch_size if isinstance(SCREAMING_SNAKE_CASE_       ,							collections.abc.Iterable       ) else (patch_size, patch_size)
					lowercase_       =					stride if isinstance(SCREAMING_SNAKE_CASE_       ,							collections.abc.Iterable       ) else (stride, stride)
					lowercase_       =					padding if isinstance(SCREAMING_SNAKE_CASE_       ,							collections.abc.Iterable       ) else (padding, padding)

					lowercase_       =					nn.Convad(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							kernel_size=SCREAMING_SNAKE_CASE_       ,							stride=SCREAMING_SNAKE_CASE_       ,							padding=SCREAMING_SNAKE_CASE_       )
					lowercase_       =					norm_layer(SCREAMING_SNAKE_CASE_       ) if norm_layer else nn.Identity()




	def 						_lowercase  (       self  :			str       ,							SCREAMING_SNAKE_CASE_  :			int       )       ->     Optional[Any]:
					lowercase_       =					self.projection(SCREAMING_SNAKE_CASE_       )
					lowercase_       =					self.norm(SCREAMING_SNAKE_CASE_       )
					return embeddings

class  lowercase__(				nn.GroupNorm							):



	"""simple docstring"""




	def __init__(       self  :			Union[str, Any]       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]       ,							**SCREAMING_SNAKE_CASE_  :			int       )       ->     List[str]:
					super().__init__(1       ,							SCREAMING_SNAKE_CASE_       ,							**SCREAMING_SNAKE_CASE_       )

class  lowercase__(				nn.Module							):



	"""simple docstring"""




	def __init__(       self  :			Optional[Any]       ,							SCREAMING_SNAKE_CASE_  :			int       )       ->     Optional[int]:
					super().__init__()
					lowercase_       =					nn.AvgPoolad(SCREAMING_SNAKE_CASE_       ,							stride=1       ,							padding=pool_size // 2       ,							count_include_pad=SCREAMING_SNAKE_CASE_       )




	def 						_lowercase  (       self  :			Optional[Any]       ,							SCREAMING_SNAKE_CASE_  :			Dict       )       ->     List[str]:
					return self.pool(SCREAMING_SNAKE_CASE_       ) - hidden_states

class  lowercase__(				nn.Module							):



	"""simple docstring"""




	def __init__(       self  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			Dict       ,							SCREAMING_SNAKE_CASE_  :			Any       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]       )       ->     List[str]:
					super().__init__()
					lowercase_       =					nn.Convad(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							1       )
					lowercase_       =					nn.Convad(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							1       )
					lowercase_       =					PoolFormerDropPath(SCREAMING_SNAKE_CASE_       )
					if isinstance(config.hidden_act       ,							SCREAMING_SNAKE_CASE_       ):
									lowercase_       =					ACTaFN[config.hidden_act]
					else:
									lowercase_       =					config.hidden_act




	def 						_lowercase  (       self  :			List[str]       ,							SCREAMING_SNAKE_CASE_  :			Tuple       )       ->     Optional[int]:
					lowercase_       =					self.conva(SCREAMING_SNAKE_CASE_       )
					lowercase_       =					self.act_fn(SCREAMING_SNAKE_CASE_       )
					lowercase_       =					self.drop(SCREAMING_SNAKE_CASE_       )
					lowercase_       =					self.conva(SCREAMING_SNAKE_CASE_       )
					lowercase_       =					self.drop(SCREAMING_SNAKE_CASE_       )

					return hidden_states

class  lowercase__(				nn.Module							):



	"""simple docstring"""




	def __init__(       self  :			Any       ,							SCREAMING_SNAKE_CASE_  :			str       ,							SCREAMING_SNAKE_CASE_  :			List[str]       ,							SCREAMING_SNAKE_CASE_  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			List[str]       ,							SCREAMING_SNAKE_CASE_  :			List[str]       )       ->     List[Any]:
					super().__init__()
					lowercase_       =					PoolFormerPooling(SCREAMING_SNAKE_CASE_       )
					lowercase_       =					PoolFormerOutput(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       )
					lowercase_       =					PoolFormerGroupNorm(SCREAMING_SNAKE_CASE_       )
					lowercase_       =					PoolFormerGroupNorm(SCREAMING_SNAKE_CASE_       )

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




	def 						_lowercase  (       self  :			List[str]       ,							SCREAMING_SNAKE_CASE_  :			Union[str, Any]       )       ->     Dict:
					if self.use_layer_scale:
									lowercase_       =					self.pooling(self.before_norm(SCREAMING_SNAKE_CASE_       )       )
									lowercase_       =					self.layer_scale_a.unsqueeze(-1       ).unsqueeze(-1       ) * pooling_output
									# First residual connection
									lowercase_       =					hidden_states + self.drop_path(SCREAMING_SNAKE_CASE_       )
									lowercase_       =					()

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

									lowercase_       =					(output,) + outputs
									return outputs

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

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

									lowercase_       =					(output,) + outputs
									return outputs

class  lowercase__(				nn.Module							):



	"""simple docstring"""




	def __init__(       self  :			List[str]       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]       )       ->     Dict:
					super().__init__()
					lowercase_       =					config
					# stochastic depth decay rule
					lowercase_       =					[x.item() for x in torch.linspace(0       ,							config.drop_path_rate       ,							sum(config.depths       )       )]

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

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

					lowercase_       =					nn.ModuleList(SCREAMING_SNAKE_CASE_       )




	def 						_lowercase  (       self  :			Optional[int]       ,							SCREAMING_SNAKE_CASE_  :			Optional[int]       ,							SCREAMING_SNAKE_CASE_  :			Optional[Any]=False       ,							SCREAMING_SNAKE_CASE_  :			int=True       )       ->     Dict:
					lowercase_       =					() if output_hidden_states else None

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

									if output_hidden_states:
													lowercase_       =					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=SCREAMING_SNAKE_CASE_       ,							hidden_states=SCREAMING_SNAKE_CASE_       )

class  lowercase__(				UpperCAmelCase							):



	"""simple docstring"""




	a			:List[str]         =	PoolFormerConfig
	a			:int         =	'poolformer'
	a			:List[Any]         =	'pixel_values'
	a			:Union[str, Any]         =	True
	def 						_lowercase  (       self  :			List[str]       ,							SCREAMING_SNAKE_CASE_  :			Any       )       ->     List[Any]:
					if isinstance(SCREAMING_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(SCREAMING_SNAKE_CASE_       ,							nn.LayerNorm       ):
									module.bias.data.zero_()
									module.weight.data.fill_(1.0       )




	def 						_lowercase  (       self  :			Dict       ,							SCREAMING_SNAKE_CASE_  :			str       ,							SCREAMING_SNAKE_CASE_  :			List[str]=False       )       ->     Any:
					if isinstance(SCREAMING_SNAKE_CASE_       ,							SCREAMING_SNAKE_CASE_       ):
									lowercase_       =					value

__a  			=					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'

__a  			=					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.' ,						UpperCAmelCase ,						)
class  lowercase__(				UpperCAmelCase							):



	"""simple docstring"""




	def __init__(       self  :			List[str]       ,							SCREAMING_SNAKE_CASE_  :			int       )       ->     Optional[Any]:
					super().__init__(SCREAMING_SNAKE_CASE_       )
					lowercase_       =					config

					lowercase_       =					PoolFormerEncoder(SCREAMING_SNAKE_CASE_       )

					# Initialize weights and apply final processing
					self.post_init()
	def 						_lowercase  (       self  :			Dict       )       ->     Any:
					return self.embeddings.patch_embeddings




	@add_start_docstrings_to_model_forward(SCREAMING_SNAKE_CASE_       )
	@add_code_sample_docstrings(
	    checkpoint=_CHECKPOINT_FOR_DOC       ,							output_type=SCREAMING_SNAKE_CASE_       ,							config_class=_CONFIG_FOR_DOC       ,							modality='''vision'''       ,							expected_output=_EXPECTED_OUTPUT_SHAPE       ,							)
	def 						_lowercase  (       self  :			int       ,							SCREAMING_SNAKE_CASE_  :			Optional[torch.FloatTensor] = None       ,							SCREAMING_SNAKE_CASE_  :			Optional[bool] = None       ,							SCREAMING_SNAKE_CASE_  :			Optional[bool] = None       ,							)       ->     Union[Tuple, BaseModelOutputWithNoAttention]:
					lowercase_       =					(
					    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
					)
					lowercase_       =					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'''       )

					lowercase_       =					self.encoder(
					    SCREAMING_SNAKE_CASE_       ,							output_hidden_states=SCREAMING_SNAKE_CASE_       ,							return_dict=SCREAMING_SNAKE_CASE_       ,							)
					lowercase_       =					encoder_outputs[0]

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

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

class  lowercase__(				nn.Module							):



	"""simple docstring"""




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




	def 						_lowercase  (       self  :			str       ,							SCREAMING_SNAKE_CASE_  :			Union[str, Any]       )       ->     Union[str, Any]:
					lowercase_       =					self.dense(SCREAMING_SNAKE_CASE_       )
					return output

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



	"""simple docstring"""




	def __init__(       self  :			Dict       ,							SCREAMING_SNAKE_CASE_  :			str       )       ->     Union[str, Any]:
					super().__init__(SCREAMING_SNAKE_CASE_       )
					lowercase_       =					config.num_labels
					lowercase_       =					PoolFormerModel(SCREAMING_SNAKE_CASE_       )

					# Final norm
					lowercase_       =					PoolFormerGroupNorm(config.hidden_sizes[-1]       )
					# Classifier head
					lowercase_       =					(
					    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(SCREAMING_SNAKE_CASE_       )
	@add_code_sample_docstrings(
	    checkpoint=_IMAGE_CLASS_CHECKPOINT       ,							output_type=SCREAMING_SNAKE_CASE_       ,							config_class=_CONFIG_FOR_DOC       ,							expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT       ,							)
	def 						_lowercase  (       self  :			List[Any]       ,							SCREAMING_SNAKE_CASE_  :			Optional[torch.FloatTensor] = None       ,							SCREAMING_SNAKE_CASE_  :			Optional[torch.LongTensor] = None       ,							SCREAMING_SNAKE_CASE_  :			Optional[bool] = None       ,							SCREAMING_SNAKE_CASE_  :			Optional[bool] = None       ,							)       ->     Union[Tuple, ImageClassifierOutputWithNoAttention]:
					lowercase_       =					return_dict if return_dict is not None else self.config.use_return_dict

					lowercase_       =					self.poolformer(
					    SCREAMING_SNAKE_CASE_       ,							output_hidden_states=SCREAMING_SNAKE_CASE_       ,							return_dict=SCREAMING_SNAKE_CASE_       ,							)

					lowercase_       =					outputs[0]

					lowercase_       =					self.classifier(self.norm(SCREAMING_SNAKE_CASE_       ).mean([-2, -1]       )       )

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

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

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

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






import argparse
from collections import defaultdict

import yaml


__a  			=					'docs/source/en/_toctree.yml'




def 				a						(							snake_case__:					Dict					):


				'''simple docstring'''
				lowercase_       =					defaultdict(snake_case__					)
				for doc in model_doc:
								counts[doc["local"]] += 1
				lowercase_       =					[key for key, value in counts.items() if value > 1]

				lowercase_       =					[]
				for duplicate_key in duplicates:
								lowercase_       =					list({doc['''title'''] for doc in model_doc if doc['''local'''] == duplicate_key}					)
								if len(snake_case__					) > 1:
												raise ValueError(
												    F'''{duplicate_key} is present several times in the documentation table of content at '''
												    '''`docs/source/en/_toctree.yml` with different *Title* values. Choose one of those and remove the '''
												    '''others.'''					)
								# Only add this once
								new_doc.append({'''local''': duplicate_key, '''title''': titles[0]}					)

				# Add none duplicate-keys
				new_doc.extend([doc for doc in model_doc if counts[doc['''local''']] == 1]					)

				# Sort
				return sorted(snake_case__						,   key=lambda snake_case__					: s["title"].lower()					)




def 				a						(							snake_case__:					List[Any]=False					):


				'''simple docstring'''
				with open(snake_case__						,   encoding='''utf-8'''					) as f:
								lowercase_       =					yaml.safe_load(f.read()					)

				# Get to the API doc
				lowercase_       =					0
				while content[api_idx]["title"] != "API":
								api_idx += 1
				lowercase_       =					content[api_idx]['''sections''']

				# Then to the model doc
				lowercase_       =					0
				while api_doc[model_idx]["title"] != "Models":
								model_idx += 1

				lowercase_       =					api_doc[model_idx]['''sections''']

				lowercase_       =					[(idx, section) for idx, section in enumerate(snake_case__					) if '''sections''' in section]
				lowercase_       =					False
				for idx, modality_doc in modalities_docs:
								lowercase_       =					modality_doc['''sections''']
								lowercase_       =					clean_model_doc_toc(snake_case__					)

								if old_modality_doc != new_modality_doc:
												lowercase_       =					True
												if overwrite:
																lowercase_       =					new_modality_doc

				if diff:
								if overwrite:
												lowercase_       =					model_doc
												lowercase_       =					api_doc
												with open(snake_case__						,   '''w'''						,   encoding='''utf-8'''					) as f:
																f.write(yaml.dump(snake_case__						,   allow_unicode=snake_case__					)					)
								else:
												raise ValueError(
												    '''The model doc part of the table of content is not properly sorted, run `make style` to fix this.'''					)


if __name__ == "__main__":
						__a  			=					argparse.ArgumentParser()
						parser.add_argument('--fix_and_overwrite', action='store_true', help='Whether to fix inconsistencies.')
						__a  			=					parser.parse_args()

						check_model_doc(args.fix_and_overwrite)




