
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__    :		Dict          =   logging.get_logger(__name__)

# General docstring
lowerCamelCase__    :		Any          =   """PoolFormerConfig"""

# Base docstring
lowerCamelCase__    :		List[str]          =   """sail/poolformer_s12"""
lowerCamelCase__    :		Union[str, Any]          =   [1, 5_1_2, 7, 7]

# Image classification docstring
lowerCamelCase__    :		List[str]          =   """sail/poolformer_s12"""
lowerCamelCase__    :		Optional[Any]          =   """tabby, tabby cat"""

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

def 		UpperCamelCase	(      lowercase_ , lowercase_ = 0.0 , lowercase_ = False   )		->	List[Any]:


		'''simple docstring'''






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



class       _snake_case   (	nn.Module							):


	def __init__(					self					,  SCREAMING_SNAKE_CASE_ = None):


			'''simple docstring'''




			super().__init__()
			lowercase__  :   List[str]    		=      drop_prob


	def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_):


			'''simple docstring'''




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



	def      lowercase__     (					self):


			'''simple docstring'''




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



class       _snake_case   (	nn.Module							):


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


			'''simple docstring'''




			super().__init__()
			lowercase__  :   int    		=      patch_size if isinstance(_a					,  collections.abc.Iterable) else (patch_size, patch_size)
			lowercase__  :   Optional[Any]    		=      stride if isinstance(_a					,  collections.abc.Iterable) else (stride, stride)
			lowercase__  :   Dict    		=      padding if isinstance(_a					,  collections.abc.Iterable) else (padding, padding)

			lowercase__  :   Optional[Any]    		=      nn.Convad(_a					,  _a					,  kernel_size=_a					,  stride=_a					,  padding=_a)
			lowercase__  :   str    		=      norm_layer(_a) if norm_layer else nn.Identity()



	def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_):


			'''simple docstring'''




			lowercase__  :   Any    		=      self.projection(_a)
			lowercase__  :   List[Any]    		=      self.norm(_a)
			return embeddings



class       _snake_case   (	nn.GroupNorm							):


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


			'''simple docstring'''




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



class       _snake_case   (	nn.Module							):


	def __init__(					self					,  SCREAMING_SNAKE_CASE_):


			'''simple docstring'''




			super().__init__()
			lowercase__  :   Union[str, Any]    		=      nn.AvgPoolad(_a					,  stride=1					,  padding=pool_size // 2					,  count_include_pad=_a)



	def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_):


			'''simple docstring'''




			return self.pool(_a) - hidden_states



class       _snake_case   (	nn.Module							):


	def __init__(					self					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_):


			'''simple docstring'''




			super().__init__()
			lowercase__  :   Tuple    		=      nn.Convad(_a					,  _a					,  1)
			lowercase__  :   Tuple    		=      nn.Convad(_a					,  _a					,  1)
			lowercase__  :   Any    		=      PoolFormerDropPath(_a)
			if isinstance(config.hidden_act					,  _a):
					lowercase__  :   Optional[int]    		=      ACTaFN[config.hidden_act]
			else:
					lowercase__  :   Optional[Any]    		=      config.hidden_act



	def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_):


			'''simple docstring'''




			lowercase__  :   Tuple    		=      self.conva(_a)
			lowercase__  :   int    		=      self.act_fn(_a)
			lowercase__  :   List[str]    		=      self.drop(_a)
			lowercase__  :   Union[str, Any]    		=      self.conva(_a)
			lowercase__  :   Optional[int]    		=      self.drop(_a)

			return hidden_states



class       _snake_case   (	nn.Module							):


	def __init__(					self					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_):


			'''simple docstring'''




			super().__init__()
			lowercase__  :   List[str]    		=      PoolFormerPooling(_a)
			lowercase__  :   int    		=      PoolFormerOutput(_a					,  _a					,  _a					,  _a)
			lowercase__  :   Optional[Any]    		=      PoolFormerGroupNorm(_a)
			lowercase__  :   Optional[Any]    		=      PoolFormerGroupNorm(_a)

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



	def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_):


			'''simple docstring'''




			if self.use_layer_scale:
					lowercase__  :   Optional[Any]    		=      self.pooling(self.before_norm(_a))
					lowercase__  :   Union[str, Any]    		=      self.layer_scale_a.unsqueeze(-1).unsqueeze(-1) * pooling_output
					# First residual connection
					lowercase__  :   Union[str, Any]    		=      hidden_states + self.drop_path(_a)
					lowercase__  :   List[Any]    		=      ()

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

					lowercase__  :   Union[str, Any]    		=      (output,) + outputs
					return outputs

			else:
					lowercase__  :   int    		=      self.drop_path(self.pooling(self.before_norm(_a)))
					# First residual connection
					lowercase__  :   int    		=      pooling_output + hidden_states
					lowercase__  :   Dict    		=      ()

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

					lowercase__  :   Dict    		=      (output,) + outputs
					return outputs



