


"""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


__SCREAMING_SNAKE_CASE      		=    logging.get_logger(__name__)

# General docstring
__SCREAMING_SNAKE_CASE      		=    """PoolFormerConfig"""

# Base docstring
__SCREAMING_SNAKE_CASE      		=    """sail/poolformer_s12"""
__SCREAMING_SNAKE_CASE      		=    [1, 5_12, 7, 7]

# Image classification docstring
__SCREAMING_SNAKE_CASE      		=    """sail/poolformer_s12"""
__SCREAMING_SNAKE_CASE      		=    """tabby, tabby cat"""

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



def     __a   (	a,		a = 0.0,		a = False       ):





       """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(__a,		dtype=input.dtype,		device=input.device       )
       random_tensor.floor_()  # binarize
       _a							 =					input.div(__a       ) * random_tensor
       return output





class __snake_case		(       nn.Module   ):







   """simple docstring"""




   def __init__(		self				:Any					,	UpperCamelCase__				:Tuple = None							):
          super().__init__()
          _a							 =					drop_prob




   def SCREAMING_SNAKE_CASE_   (		self				:List[Any]					,	UpperCamelCase__				:Dict							):
          return drop_path(SCREAMING_SNAKE_CASE_					,	self.drop_prob					,	self.training							)




   def SCREAMING_SNAKE_CASE_   (		self				:int							):
          return "p={}".format(self.drop_prob							)






class __snake_case		(       nn.Module   ):







   """simple docstring"""




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

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




   def SCREAMING_SNAKE_CASE_   (		self				:List[Any]					,	UpperCamelCase__				:List[Any]							):
          _a							 =					self.projection(SCREAMING_SNAKE_CASE_							)
          _a							 =					self.norm(SCREAMING_SNAKE_CASE_							)
          return embeddings






class __snake_case		(       nn.GroupNorm   ):







   """simple docstring"""




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






class __snake_case		(       nn.Module   ):







   """simple docstring"""




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




   def SCREAMING_SNAKE_CASE_   (		self				:List[Any]					,	UpperCamelCase__				:Optional[int]							):
          return self.pool(SCREAMING_SNAKE_CASE_							) - hidden_states






class __snake_case		(       nn.Module   ):







   """simple docstring"""




   def __init__(		self				:Union[str, Any]					,	UpperCamelCase__				:Tuple					,	UpperCamelCase__				:Optional[Any]					,	UpperCamelCase__				:str					,	UpperCamelCase__				:Dict							):
          super().__init__()
          _a							 =					nn.Convad(SCREAMING_SNAKE_CASE_					,	SCREAMING_SNAKE_CASE_					,	1							)
          _a							 =					nn.Convad(SCREAMING_SNAKE_CASE_					,	SCREAMING_SNAKE_CASE_					,	1							)
          _a							 =					PoolFormerDropPath(SCREAMING_SNAKE_CASE_							)
          if isinstance(config.hidden_act					,	SCREAMING_SNAKE_CASE_							):
                 _a							 =					ACTaFN[config.hidden_act]
          else:
                 _a							 =					config.hidden_act




   def SCREAMING_SNAKE_CASE_   (		self				:List[Any]					,	UpperCamelCase__				:Any							):
          _a							 =					self.conva(SCREAMING_SNAKE_CASE_							)
          _a							 =					self.act_fn(SCREAMING_SNAKE_CASE_							)
          _a							 =					self.drop(SCREAMING_SNAKE_CASE_							)
          _a							 =					self.conva(SCREAMING_SNAKE_CASE_							)
          _a							 =					self.drop(SCREAMING_SNAKE_CASE_							)

          return hidden_states






class __snake_case		(       nn.Module   ):







   """simple docstring"""




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

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




   def SCREAMING_SNAKE_CASE_   (		self				:List[Any]					,	UpperCamelCase__				:Dict							):







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

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

                 _a							 =					(output,) + outputs
                 return outputs

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

                 _a							 =					(output,) + outputs
                 return outputs






class __snake_case		(       nn.Module   ):







   """simple docstring"""




   def __init__(		self				:Optional[int]					,	UpperCamelCase__				:Union[str, Any]							):
          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(SCREAMING_SNAKE_CASE_							)

          # 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(
                                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_							)							)

          _a							 =					nn.ModuleList(SCREAMING_SNAKE_CASE_							)




   def SCREAMING_SNAKE_CASE_   (		self				:Any					,	UpperCamelCase__				:Any					,	UpperCamelCase__				:List[str]=False					,	UpperCamelCase__				: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(SCREAMING_SNAKE_CASE_							)
                 # Send the embeddings through the blocks
                 for _, blk in enumerate(SCREAMING_SNAKE_CASE_							):
                        _a							 =					blk(SCREAMING_SNAKE_CASE_							)
                        _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=SCREAMING_SNAKE_CASE_					,	hidden_states=SCREAMING_SNAKE_CASE_							)






