
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


SCREAMING_SNAKE_CASE							:int										=    logging.get_logger(__name__)

# General docstring
SCREAMING_SNAKE_CASE							:int										=    'PoolFormerConfig'

# Base docstring
SCREAMING_SNAKE_CASE							:Tuple										=    'sail/poolformer_s12'
SCREAMING_SNAKE_CASE							:List[Any]										=    [1, 512, 7, 7]

# Image classification docstring
SCREAMING_SNAKE_CASE							:Optional[Any]										=    'sail/poolformer_s12'
SCREAMING_SNAKE_CASE							:Union[str, Any]										=    'tabby, tabby cat'

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


def 						UpperCAmelCase      (     a_		,       a_ = 0.0		,       a_ = False						)		->     Union[str, 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(__UpperCamelCase		,       dtype=input.dtype		,       device=input.device						)
							random_tensor.floor_()  # binarize
							__A             =			input.div(__UpperCamelCase						) * random_tensor
							return output





class 							UpperCAmelCase				(  nn.Module  ):




	'''simple docstring'''





	def __init__(	self :			Tuple							,A :			Optional[float] = None			):
								super().__init__()
								__A             =			drop_prob






	def     UpperCamelCase_  (	self :			Dict							,A :			torch.Tensor			):
								return drop_path(UpperCamelCase__							,self.drop_prob							,self.training			)







	def     UpperCamelCase_  (	self :			Any			):
								return "p={}".format(self.drop_prob			)





class 							UpperCAmelCase				(  nn.Module  ):




	'''simple docstring'''





	def __init__(	self :			str							,A :			Optional[Any]							,A :			List[Any]							,A :			List[str]							,A :			List[str]							,A :			Dict							,A :			Tuple=None			):
								super().__init__()
								__A             =			patch_size if isinstance(UpperCamelCase__							,collections.abc.Iterable			) else (patch_size, patch_size)
								__A             =			stride if isinstance(UpperCamelCase__							,collections.abc.Iterable			) else (stride, stride)
								__A             =			padding if isinstance(UpperCamelCase__							,collections.abc.Iterable			) else (padding, padding)

								__A             =			nn.Convad(UpperCamelCase__							,UpperCamelCase__							,kernel_size=UpperCamelCase__							,stride=UpperCamelCase__							,padding=UpperCamelCase__			)
								__A             =			norm_layer(UpperCamelCase__			) if norm_layer else nn.Identity()







	def     UpperCamelCase_  (	self :			Union[str, Any]							,A :			int			):
								__A             =			self.projection(UpperCamelCase__			)
								__A             =			self.norm(UpperCamelCase__			)
								return embeddings





class 							UpperCAmelCase				(  nn.GroupNorm  ):




	'''simple docstring'''





	def __init__(	self :			Union[str, Any]							,A :			Optional[int]							,**A :			Dict			):
								super().__init__(1							,UpperCamelCase__							,**UpperCamelCase__			)





class 							UpperCAmelCase				(  nn.Module  ):




	'''simple docstring'''





	def __init__(	self :			Optional[int]							,A :			int			):
								super().__init__()
								__A             =			nn.AvgPoolad(UpperCamelCase__							,stride=1							,padding=pool_size // 2							,count_include_pad=UpperCamelCase__			)







	def     UpperCamelCase_  (	self :			Union[str, Any]							,A :			Union[str, Any]			):
								return self.pool(UpperCamelCase__			) - hidden_states





class 							UpperCAmelCase				(  nn.Module  ):




	'''simple docstring'''





	def __init__(	self :			List[str]							,A :			List[str]							,A :			Optional[int]							,A :			Dict							,A :			int			):
								super().__init__()
								__A             =			nn.Convad(UpperCamelCase__							,UpperCamelCase__							,1			)
								__A             =			nn.Convad(UpperCamelCase__							,UpperCamelCase__							,1			)
								__A             =			PoolFormerDropPath(UpperCamelCase__			)
								if isinstance(config.hidden_act							,UpperCamelCase__			):
															__A             =			ACTaFN[config.hidden_act]
								else:
															__A             =			config.hidden_act







	def     UpperCamelCase_  (	self :			Any							,A :			Optional[Any]			):
								__A             =			self.conva(UpperCamelCase__			)
								__A             =			self.act_fn(UpperCamelCase__			)
								__A             =			self.drop(UpperCamelCase__			)
								__A             =			self.conva(UpperCamelCase__			)
								__A             =			self.drop(UpperCamelCase__			)

								return hidden_states





class 							UpperCAmelCase				(  nn.Module  ):




	'''simple docstring'''





	def __init__(	self :			List[Any]							,A :			str							,A :			str							,A :			Union[str, Any]							,A :			Optional[Any]							,A :			Optional[int]							,A :			Tuple			):
								super().__init__()
								__A             =			PoolFormerPooling(UpperCamelCase__			)
								__A             =			PoolFormerOutput(UpperCamelCase__							,UpperCamelCase__							,UpperCamelCase__							,UpperCamelCase__			)
								__A             =			PoolFormerGroupNorm(UpperCamelCase__			)
								__A             =			PoolFormerGroupNorm(UpperCamelCase__			)

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







	def     UpperCamelCase_  (	self :			Union[str, Any]							,A :			Optional[int]			):



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

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

															__A             =			(output,) + outputs
															return outputs

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

															__A             =			(output,) + outputs
															return outputs





class 							UpperCAmelCase				(  nn.Module  ):




	'''simple docstring'''





	def __init__(	self :			List[Any]							,A :			Tuple			):
								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(UpperCamelCase__			)

								# 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(
																						        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             =			nn.ModuleList(UpperCamelCase__			)







	def     UpperCamelCase_  (	self :			str							,A :			Optional[Any]							,A :			Union[str, Any]=False							,A :			Optional[Any]=True			):
								__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(UpperCamelCase__			)
															# Send the embeddings through the blocks
															for _, blk in enumerate(UpperCamelCase__			):
																						__A             =			blk(UpperCamelCase__			)
																						__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=UpperCamelCase__							,hidden_states=UpperCamelCase__			)





class 							UpperCAmelCase				(  __a  ):




	'''simple docstring'''





	snake_case_    =					PoolFormerConfig
	snake_case_    =					"poolformer"
	snake_case_    =					"pixel_values"
	snake_case_    =					True






	def     UpperCamelCase_  (	self :			Optional[Any]							,A :			Tuple			):



								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     UpperCamelCase_  (	self :			str							,A :			Optional[Any]							,A :			Dict=False			):



								if isinstance(UpperCamelCase__							,UpperCamelCase__			):
															__A             =			value





SCREAMING_SNAKE_CASE							:Optional[int]										=    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'

SCREAMING_SNAKE_CASE							:Tuple										=    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."      ,      __a      ,      )
class 							UpperCAmelCase				(  __a  ):




	'''simple docstring'''





	def __init__(	self :			List[Any]							,A :			Optional[int]			):
								super().__init__(UpperCamelCase__			)
								__A             =			config

								__A             =			PoolFormerEncoder(UpperCamelCase__			)

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






	def     UpperCamelCase_  (	self :			Optional[Any]			):
								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     UpperCamelCase_  (	self :			Dict							,A :			Optional[torch.FloatTensor] = None							,A :			Optional[bool] = None							,A :			Optional[bool] = None							,):
								__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(
								    UpperCamelCase__							,output_hidden_states=UpperCamelCase__							,return_dict=UpperCamelCase__							,)
								__A             =			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 							UpperCAmelCase				(  nn.Module  ):




	'''simple docstring'''





	def __init__(	self :			str							,A :			Union[str, Any]			):
								super().__init__()
								__A             =			nn.Linear(config.hidden_size							,config.hidden_size			)







	def     UpperCamelCase_  (	self :			List[Any]							,A :			Any			):
								__A             =			self.dense(UpperCamelCase__			)
								return output





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




	'''simple docstring'''





	def __init__(	self :			Optional[Any]							,A :			Any			):
								super().__init__(UpperCamelCase__			)
								__A             =			config.num_labels
								__A             =			PoolFormerModel(UpperCamelCase__			)

								# 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(UpperCamelCase__			)
	@add_code_sample_docstrings(
	    checkpoint=_IMAGE_CLASS_CHECKPOINT							,output_type=UpperCamelCase__							,config_class=_CONFIG_FOR_DOC							,expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT							,)
	def     UpperCamelCase_  (	self :			Dict							,A :			Optional[torch.FloatTensor] = None							,A :			Optional[torch.LongTensor] = None							,A :			Optional[bool] = None							,A :			Optional[bool] = None							,):
								__A             =			return_dict if return_dict is not None else self.config.use_return_dict

								__A             =			self.poolformer(
								    UpperCamelCase__							,output_hidden_states=UpperCamelCase__							,return_dict=UpperCamelCase__							,)

								__A             =			outputs[0]

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

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

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


from typing import TYPE_CHECKING

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


SCREAMING_SNAKE_CASE							:Any										=    {
    'configuration_rag': ['RagConfig'],
    'retrieval_rag': ['RagRetriever'],
    'tokenization_rag': ['RagTokenizer'],
}

try:
					if not is_torch_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					SCREAMING_SNAKE_CASE							:List[Any]										=    [
					    'RagModel',
					    'RagPreTrainedModel',
					    'RagSequenceForGeneration',
					    'RagTokenForGeneration',
					]

try:
					if not is_tf_available():
										raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
					pass
else:
					SCREAMING_SNAKE_CASE							:Any										=    [
					    'TFRagModel',
					    'TFRagPreTrainedModel',
					    'TFRagSequenceForGeneration',
					    'TFRagTokenForGeneration',
					]


if TYPE_CHECKING:
					from .configuration_rag import RagConfig
					from .retrieval_rag import RagRetriever
					from .tokenization_rag import RagTokenizer

					try:
										if not is_torch_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_rag import RagModel, RagPreTrainedModel, RagSequenceForGeneration, RagTokenForGeneration

					try:
										if not is_tf_available():
															raise OptionalDependencyNotAvailable()
					except OptionalDependencyNotAvailable:
										pass
					else:
										from .modeling_tf_rag import (
										    TFRagModel,
										    TFRagPreTrainedModel,
										    TFRagSequenceForGeneration,
										    TFRagTokenForGeneration,
										)

else:
					import sys

					SCREAMING_SNAKE_CASE							:int										=    _LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)