class       _snake_case   (	nn.Module							):


	def __init__(					self					,  SCREAMING_SNAKE_CASE_):


			'''simple docstring'''




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

			# patch embeddings
			lowercase__  :   Optional[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]					,  ))
			lowercase__  :   Optional[int]    		=      nn.ModuleList(_a)

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

			lowercase__  :   List[Any]    		=      nn.ModuleList(_a)



	def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_=False					,  SCREAMING_SNAKE_CASE_=True):


			'''simple docstring'''




			lowercase__  :   Tuple    		=      () if output_hidden_states else None

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

					if output_hidden_states:
							lowercase__  :   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=_a					,  hidden_states=_a)



class       _snake_case   (	UpperCamelCase__							):
	__lowerCAmelCase  :     Optional[Any]        =			PoolFormerConfig
	__lowerCAmelCase  :     int        =			'poolformer'
	__lowerCAmelCase  :     Tuple        =			'pixel_values'
	__lowerCAmelCase  :     Tuple        =			True


	def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_):


			'''simple docstring'''




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



	def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_=False):


			'''simple docstring'''




			if isinstance(_a					,  _a):
					lowercase__  :   List[str]    		=      value

lowerCamelCase__    :		List[str]          =   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"""

lowerCamelCase__    :		Union[str, Any]          =   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       _snake_case   (	UpperCamelCase__							):


	def __init__(					self					,  SCREAMING_SNAKE_CASE_):


			'''simple docstring'''




			super().__init__(_a)
			lowercase__  :   Any    		=      config

			lowercase__  :   Union[str, Any]    		=      PoolFormerEncoder(_a)

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


	def      lowercase__     (					self):


			'''simple docstring'''




			return self.embeddings.patch_embeddings



	@add_start_docstrings_to_model_forward(_a)
	@add_code_sample_docstrings(
	    checkpoint=_CHECKPOINT_FOR_DOC					,  output_type=_a					,  config_class=_CONFIG_FOR_DOC					,  modality="""vision"""					,  expected_output=_EXPECTED_OUTPUT_SHAPE					,  )
	def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_ = None					,  SCREAMING_SNAKE_CASE_ = None					,  SCREAMING_SNAKE_CASE_ = None					,  ):


			'''simple docstring'''




			lowercase__  :   Optional[int]    		=      (
			    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
			)
			lowercase__  :   List[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""")

			lowercase__  :   Any    		=      self.encoder(
			    _a					,  output_hidden_states=_a					,  return_dict=_a					,  )
			lowercase__  :   Optional[int]    		=      encoder_outputs[0]

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

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



class       _snake_case   (	nn.Module							):


	def __init__(					self					,  SCREAMING_SNAKE_CASE_):


			'''simple docstring'''




			super().__init__()
			lowercase__  :   List[str]    		=      nn.Linear(config.hidden_size					,  config.hidden_size)



	def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_):


			'''simple docstring'''




			lowercase__  :   str    		=      self.dense(_a)
			return output



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


	def __init__(					self					,  SCREAMING_SNAKE_CASE_):


			'''simple docstring'''




			super().__init__(_a)
			lowercase__  :   Tuple    		=      config.num_labels
			lowercase__  :   Optional[Any]    		=      PoolFormerModel(_a)

			# Final norm
			lowercase__  :   Optional[int]    		=      PoolFormerGroupNorm(config.hidden_sizes[-1])
			# Classifier head
			lowercase__  :   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(_a)
	@add_code_sample_docstrings(
	    checkpoint=_IMAGE_CLASS_CHECKPOINT					,  output_type=_a					,  config_class=_CONFIG_FOR_DOC					,  expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT					,  )
	def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_ = None					,  SCREAMING_SNAKE_CASE_ = None					,  SCREAMING_SNAKE_CASE_ = None					,  SCREAMING_SNAKE_CASE_ = None					,  ):


			'''simple docstring'''




			lowercase__  :   Optional[Any]    		=      return_dict if return_dict is not None else self.config.use_return_dict

			lowercase__  :   Optional[Any]    		=      self.poolformer(
			    _a					,  output_hidden_states=_a					,  return_dict=_a					,  )

			lowercase__  :   Tuple    		=      outputs[0]

			lowercase__  :   Any    		=      self.classifier(self.norm(_a).mean([-2, -1]))

			lowercase__  :   Dict    		=      None
			if labels is not None:
					if self.config.problem_type is None:
							if self.num_labels == 1:
									lowercase__  :   Union[str, Any]    		=      """regression"""
							elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
									lowercase__  :   Optional[int]    		=      """single_label_classification"""
							else:
									lowercase__  :   List[Any]    		=      """multi_label_classification"""

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

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

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








