



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


A_		:     Optional[Any]        =     logging.get_logger(__name__)

# General docstring
A_		:     List[str]        =     '''PoolFormerConfig'''

# Base docstring
A_		:     Tuple        =     '''sail/poolformer_s12'''
A_		:     Optional[int]        =     [1, 5_12, 7, 7]

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

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


def 	UpperCAmelCase__      (      UpperCAmelCase__ :Optional[Any] ,					UpperCAmelCase__ :float = 0.0 ,					UpperCAmelCase__ :bool = 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(UpperCAmelCase__ ,					dtype=input.dtype ,					device=input.device					)
 random_tensor.floor_()  # binarize
 a						=		input.div(UpperCAmelCase__					) * random_tensor
 return output



class 		_lowercase		(       nn.Module	):

   def __init__(      self			:       Tuple							,   __lowerCAmelCase			:       Optional[float] = None							)       ->				None:


    """simple docstring"""






    super().__init__()
    a						=		drop_prob

   def 		A							(      self			:       Union[str, Any]							,   __lowerCAmelCase			:       torch.Tensor							)       ->				torch.Tensor:


    """simple docstring"""






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




   def 		A							(      self			:       List[str]							)       ->				str:


    """simple docstring"""






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



class 		_lowercase		(       nn.Module	):

   def __init__(      self			:       Union[str, Any]							,   __lowerCAmelCase			:       Tuple							,   __lowerCAmelCase			:       Union[str, Any]							,   __lowerCAmelCase			:       Any							,   __lowerCAmelCase			:       Union[str, Any]							,   __lowerCAmelCase			:       Dict							,   __lowerCAmelCase			:       Dict=None							)       ->				Optional[Any]:


    """simple docstring"""






    super().__init__()
    a						=		patch_size if isinstance(__lowerCAmelCase							,   collections.abc.Iterable							) else (patch_size, patch_size)
    a						=		stride if isinstance(__lowerCAmelCase							,   collections.abc.Iterable							) else (stride, stride)
    a						=		padding if isinstance(__lowerCAmelCase							,   collections.abc.Iterable							) else (padding, padding)

    a						=		nn.Convad(__lowerCAmelCase							,   __lowerCAmelCase							,   kernel_size=__lowerCAmelCase							,   stride=__lowerCAmelCase							,   padding=__lowerCAmelCase							)
    a						=		norm_layer(__lowerCAmelCase							) if norm_layer else nn.Identity()




   def 		A							(      self			:       List[str]							,   __lowerCAmelCase			:       Union[str, Any]							)       ->				Tuple:


    """simple docstring"""






    a						=		self.projection(__lowerCAmelCase							)
    a						=		self.norm(__lowerCAmelCase							)
    return embeddings



class 		_lowercase		(       nn.GroupNorm	):

   def __init__(      self			:       Tuple							,   __lowerCAmelCase			:       int							,   **__lowerCAmelCase			:       Optional[Any]							)       ->				Optional[int]:


    """simple docstring"""






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



class 		_lowercase		(       nn.Module	):

   def __init__(      self			:       int							,   __lowerCAmelCase			:       Optional[Any]							)       ->				str:


    """simple docstring"""






    super().__init__()
    a						=		nn.AvgPoolad(__lowerCAmelCase							,   stride=1							,   padding=pool_size // 2							,   count_include_pad=__lowerCAmelCase							)




   def 		A							(      self			:       Optional[Any]							,   __lowerCAmelCase			:       List[Any]							)       ->				Any:


    """simple docstring"""






    return self.pool(__lowerCAmelCase							) - hidden_states



class 		_lowercase		(       nn.Module	):

   def __init__(      self			:       Tuple							,   __lowerCAmelCase			:       int							,   __lowerCAmelCase			:       Optional[int]							,   __lowerCAmelCase			:       Optional[int]							,   __lowerCAmelCase			:       Any							)       ->				List[Any]:


    """simple docstring"""






    super().__init__()
    a						=		nn.Convad(__lowerCAmelCase							,   __lowerCAmelCase							,   1							)
    a						=		nn.Convad(__lowerCAmelCase							,   __lowerCAmelCase							,   1							)
    a						=		PoolFormerDropPath(__lowerCAmelCase							)
    if isinstance(config.hidden_act							,   __lowerCAmelCase							):
     a						=		ACTaFN[config.hidden_act]
    else:
     a						=		config.hidden_act




   def 		A							(      self			:       Tuple							,   __lowerCAmelCase			:       List[str]							)       ->				Optional[Any]:


    """simple docstring"""






    a						=		self.conva(__lowerCAmelCase							)
    a						=		self.act_fn(__lowerCAmelCase							)
    a						=		self.drop(__lowerCAmelCase							)
    a						=		self.conva(__lowerCAmelCase							)
    a						=		self.drop(__lowerCAmelCase							)

    return hidden_states



class 		_lowercase		(       nn.Module	):

   def __init__(      self			:       Dict							,   __lowerCAmelCase			:       Optional[Any]							,   __lowerCAmelCase			:       int							,   __lowerCAmelCase			:       Dict							,   __lowerCAmelCase			:       Dict							,   __lowerCAmelCase			:       Tuple							,   __lowerCAmelCase			:       Union[str, Any]							)       ->				Optional[Any]:


    """simple docstring"""






    super().__init__()
    a						=		PoolFormerPooling(__lowerCAmelCase							)
    a						=		PoolFormerOutput(__lowerCAmelCase							,   __lowerCAmelCase							,   __lowerCAmelCase							,   __lowerCAmelCase							)
    a						=		PoolFormerGroupNorm(__lowerCAmelCase							)
    a						=		PoolFormerGroupNorm(__lowerCAmelCase							)

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




   def 		A							(      self			:       Dict							,   __lowerCAmelCase			:       Dict							)       ->				str:


    """simple docstring"""






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

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

     a						=		(output,) + outputs
     return outputs

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

     a						=		(output,) + outputs
     return outputs



class 		_lowercase		(       nn.Module	):

   def __init__(      self			:       Any							,   __lowerCAmelCase			:       Dict							)       ->				Union[str, Any]:


    """simple docstring"""






    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(__lowerCAmelCase							)

    # 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(
              __lowerCAmelCase							,   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(__lowerCAmelCase							)							)

    a						=		nn.ModuleList(__lowerCAmelCase							)




   def 		A							(      self			:       Any							,   __lowerCAmelCase			:       Dict							,   __lowerCAmelCase			:       Optional[Any]=False							,   __lowerCAmelCase			:       List[str]=True							)       ->				Tuple:


    """simple docstring"""






    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(__lowerCAmelCase							)
     # Send the embeddings through the blocks
     for _, blk in enumerate(__lowerCAmelCase							):
      a						=		blk(__lowerCAmelCase							)
      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=__lowerCAmelCase							,   hidden_states=__lowerCAmelCase							)



