



"""simple docstring"""






import collections.abc
from typing import Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithNoAttention, ImageClassifierOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_poolformer import PoolFormerConfig


_lowercase			:			int										=						logging.get_logger(__name__)

# General docstring
_lowercase			:			List[Any]										=						"""PoolFormerConfig"""

# Base docstring
_lowercase			:			Optional[int]										=						"""sail/poolformer_s12"""
_lowercase			:			Union[str, Any]										=						[1, 5_1_2, 7, 7]

# Image classification docstring
_lowercase			:			Tuple										=						"""sail/poolformer_s12"""
_lowercase			:			Tuple										=						"""tabby, tabby cat"""

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






def 					snake_case__       (		__lowerCamelCase    :   List[str]							,		__lowerCamelCase    :   List[str] = 0.0							,		__lowerCamelCase    :   str = False     ):
			"""simple docstring"""
			if drop_prob == 0.0 or not training:
						return input
			lowerCamelCase__		:		List[str]		     =1 - drop_prob
			lowerCamelCase__		:		Any		     =(input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
			lowerCamelCase__		:		Dict		     =keep_prob + torch.rand(__lowerCamelCase							,		dtype=input.dtype							,		device=input.device     )
			random_tensor.floor_()  # binarize
			lowerCamelCase__		:		List[str]		     =input.div(__lowerCamelCase     ) * random_tensor
			return output

class 	__SCREAMING_SNAKE_CASE     (				nn.Module       ):


					'''simple docstring'''





					def __init__(	self							:  Any,  lowerCamelCase							:  Union[str, Any] = None					)-> Dict:
								super().__init__()
								lowerCamelCase__		:		Optional[int]		     =drop_prob



					def 		snake_case	(	self							:  List[str],  lowerCamelCase							:  Tuple					)-> str:
								return drop_path(__lowercase,  self.drop_prob,  self.training					)




					def 		snake_case	(	self							:  Union[str, Any]					)-> List[str]:
								return "p={}".format(self.drop_prob					)





class 	__SCREAMING_SNAKE_CASE     (				nn.Module       ):


					'''simple docstring'''





					def __init__(	self							:  List[Any],  lowerCamelCase							:  Dict,  lowerCamelCase							:  int,  lowerCamelCase							:  Union[str, Any],  lowerCamelCase							:  Any,  lowerCamelCase							:  Any,  lowerCamelCase							:  Optional[Any]=None					)-> str:
								super().__init__()
								lowerCamelCase__		:		Optional[Any]		     =patch_size if isinstance(__lowercase,  collections.abc.Iterable					) else (patch_size, patch_size)
								lowerCamelCase__		:		int		     =stride if isinstance(__lowercase,  collections.abc.Iterable					) else (stride, stride)
								lowerCamelCase__		:		Dict		     =padding if isinstance(__lowercase,  collections.abc.Iterable					) else (padding, padding)

								lowerCamelCase__		:		str		     =nn.Convad(__lowercase,  __lowercase,  kernel_size=__lowercase,  stride=__lowercase,  padding=__lowercase					)
								lowerCamelCase__		:		Tuple		     =norm_layer(__lowercase					) if norm_layer else nn.Identity()




					def 		snake_case	(	self							:  Optional[int],  lowerCamelCase							:  Any					)-> Dict:
								lowerCamelCase__		:		Union[str, Any]		     =self.projection(__lowercase					)
								lowerCamelCase__		:		Dict		     =self.norm(__lowercase					)
								return embeddings





class 	__SCREAMING_SNAKE_CASE     (				nn.GroupNorm       ):


					'''simple docstring'''





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





class 	__SCREAMING_SNAKE_CASE     (				nn.Module       ):


					'''simple docstring'''





					def __init__(	self							:  List[str],  lowerCamelCase							:  Optional[Any]					)-> List[str]:
								super().__init__()
								lowerCamelCase__		:		List[Any]		     =nn.AvgPoolad(__lowercase,  stride=1,  padding=pool_size // 2,  count_include_pad=__lowercase					)




					def 		snake_case	(	self							:  int,  lowerCamelCase							:  str					)-> Optional[Any]:
								return self.pool(__lowercase					) - hidden_states





class 	__SCREAMING_SNAKE_CASE     (				nn.Module       ):


					'''simple docstring'''





					def __init__(	self							:  Dict,  lowerCamelCase							:  Optional[Any],  lowerCamelCase							:  int,  lowerCamelCase							:  Optional[int],  lowerCamelCase							:  Optional[int]					)-> int:
								super().__init__()
								lowerCamelCase__		:		Tuple		     =nn.Convad(__lowercase,  __lowercase,  1					)
								lowerCamelCase__		:		Tuple		     =nn.Convad(__lowercase,  __lowercase,  1					)
								lowerCamelCase__		:		Optional[int]		     =PoolFormerDropPath(__lowercase					)
								if isinstance(config.hidden_act,  __lowercase					):
											lowerCamelCase__		:		Tuple		     =ACTaFN[config.hidden_act]
								else:
											lowerCamelCase__		:		Any		     =config.hidden_act




					def 		snake_case	(	self							:  Optional[int],  lowerCamelCase							:  Optional[int]					)-> int:
								lowerCamelCase__		:		Optional[int]		     =self.conva(__lowercase					)
								lowerCamelCase__		:		Dict		     =self.act_fn(__lowercase					)
								lowerCamelCase__		:		int		     =self.drop(__lowercase					)
								lowerCamelCase__		:		Dict		     =self.conva(__lowercase					)
								lowerCamelCase__		:		str		     =self.drop(__lowercase					)

								return hidden_states





class 	__SCREAMING_SNAKE_CASE     (				nn.Module       ):


					'''simple docstring'''





					def __init__(	self							:  Tuple,  lowerCamelCase							:  Any,  lowerCamelCase							:  Tuple,  lowerCamelCase							:  str,  lowerCamelCase							:  Tuple,  lowerCamelCase							:  Tuple,  lowerCamelCase							:  List[Any]					)-> Tuple:
								super().__init__()
								lowerCamelCase__		:		List[str]		     =PoolFormerPooling(__lowercase					)
								lowerCamelCase__		:		Dict		     =PoolFormerOutput(__lowercase,  __lowercase,  __lowercase,  __lowercase					)
								lowerCamelCase__		:		str		     =PoolFormerGroupNorm(__lowercase					)
								lowerCamelCase__		:		List[str]		     =PoolFormerGroupNorm(__lowercase					)

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




					def 		snake_case	(	self							:  Dict,  lowerCamelCase							:  List[Any]					)-> Union[str, Any]:
								if self.use_layer_scale:
											lowerCamelCase__		:		Optional[int]		     =self.pooling(self.before_norm(__lowercase					)					)
											lowerCamelCase__		:		Union[str, Any]		     =self.layer_scale_a.unsqueeze(-1					).unsqueeze(-1					) * pooling_output
											# First residual connection
											lowerCamelCase__		:		Dict		     =hidden_states + self.drop_path(__lowercase					)
											lowerCamelCase__		:		str		     =()

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

											lowerCamelCase__		:		int		     =(output,) + outputs
											return outputs

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

											# Second residual connection inside the PoolFormerOutput block
											lowerCamelCase__		:		Optional[int]		     =self.drop_path(self.output(self.after_norm(__lowercase					)					)					)
											lowerCamelCase__		:		Union[str, Any]		     =hidden_states + layer_output

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





class 	__SCREAMING_SNAKE_CASE     (				nn.Module       ):


					'''simple docstring'''





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

								# patch embeddings
								lowerCamelCase__		:		Union[str, 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],  )					)
								lowerCamelCase__		:		Any		     =nn.ModuleList(__lowercase					)

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

								lowerCamelCase__		:		Dict		     =nn.ModuleList(__lowercase					)




					def 		snake_case	(	self							:  Optional[int],  lowerCamelCase							:  List[Any],  lowerCamelCase							:  Union[str, Any]=False,  lowerCamelCase							:  Tuple=True					)-> Union[str, Any]:
								lowerCamelCase__		:		str		     =() if output_hidden_states else None

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

											if output_hidden_states:
														lowerCamelCase__		:		Union[str, 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=__lowercase,  hidden_states=__lowercase					)





