


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

# General docstring
lowercase  							=					'''PoolFormerConfig'''

# Base docstring
lowercase  							=					'''sail/poolformer_s12'''
lowercase  							=					[1, 5_12, 7, 7]

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

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




def 	UpperCAmelCase     (		A	:						str     ,				A	:						float = 0.0     ,				A	:						bool = False       ):

		'''simple docstring'''


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




class 			lowercase__			(      nn.Module    ):
							'''simple docstring'''







							def __init__(  self	,       snake_case = None						)			->							List[Any]:
									super().__init__()
									_UpperCAmelCase							 =		drop_prob






							def        lowerCamelCase_   (  self	,       snake_case						)			->							str:
									return drop_path(snake_case	,       self.drop_prob	,       self.training						)






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







class 			lowercase__			(      nn.Module    ):
							'''simple docstring'''







							def __init__(  self	,       snake_case	,       snake_case	,       snake_case	,       snake_case	,       snake_case	,       snake_case=None						)			->							Any:
									super().__init__()
									_UpperCAmelCase							 =		patch_size if isinstance(snake_case	,       collections.abc.Iterable						) else (patch_size, patch_size)
									_UpperCAmelCase							 =		stride if isinstance(snake_case	,       collections.abc.Iterable						) else (stride, stride)
									_UpperCAmelCase							 =		padding if isinstance(snake_case	,       collections.abc.Iterable						) else (padding, padding)

									_UpperCAmelCase							 =		nn.Convad(snake_case	,       snake_case	,       kernel_size=snake_case	,       stride=snake_case	,       padding=snake_case						)
									_UpperCAmelCase							 =		norm_layer(snake_case						) if norm_layer else nn.Identity()






							def        lowerCamelCase_   (  self	,       snake_case						)			->							Optional[int]:
									_UpperCAmelCase							 =		self.projection(snake_case						)
									_UpperCAmelCase							 =		self.norm(snake_case						)
									return embeddings







class 			lowercase__			(      nn.GroupNorm    ):
							'''simple docstring'''







							def __init__(  self	,       snake_case	,       **snake_case						)			->							str:
									super().__init__(1	,       snake_case	,       **snake_case						)







class 			lowercase__			(      nn.Module    ):
							'''simple docstring'''







							def __init__(  self	,       snake_case						)			->							int:
									super().__init__()
									_UpperCAmelCase							 =		nn.AvgPoolad(snake_case	,       stride=1	,       padding=pool_size // 2	,       count_include_pad=snake_case						)






							def        lowerCamelCase_   (  self	,       snake_case						)			->							Union[str, Any]:
									return self.pool(snake_case						) - hidden_states







class 			lowercase__			(      nn.Module    ):
							'''simple docstring'''







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






							def        lowerCamelCase_   (  self	,       snake_case						)			->							List[str]:
									_UpperCAmelCase							 =		self.conva(snake_case						)
									_UpperCAmelCase							 =		self.act_fn(snake_case						)
									_UpperCAmelCase							 =		self.drop(snake_case						)
									_UpperCAmelCase							 =		self.conva(snake_case						)
									_UpperCAmelCase							 =		self.drop(snake_case						)

									return hidden_states







class 			lowercase__			(      nn.Module    ):
							'''simple docstring'''







							def __init__(  self	,       snake_case	,       snake_case	,       snake_case	,       snake_case	,       snake_case	,       snake_case						)			->							Dict:
									super().__init__()
									_UpperCAmelCase							 =		PoolFormerPooling(snake_case						)
									_UpperCAmelCase							 =		PoolFormerOutput(snake_case	,       snake_case	,       snake_case	,       snake_case						)
									_UpperCAmelCase							 =		PoolFormerGroupNorm(snake_case						)
									_UpperCAmelCase							 =		PoolFormerGroupNorm(snake_case						)

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






							def        lowerCamelCase_   (  self	,       snake_case						)			->							Dict:





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

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

											_UpperCAmelCase							 =		(output,) + outputs
											return outputs

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

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

											_UpperCAmelCase							 =		(output,) + outputs
											return outputs







class 			lowercase__			(      nn.Module    ):
							'''simple docstring'''







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

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

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

									_UpperCAmelCase							 =		nn.ModuleList(snake_case						)






							def        lowerCamelCase_   (  self	,       snake_case	,       snake_case=False	,       snake_case=True						)			->							str:
									_UpperCAmelCase							 =		() if output_hidden_states else None

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

											if output_hidden_states:
													_UpperCAmelCase							 =		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=snake_case	,       hidden_states=snake_case						)