class __snake_case		(       __lowercase   ):







   """simple docstring"""




   lowerCAmelCase_					: Optional[Any]  						=      PoolFormerConfig
   lowerCAmelCase_					: int  						=      'poolformer'
   lowerCAmelCase_					: Optional[Any]  						=      'pixel_values'
   lowerCAmelCase_					: Optional[Any]  						=      True




   def SCREAMING_SNAKE_CASE_   (		self				:Optional[Any]					,	UpperCamelCase__				: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 SCREAMING_SNAKE_CASE_   (		self				:Union[str, Any]					,	UpperCamelCase__				:List[Any]					,	UpperCamelCase__				:Optional[Any]=False							):







          if isinstance(SCREAMING_SNAKE_CASE_					,	SCREAMING_SNAKE_CASE_							):
                 _a							 =					value






__SCREAMING_SNAKE_CASE      		=    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      		=    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.'					,							__lowercase					,							)
class __snake_case		(       __lowercase   ):







   """simple docstring"""




   def __init__(		self				:Any					,	UpperCamelCase__				:Optional[int]							):
          super().__init__(SCREAMING_SNAKE_CASE_							)
          _a							 =					config

          _a							 =					PoolFormerEncoder(SCREAMING_SNAKE_CASE_							)

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




   def SCREAMING_SNAKE_CASE_   (		self				:Dict							):
          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 SCREAMING_SNAKE_CASE_   (		self				:List[Any]					,	UpperCamelCase__				:Optional[int] = None					,	UpperCamelCase__				:Any = None					,	UpperCamelCase__				:List[Any] = 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(
              SCREAMING_SNAKE_CASE_					,	output_hidden_states=SCREAMING_SNAKE_CASE_					,	return_dict=SCREAMING_SNAKE_CASE_					,	)
          _a							 =					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 __snake_case		(       nn.Module   ):







   """simple docstring"""




   def __init__(		self				:List[Any]					,	UpperCamelCase__				:Optional[Any]							):
          super().__init__()
          _a							 =					nn.Linear(config.hidden_size					,	config.hidden_size							)




   def SCREAMING_SNAKE_CASE_   (		self				:Any					,	UpperCamelCase__				:Dict							):
          _a							 =					self.dense(SCREAMING_SNAKE_CASE_							)
          return output






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







   """simple docstring"""




   def __init__(		self				:List[str]					,	UpperCamelCase__				:Union[str, Any]							):
          super().__init__(SCREAMING_SNAKE_CASE_							)
          _a							 =					config.num_labels
          _a							 =					PoolFormerModel(SCREAMING_SNAKE_CASE_							)

          # 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(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 SCREAMING_SNAKE_CASE_   (		self				:Optional[int]					,	UpperCamelCase__				:List[Any] = None					,	UpperCamelCase__				:int = None					,	UpperCamelCase__				:List[str] = None					,	UpperCamelCase__				:int = None					,	):
          _a							 =					return_dict if return_dict is not None else self.config.use_return_dict

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

          _a							 =					outputs[0]

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

          if not return_dict:
                 _a							 =					(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							)


"""simple docstring"""




import datasets

from .evaluate import evaluate


lowerCAmelCase__				=					'\\n@inproceedings{Rajpurkar2016SQuAD10,\n  title={SQuAD: 100, 000+ Questions for Machine Comprehension of Text},\n  author={Pranav Rajpurkar and Jian Zhang and Konstantin Lopyrev and Percy Liang},\n  booktitle={EMNLP},\n  year={2016}\n}\n'

lowerCAmelCase__				=					'\nThis metric wrap the official scoring script for version 1 of the Stanford Question Answering Dataset (SQuAD).\n\nStanford Question Answering Dataset (SQuAD) is a reading comprehension dataset, consisting of questions posed by\ncrowdworkers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span,\nfrom the corresponding reading passage, or the question might be unanswerable.\n'

lowerCAmelCase__				=					'\nComputes SQuAD scores (F1 and EM).\nArgs:\n    predictions: List of question-answers dictionaries with the following key-values:\n        - \'id\': id of the question-answer pair as given in the references (see below)\n        - \'prediction_text\': the text of the answer\n    references: List of question-answers dictionaries with the following key-values:\n        - \'id\': id of the question-answer pair (see above),\n        - \'answers\': a Dict in the SQuAD dataset format\n            {\n                \'text\': list of possible texts for the answer, as a list of strings\n                \'answer_start\': list of start positions for the answer, as a list of ints\n            }\n            Note that answer_start values are not taken into account to compute the metric.\nReturns:\n    \'exact_match\': Exact match (the normalized answer exactly match the gold answer)\n    \'f1\': The F-score of predicted tokens versus the gold answer\nExamples:\n\n    >>> predictions = [{\'prediction_text\': \'1976\', \'id\': \'56e10a3be3433e1400422b22\'}]\n    >>> references = [{\'answers\': {\'answer_start\': [97], \'text\': [\'1976\']}, \'id\': \'56e10a3be3433e1400422b22\'}]\n    >>> squad_metric = datasets.load_metric("squad")\n    >>> results = squad_metric.compute(predictions=predictions, references=references)\n    >>> print(results)\n    {\'exact_match\': 100.0, \'f1\': 100.0}\n'





@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''': {'''id''': datasets.Value('''string'''     ), '''prediction_text''': datasets.Value('''string'''     )},
									            '''references''': {
									                '''id''': datasets.Value('''string'''     ),
									                '''answers''': datasets.features.Sequence(
									                    {
									                        '''text''': datasets.Value('''string'''     ),
									                        '''answer_start''': datasets.Value('''int32'''     ),
									                    }     ),
									            },
									        }     )							,							codebase_urls=['''https://rajpurkar.github.io/SQuAD-explorer/''']							,							reference_urls=['''https://rajpurkar.github.io/SQuAD-explorer/''']							,							)







							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			{prediction['''id''']: prediction['''prediction_text'''] for prediction in predictions}
									SCREAMING_SNAKE_CASE_    			=			[
									    {
									        '''paragraphs''': [
									            {
									                '''qas''': [
									                    {
									                        '''answers''': [{'''text''': answer_text} for answer_text in ref['''answers''']['''text''']],
									                        '''id''': ref['''id'''],
									                    }
									                    for ref in references
									                ]
									            }
									        ]
									    }
									]
									SCREAMING_SNAKE_CASE_    			=			evaluate(dataset=SCREAMING_SNAKE_CASE_							,							predictions=SCREAMING_SNAKE_CASE_     )
									return score