class 	__SCREAMING_SNAKE_CASE     (				__A       ):


					'''simple docstring'''





					_a             =							PoolFormerConfig
					_a             =							'poolformer'
					_a             =							'pixel_values'
					_a             =							True



					def 		snake_case	(	self							:  str,  lowerCamelCase							:  Optional[int]					)-> Any:
								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 		snake_case	(	self							:  Any,  lowerCamelCase							:  Optional[Any],  lowerCamelCase							:  List[Any]=False					)-> Tuple:
								if isinstance(__lowercase,  __lowercase					):
											lowerCamelCase__		:		List[str]		     =value





_lowercase			:			Any										=						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			:			List[Any]										=						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 	__SCREAMING_SNAKE_CASE     (				__A       ):


					'''simple docstring'''





					def __init__(	self							:  str,  lowerCamelCase							:  Tuple					)-> Dict:
								super().__init__(__lowercase					)
								lowerCamelCase__		:		List[Any]		     =config

								lowerCamelCase__		:		List[str]		     =PoolFormerEncoder(__lowercase					)

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



					def 		snake_case	(	self							:  Any					)-> List[Any]:
								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 		snake_case	(	self							:  Union[str, Any],  lowerCamelCase							:  str = None,  lowerCamelCase							:  Tuple = None,  lowerCamelCase							:  Any = None,  )-> List[str]:
								lowerCamelCase__		:		Dict		     =(
								    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
								)
								lowerCamelCase__		:		Union[str, Any]		     =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'''					)

								lowerCamelCase__		:		List[Any]		     =self.encoder(
								    __lowercase,  output_hidden_states=__lowercase,  return_dict=__lowercase,  )
								lowerCamelCase__		:		Union[str, Any]		     =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 	__SCREAMING_SNAKE_CASE     (				nn.Module       ):


					'''simple docstring'''





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




					def 		snake_case	(	self							:  List[str],  lowerCamelCase							:  Union[str, Any]					)-> Dict:
								lowerCamelCase__		:		int		     =self.dense(__lowercase					)
								return output





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


					'''simple docstring'''





					def __init__(	self							:  Any,  lowerCamelCase							:  Union[str, Any]					)-> List[Any]:
								super().__init__(__lowercase					)
								lowerCamelCase__		:		Optional[Any]		     =config.num_labels
								lowerCamelCase__		:		str		     =PoolFormerModel(__lowercase					)

								# Final norm
								lowerCamelCase__		:		List[Any]		     =PoolFormerGroupNorm(config.hidden_sizes[-1]					)
								# Classifier head
								lowerCamelCase__		:		str		     =(
								    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 		snake_case	(	self							:  List[str],  lowerCamelCase							:  Union[str, Any] = None,  lowerCamelCase							:  int = None,  lowerCamelCase							:  Dict = None,  lowerCamelCase							:  List[str] = None,  )-> List[str]:
								lowerCamelCase__		:		Any		     =return_dict if return_dict is not None else self.config.use_return_dict

								lowerCamelCase__		:		Optional[Any]		     =self.poolformer(
								    __lowercase,  output_hidden_states=__lowercase,  return_dict=__lowercase,  )

								lowerCamelCase__		:		int		     =outputs[0]

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

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

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

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

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









"""simple docstring"""






from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_sentencepiece_available,
    is_speech_available,
    is_torch_available,
)


_lowercase			:			Any										=						{
    "configuration_trocr": ["TROCR_PRETRAINED_CONFIG_ARCHIVE_MAP", "TrOCRConfig"],
    "processing_trocr": ["TrOCRProcessor"],
}


try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		_lowercase			:			List[Any]										=						[
		    "TROCR_PRETRAINED_MODEL_ARCHIVE_LIST",
		    "TrOCRForCausalLM",
		    "TrOCRPreTrainedModel",
		]


if TYPE_CHECKING:
		from .configuration_trocr import TROCR_PRETRAINED_CONFIG_ARCHIVE_MAP, TrOCRConfig
		from .processing_trocr import TrOCRProcessor

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_trocr import TROCR_PRETRAINED_MODEL_ARCHIVE_LIST, TrOCRForCausalLM, TrOCRPreTrainedModel

else:
		import sys

		_lowercase			:			Tuple										=						_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)