class 			lowercase__			(      __A    ):
							'''simple docstring'''







							_UpperCAmelCase      				=					PoolFormerConfig
							_UpperCAmelCase      				=					'''poolformer'''
							_UpperCAmelCase      				=					'''pixel_values'''
							_UpperCAmelCase      				=					True






							def        lowerCamelCase_   (  self	,       snake_case						)			->							Any:





									if isinstance(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(snake_case	,       nn.LayerNorm						):
											module.bias.data.zero_()
											module.weight.data.fill_(1.0						)






							def        lowerCamelCase_   (  self	,       snake_case	,       snake_case=False						)			->							Dict:





									if isinstance(snake_case	,       snake_case						):
											_UpperCAmelCase							 =		value







lowercase  							=					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.
'''

lowercase  							=					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 			lowercase__			(      __A    ):
							'''simple docstring'''







							def __init__(  self	,       snake_case						)			->							Dict:
									super().__init__(snake_case						)
									_UpperCAmelCase							 =		config

									_UpperCAmelCase							 =		PoolFormerEncoder(snake_case						)

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






							def        lowerCamelCase_   (  self						)			->							List[Any]:
									return self.embeddings.patch_embeddings






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

									_UpperCAmelCase							 =		self.encoder(
									    snake_case	,       output_hidden_states=snake_case	,       return_dict=snake_case	,       )
									_UpperCAmelCase							 =		encoder_outputs[0]

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

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







class 			lowercase__			(      nn.Module    ):
							'''simple docstring'''







							def __init__(  self	,       snake_case						)			->							str:
									super().__init__()
									_UpperCAmelCase							 =		nn.Linear(config.hidden_size	,       config.hidden_size						)






							def        lowerCamelCase_   (  self	,       snake_case						)			->							List[Any]:
									_UpperCAmelCase							 =		self.dense(snake_case						)
									return output







@add_start_docstrings(
    '''
    PoolFormer Model transformer with an image classification head on top
    ''',					__A,					)
class 			lowercase__			(      __A    ):
							'''simple docstring'''







							def __init__(  self	,       snake_case						)			->							Any:
									super().__init__(snake_case						)
									_UpperCAmelCase							 =		config.num_labels
									_UpperCAmelCase							 =		PoolFormerModel(snake_case						)

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

									_UpperCAmelCase							 =		self.poolformer(
									    snake_case	,       output_hidden_states=snake_case	,       return_dict=snake_case	,       )

									_UpperCAmelCase							 =		outputs[0]

									_UpperCAmelCase							 =		self.classifier(self.norm(snake_case						).mean([-2, -1]						)						)

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

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

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

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







"""simple docstring"""





from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tokenizers_available, is_torch_available


lowercase  							=					{
    '''configuration_mvp''': ['''MVP_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''MvpConfig''', '''MvpOnnxConfig'''],
    '''tokenization_mvp''': ['''MvpTokenizer'''],
}

try:
					if not is_tokenizers_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowercase  							=					['''MvpTokenizerFast''']

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					lowercase  							=					[
					    '''MVP_PRETRAINED_MODEL_ARCHIVE_LIST''',
					    '''MvpForCausalLM''',
					    '''MvpForConditionalGeneration''',
					    '''MvpForQuestionAnswering''',
					    '''MvpForSequenceClassification''',
					    '''MvpModel''',
					    '''MvpPreTrainedModel''',
					]

if TYPE_CHECKING:
					from .configuration_mvp import MVP_PRETRAINED_CONFIG_ARCHIVE_MAP, MvpConfig, MvpOnnxConfig
					from .tokenization_mvp import MvpTokenizer

					try:
										if not is_tokenizers_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .tokenization_mvp_fast import MvpTokenizerFast

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_mvp import (
										    MVP_PRETRAINED_MODEL_ARCHIVE_LIST,
										    MvpForCausalLM,
										    MvpForConditionalGeneration,
										    MvpForQuestionAnswering,
										    MvpForSequenceClassification,
										    MvpModel,
										    MvpPreTrainedModel,
										)

else:
					import sys

					lowercase  							=					_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)