class 		_lowercase		(       UpperCAmelCase__	):
   _UpperCAmelCase     	=			PoolFormerConfig
   _UpperCAmelCase     	=			'''poolformer'''
   _UpperCAmelCase     	=			'''pixel_values'''
   _UpperCAmelCase     	=			True

   def 		A							(      self			:       List[str]							,   __lowerCAmelCase			:       List[str]							)       ->				List[Any]:


    """simple docstring"""






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




   def 		A							(      self			:       Optional[int]							,   __lowerCAmelCase			:       Optional[int]							,   __lowerCAmelCase			:       Optional[Any]=False							)       ->				List[Any]:


    """simple docstring"""






    if isinstance(__lowerCAmelCase							,   __lowerCAmelCase							):
     a						=		value

A_		:     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.
'''

A_		:     Optional[int]        =     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.''',   UpperCAmelCase__,   )
class 		_lowercase		(       UpperCAmelCase__	):

   def __init__(      self			:       Any							,   __lowerCAmelCase			:       Union[str, Any]							)       ->				List[Any]:


    """simple docstring"""






    super().__init__(__lowerCAmelCase							)
    a						=		config

    a						=		PoolFormerEncoder(__lowerCAmelCase							)

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

   def 		A							(      self			:       List[Any]							)       ->				Any:


    """simple docstring"""






    return self.embeddings.patch_embeddings




   @add_start_docstrings_to_model_forward(__lowerCAmelCase							)
   @add_code_sample_docstrings(
       checkpoint=_CHECKPOINT_FOR_DOC							,   output_type=__lowerCAmelCase							,   config_class=_CONFIG_FOR_DOC							,   modality="vision"							,   expected_output=_EXPECTED_OUTPUT_SHAPE							,   )
   def 		A							(      self			:       int							,   __lowerCAmelCase			:       Optional[torch.FloatTensor] = None							,   __lowerCAmelCase			:       Optional[bool] = None							,   __lowerCAmelCase			:       Optional[bool] = None							,   )       ->				Union[Tuple, BaseModelOutputWithNoAttention]:


    """simple docstring"""






    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(
        __lowerCAmelCase							,   output_hidden_states=__lowerCAmelCase							,   return_dict=__lowerCAmelCase							,   )
    a						=		encoder_outputs[0]

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

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