from scipy.stats import spearmanr

import datasets


lowerCamelCase__    :		Union[str, Any]          =   """
The Spearman rank-order correlation coefficient is a measure of the
relationship between two datasets. Like other correlation coefficients,
this one varies between -1 and +1 with 0 implying no correlation.
Positive correlations imply that as data in dataset x increases, so
does data in dataset y. Negative correlations imply that as x increases,
y decreases. Correlations of -1 or +1 imply an exact monotonic relationship.

Unlike the Pearson correlation, the Spearman correlation does not
assume that both datasets are normally distributed.

The p-value roughly indicates the probability of an uncorrelated system
producing datasets that have a Spearman correlation at least as extreme
as the one computed from these datasets. The p-values are not entirely
reliable but are probably reasonable for datasets larger than 500 or so.
"""

lowerCamelCase__    :		int          =   """
Args:
    predictions (`List[float]`): Predicted labels, as returned by a model.
    references (`List[float]`): Ground truth labels.
    return_pvalue (`bool`): If `True`, returns the p-value. If `False`, returns
            only the spearmanr score. Defaults to `False`.
Returns:
    spearmanr (`float`): Spearman correlation coefficient.
    p-value (`float`): p-value. **Note**: is only returned if `return_pvalue=True` is input.
Examples:
    Example 1:
        >>> spearmanr_metric = datasets.load_metric(\"spearmanr\")
        >>> results = spearmanr_metric.compute(references=[1, 2, 3, 4, 5], predictions=[10, 9, 2.5, 6, 4])
        >>> print(results)
        {'spearmanr': -0.7}

    Example 2:
        >>> spearmanr_metric = datasets.load_metric(\"spearmanr\")
        >>> results = spearmanr_metric.compute(references=[1, 2, 3, 4, 5],
        ...                                     predictions=[10, 9, 2.5, 6, 4],
        ...                                     return_pvalue=True)
        >>> print(results['spearmanr'])
        -0.7
        >>> print(round(results['spearmanr_pvalue'], 2))
        0.19
"""

lowerCamelCase__    :		str          =   R"""\
@book{kokoska2000crc,
  title={CRC standard probability and statistics tables and formulae},
  author={Kokoska, Stephen and Zwillinger, Daniel},
  year={2000},
  publisher={Crc Press}
}
@article{2020SciPy-NMeth,
  author  = {Virtanen, Pauli and Gommers, Ralf and Oliphant, Travis E. and
            Haberland, Matt and Reddy, Tyler and Cournapeau, David and
            Burovski, Evgeni and Peterson, Pearu and Weckesser, Warren and
            Bright, Jonathan and {van der Walt}, St{\'e}fan J. and
            Brett, Matthew and Wilson, Joshua and Millman, K. Jarrod and
            Mayorov, Nikolay and Nelson, Andrew R. J. and Jones, Eric and
            Kern, Robert and Larson, Eric and Carey, C J and
            Polat, {\.I}lhan and Feng, Yu and Moore, Eric W. and
            {VanderPlas}, Jake and Laxalde, Denis and Perktold, Josef and
            Cimrman, Robert and Henriksen, Ian and Quintero, E. A. and
            Harris, Charles R. and Archibald, Anne M. and
            Ribeiro, Ant{\^o}nio H. and Pedregosa, Fabian and
            {van Mulbregt}, Paul and {SciPy 1.0 Contributors}},
  title   = {{{SciPy} 1.0: Fundamental Algorithms for Scientific
            Computing in Python}},
  journal = {Nature Methods},
  year    = {2020},
  volume  = {17},
  pages   = {261--272},
  adsurl  = {https://rdcu.be/b08Wh},
  doi     = {10.1038/s41592-019-0686-2},
}
"""



@datasets.utils.file_utils.add_start_docstrings(_DESCRIPTION  ,						_KWARGS_DESCRIPTION							)
class       _snake_case   (	datasets.Metric							):


 def      lowercase__     (					self):


   '''simple docstring'''




   return datasets.MetricInfo(
       description=_DESCRIPTION					,  citation=_CITATION					,  inputs_description=_KWARGS_DESCRIPTION					,  features=datasets.Features(
           {
               """predictions""": datasets.Value("""float"""),
               """references""": datasets.Value("""float"""),
           })					,  reference_urls=["""https://docs.scipy.org/doc/scipy/reference/generated/scipy.stats.spearmanr.html"""]					,  )



 def      lowercase__     (					self					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_=False):


   '''simple docstring'''




   lowercase__  :   List[str]    		=      spearmanr(SCREAMING_SNAKE_CASE_					,  SCREAMING_SNAKE_CASE_)
   if return_pvalue:
     return {"spearmanr": results[0], "spearmanr_pvalue": results[1]}
   else:
     return {"spearmanr": results[0]}