class 		_lowercase		(       nn.Module	):

   def __init__(      self			:       Union[str, Any]							,   __lowerCAmelCase			:       Any							)       ->				str:


    """simple docstring"""






    super().__init__()
    a						=		nn.Linear(config.hidden_size							,   config.hidden_size							)




   def 		A							(      self			:       int							,   __lowerCAmelCase			:       List[Any]							)       ->				Any:


    """simple docstring"""






    a						=		self.dense(__lowerCAmelCase							)
    return output



@add_start_docstrings(
    '''
    PoolFormer Model transformer with an image classification head on top
    ''',   UpperCAmelCase__,   )
class 		_lowercase		(       UpperCAmelCase__	):

   def __init__(      self			:       Dict							,   __lowerCAmelCase			:       Optional[int]							)       ->				Any:


    """simple docstring"""






    super().__init__(__lowerCAmelCase							)
    a						=		config.num_labels
    a						=		PoolFormerModel(__lowerCAmelCase							)

    # 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(__lowerCAmelCase							)
   @add_code_sample_docstrings(
       checkpoint=_IMAGE_CLASS_CHECKPOINT							,   output_type=__lowerCAmelCase							,   config_class=_CONFIG_FOR_DOC							,   expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT							,   )
   def 		A							(      self			:       Union[str, Any]							,   __lowerCAmelCase			:       Optional[torch.FloatTensor] = None							,   __lowerCAmelCase			:       Optional[torch.LongTensor] = None							,   __lowerCAmelCase			:       Optional[bool] = None							,   __lowerCAmelCase			:       Optional[bool] = None							,   )       ->				Union[Tuple, ImageClassifierOutputWithNoAttention]:


    """simple docstring"""






    a						=		return_dict if return_dict is not None else self.config.use_return_dict

    a						=		self.poolformer(
        __lowerCAmelCase							,   output_hidden_states=__lowerCAmelCase							,   return_dict=__lowerCAmelCase							,   )

    a						=		outputs[0]

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

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

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






from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


A_		:     Optional[int]        =     {
    '''configuration_instructblip''': [
        '''INSTRUCTBLIP_PRETRAINED_CONFIG_ARCHIVE_MAP''',
        '''InstructBlipConfig''',
        '''InstructBlipQFormerConfig''',
        '''InstructBlipVisionConfig''',
    ],
    '''processing_instructblip''': ['''InstructBlipProcessor'''],
}

try:
     if not is_torch_available():
          raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
     pass
else:
     A_		:     List[Any]        =     [
         '''INSTRUCTBLIP_PRETRAINED_MODEL_ARCHIVE_LIST''',
         '''InstructBlipQFormerModel''',
         '''InstructBlipPreTrainedModel''',
         '''InstructBlipForConditionalGeneration''',
         '''InstructBlipVisionModel''',
     ]

if TYPE_CHECKING:
     from .configuration_instructblip import (
         INSTRUCTBLIP_PRETRAINED_CONFIG_ARCHIVE_MAP,
         InstructBlipConfig,
         InstructBlipQFormerConfig,
         InstructBlipVisionConfig,
     )
     from .processing_instructblip import InstructBlipProcessor

     try:
          if not is_torch_available():
               raise OptionalDependencyNotAvailable()
     except OptionalDependencyNotAvailable:
          pass
     else:
          from .modeling_instructblip import (
              INSTRUCTBLIP_PRETRAINED_MODEL_ARCHIVE_LIST,
              InstructBlipForConditionalGeneration,
              InstructBlipPreTrainedModel,
              InstructBlipQFormerModel,
              InstructBlipVisionModel,
          )

else:
     import sys

     A_		:     Tuple        =     _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)

