





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


lowerCamelCase_								=		logging.get_logger(__name__)

# General docstring
lowerCamelCase_								=		"PoolFormerConfig"

# Base docstring
lowerCamelCase_								=		"sail/poolformer_s12"
lowerCamelCase_								=		[1, 5_1_2, 7, 7]

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

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



def      __lowerCamelCase   (							a_     :		List[Any] ,   a_     :		float = 0.0 ,   a_     :		bool = False							)			->      Optional[int]:
      if drop_prob == 0.0 or not training:
            return input
      __SCREAMING_SNAKE_CASE   :Dict						=  1 - drop_prob
      __SCREAMING_SNAKE_CASE   :Tuple						=  (input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
      __SCREAMING_SNAKE_CASE   :List[str]						=  keep_prob + torch.rand(a_ ,   dtype=input.dtype ,   device=input.device							)
      random_tensor.floor_()  # binarize
      __SCREAMING_SNAKE_CASE   :Tuple						=  input.div(a_							) * random_tensor
      return output






class 					_SCREAMING_SNAKE_CASE(   nn.Module  ):




       def __init__(						self		,SCREAMING_SNAKE_CASE__ = None							)			->      None:





             """simple docstring"""






             super().__init__()
             __SCREAMING_SNAKE_CASE   :Any						=  drop_prob




       def     _UpperCamelCase     (						self		,SCREAMING_SNAKE_CASE__							)			->      torch.Tensor:





             """simple docstring"""






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




       def     _UpperCamelCase     (						self							)			->      str:





             """simple docstring"""






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






class 					_SCREAMING_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							)			->      List[str]:





             """simple docstring"""






             super().__init__()
             __SCREAMING_SNAKE_CASE   :Optional[int]						=  patch_size if isinstance(SCREAMING_SNAKE_CASE__		,collections.abc.Iterable							) else (patch_size, patch_size)
             __SCREAMING_SNAKE_CASE   :str						=  stride if isinstance(SCREAMING_SNAKE_CASE__		,collections.abc.Iterable							) else (stride, stride)
             __SCREAMING_SNAKE_CASE   :List[str]						=  padding if isinstance(SCREAMING_SNAKE_CASE__		,collections.abc.Iterable							) else (padding, padding)

             __SCREAMING_SNAKE_CASE   :Tuple						=  nn.Convad(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,kernel_size=SCREAMING_SNAKE_CASE__		,stride=SCREAMING_SNAKE_CASE__		,padding=SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :int						=  norm_layer(SCREAMING_SNAKE_CASE__							) if norm_layer else nn.Identity()




       def     _UpperCamelCase     (						self		,SCREAMING_SNAKE_CASE__							)			->      str:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Dict						=  self.projection(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :Optional[int]						=  self.norm(SCREAMING_SNAKE_CASE__							)
             return embeddings






class 					_SCREAMING_SNAKE_CASE(   nn.GroupNorm  ):




       def __init__(						self		,SCREAMING_SNAKE_CASE__		,**SCREAMING_SNAKE_CASE__							)			->      str:





             """simple docstring"""






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






class 					_SCREAMING_SNAKE_CASE(   nn.Module  ):




       def __init__(						self		,SCREAMING_SNAKE_CASE__							)			->      int:





             """simple docstring"""






             super().__init__()
             __SCREAMING_SNAKE_CASE   :int						=  nn.AvgPoolad(SCREAMING_SNAKE_CASE__		,stride=1		,padding=pool_size // 2		,count_include_pad=SCREAMING_SNAKE_CASE__							)




       def     _UpperCamelCase     (						self		,SCREAMING_SNAKE_CASE__							)			->      int:





             """simple docstring"""






             return self.pool(SCREAMING_SNAKE_CASE__							) - hidden_states






class 					_SCREAMING_SNAKE_CASE(   nn.Module  ):




       def __init__(						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							)			->      List[str]:





             """simple docstring"""






             super().__init__()
             __SCREAMING_SNAKE_CASE   :Any						=  nn.Convad(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,1							)
             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  nn.Convad(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,1							)
             __SCREAMING_SNAKE_CASE   :int						=  PoolFormerDropPath(SCREAMING_SNAKE_CASE__							)
             if isinstance(config.hidden_act		,SCREAMING_SNAKE_CASE__							):
                   __SCREAMING_SNAKE_CASE   :Dict						=  ACTaFN[config.hidden_act]
             else:
                   __SCREAMING_SNAKE_CASE   :Tuple						=  config.hidden_act




       def     _UpperCamelCase     (						self		,SCREAMING_SNAKE_CASE__							)			->      int:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :int						=  self.conva(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :Dict						=  self.act_fn(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :Any						=  self.drop(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :Any						=  self.conva(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :str						=  self.drop(SCREAMING_SNAKE_CASE__							)

             return hidden_states






class 					_SCREAMING_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__							)			->      str:





             """simple docstring"""






             super().__init__()
             __SCREAMING_SNAKE_CASE   :Optional[int]						=  PoolFormerPooling(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :Tuple						=  PoolFormerOutput(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  PoolFormerGroupNorm(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  PoolFormerGroupNorm(SCREAMING_SNAKE_CASE__							)

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




       def     _UpperCamelCase     (						self		,SCREAMING_SNAKE_CASE__							)			->      List[str]:





             """simple docstring"""






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

                   __SCREAMING_SNAKE_CASE   :int						=  self.output(self.after_norm(SCREAMING_SNAKE_CASE__							)							)
                   __SCREAMING_SNAKE_CASE   :Tuple						=  self.layer_scale_a.unsqueeze(-1							).unsqueeze(-1							) * layer_output
                   # Second residual connection
                   __SCREAMING_SNAKE_CASE   :List[Any]						=  hidden_states + self.drop_path(SCREAMING_SNAKE_CASE__							)

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

             else:
                   __SCREAMING_SNAKE_CASE   :Optional[int]						=  self.drop_path(self.pooling(self.before_norm(SCREAMING_SNAKE_CASE__							)							)							)
                   # First residual connection
                   __SCREAMING_SNAKE_CASE   :Optional[Any]						=  pooling_output + hidden_states
                   __SCREAMING_SNAKE_CASE   :Dict						=  ()

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

                   __SCREAMING_SNAKE_CASE   :List[str]						=  (output,) + outputs
                   return outputs






class 					_SCREAMING_SNAKE_CASE(   nn.Module  ):




       def __init__(						self		,SCREAMING_SNAKE_CASE__							)			->      Any:





             """simple docstring"""






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

             # patch embeddings
             __SCREAMING_SNAKE_CASE   :List[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]		,)							)
             __SCREAMING_SNAKE_CASE   :Tuple						=  nn.ModuleList(SCREAMING_SNAKE_CASE__							)

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

             __SCREAMING_SNAKE_CASE   :str						=  nn.ModuleList(SCREAMING_SNAKE_CASE__							)




       def     _UpperCamelCase     (						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__=False		,SCREAMING_SNAKE_CASE__=True							)			->      List[Any]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :int						=  () if output_hidden_states else None

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

                   if output_hidden_states:
                         __SCREAMING_SNAKE_CASE   :Optional[int]						=  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 					_SCREAMING_SNAKE_CASE(   A  ):
       SCREAMING_SNAKE_CASE_   :      Dict							=  PoolFormerConfig
       SCREAMING_SNAKE_CASE_   :      Tuple							=  '''poolformer'''
       SCREAMING_SNAKE_CASE_   :      str							=  '''pixel_values'''
       SCREAMING_SNAKE_CASE_   :      Optional[int]							=  True




       def     _UpperCamelCase     (						self		,SCREAMING_SNAKE_CASE__							)			->      Union[str, Any]:





             """simple docstring"""






             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     _UpperCamelCase     (						self		,SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__=False							)			->      Optional[int]:





             """simple docstring"""






             if isinstance(SCREAMING_SNAKE_CASE__		,SCREAMING_SNAKE_CASE__							):
                   __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  value




lowerCamelCase_								=		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_								=		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 					_SCREAMING_SNAKE_CASE(   A  ):




       def __init__(						self		,SCREAMING_SNAKE_CASE__							)			->      Optional[int]:





             """simple docstring"""






             super().__init__(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :Optional[int]						=  config

             __SCREAMING_SNAKE_CASE   :List[str]						=  PoolFormerEncoder(SCREAMING_SNAKE_CASE__							)

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




       def     _UpperCamelCase     (						self							)			->      Optional[Any]:





             """simple docstring"""






             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     _UpperCamelCase     (						self		,SCREAMING_SNAKE_CASE__ = None		,SCREAMING_SNAKE_CASE__ = None		,SCREAMING_SNAKE_CASE__ = None		,)			->      Union[Tuple, BaseModelOutputWithNoAttention]:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :Any						=  (
                 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
             )
             __SCREAMING_SNAKE_CASE   :int						=  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'''							)

             __SCREAMING_SNAKE_CASE   :Any						=  self.encoder(
                 SCREAMING_SNAKE_CASE__		,output_hidden_states=SCREAMING_SNAKE_CASE__		,return_dict=SCREAMING_SNAKE_CASE__		,)
             __SCREAMING_SNAKE_CASE   :Optional[int]						=  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 					_SCREAMING_SNAKE_CASE(   nn.Module  ):




       def __init__(						self		,SCREAMING_SNAKE_CASE__							)			->      Tuple:





             """simple docstring"""






             super().__init__()
             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  nn.Linear(config.hidden_size		,config.hidden_size							)




       def     _UpperCamelCase     (						self		,SCREAMING_SNAKE_CASE__							)			->      Any:





             """simple docstring"""






             __SCREAMING_SNAKE_CASE   :List[Any]						=  self.dense(SCREAMING_SNAKE_CASE__							)
             return output






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




       def __init__(						self		,SCREAMING_SNAKE_CASE__							)			->      int:





             """simple docstring"""






             super().__init__(SCREAMING_SNAKE_CASE__							)
             __SCREAMING_SNAKE_CASE   :Any						=  config.num_labels
             __SCREAMING_SNAKE_CASE   :int						=  PoolFormerModel(SCREAMING_SNAKE_CASE__							)

             # Final norm
             __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  PoolFormerGroupNorm(config.hidden_sizes[-1]							)
             # Classifier head
             __SCREAMING_SNAKE_CASE   :List[Any]						=  (
                 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     _UpperCamelCase     (						self		,SCREAMING_SNAKE_CASE__ = None		,SCREAMING_SNAKE_CASE__ = None		,SCREAMING_SNAKE_CASE__ = None		,SCREAMING_SNAKE_CASE__ = None		,)			->      Union[Tuple, ImageClassifierOutputWithNoAttention]:





             """simple docstring"""






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

             __SCREAMING_SNAKE_CASE   :List[str]						=  self.poolformer(
                 SCREAMING_SNAKE_CASE__		,output_hidden_states=SCREAMING_SNAKE_CASE__		,return_dict=SCREAMING_SNAKE_CASE__		,)

             __SCREAMING_SNAKE_CASE   :Optional[int]						=  outputs[0]

             __SCREAMING_SNAKE_CASE   :Optional[Any]						=  self.classifier(self.norm(SCREAMING_SNAKE_CASE__							).mean([-2, -1]							)							)

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

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

             if not return_dict:
                   __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  (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 argparse
import os

import torch

from transformers import FlavaImageCodebook, FlavaImageCodebookConfig



def      __lowerCamelCase   (							a_     :		str ,   a_     :		Dict ,   a_     :		Any ,   a_     :		str							)			->      str:
      __SCREAMING_SNAKE_CASE   :int						=  s.rsplit(a_ ,   a_							)
      return new.join(a_							)



def      __lowerCamelCase   (							a_     :		List[str]							)			->      Dict:
      # encoder.embeddings are double copied in original FLAVA
      return sum(param.float().sum() if '''encoder.embeddings''' not in key else 0 for key, param in state_dict.items()							)



def      __lowerCamelCase   (							a_     :		Optional[int]							)			->      Any:
      __SCREAMING_SNAKE_CASE   :Optional[int]						=  {}

      __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  ['''group_1''', '''group_2''', '''group_3''', '''group_4''']
      for key, value in state_dict.items():
            for group_key in group_keys:
                  if group_key in key:
                        __SCREAMING_SNAKE_CASE   :Optional[Any]						=  key.replace(f'''{group_key}.''' ,   f'''{group_key}.group.'''							)

            if "res_path" in key:
                  __SCREAMING_SNAKE_CASE   :str						=  key.replace('''res_path.''' ,   '''res_path.path.'''							)

            if key.endswith('''.w'''							):
                  __SCREAMING_SNAKE_CASE   :List[Any]						=  rreplace(a_ ,   '''.w''' ,   '''.weight''' ,   1							)
            if key.endswith('''.b'''							):
                  __SCREAMING_SNAKE_CASE   :List[Any]						=  rreplace(a_ ,   '''.b''' ,   '''.bias''' ,   1							)

            __SCREAMING_SNAKE_CASE   :Optional[Any]						=  value.float()

      return upgrade



@torch.no_grad()
def      __lowerCamelCase   (							a_     :		List[Any] ,   a_     :		Optional[int] ,   a_     :		Optional[int]=None ,   a_     :		Dict=True							)			->      Union[str, Any]:
      from dall_e import Encoder

      __SCREAMING_SNAKE_CASE   :int						=  Encoder()
      if os.path.exists(a_							):
            __SCREAMING_SNAKE_CASE   :Dict						=  torch.load(a_							)
      else:
            __SCREAMING_SNAKE_CASE   :List[str]						=  torch.hub.load_state_dict_from_url(a_							)

      if isinstance(a_ ,   a_							):
            __SCREAMING_SNAKE_CASE   :List[str]						=  ckpt.state_dict()
      encoder.load_state_dict(a_							)

      if config_path is not None:
            __SCREAMING_SNAKE_CASE   :Any						=  FlavaImageCodebookConfig.from_pretrained(a_							)
      else:
            __SCREAMING_SNAKE_CASE   :Optional[int]						=  FlavaImageCodebookConfig()

      __SCREAMING_SNAKE_CASE   :Tuple						=  FlavaImageCodebook(a_							).eval()
      __SCREAMING_SNAKE_CASE   :List[str]						=  encoder.state_dict()

      __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  upgrade_state_dict(a_							)
      hf_model.load_state_dict(a_							)
      __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  hf_model.state_dict()
      __SCREAMING_SNAKE_CASE   :Union[str, Any]						=  count_parameters(a_							)
      __SCREAMING_SNAKE_CASE   :Any						=  count_parameters(a_							)

      assert torch.allclose(a_ ,   a_ ,   atol=1e-3							)

      if save_checkpoint:
            hf_model.save_pretrained(a_							)
      else:
            return hf_state_dict


if __name__ == "__main__":
  lowerCamelCase_								=		argparse.ArgumentParser()
  parser.add_argument("--pytorch_dump_folder_path", default=None, type=str, help="Path to the output PyTorch model.")
  parser.add_argument("--checkpoint_path", default=None, type=str, help="Path to flava checkpoint")
  parser.add_argument("--config_path", default=None, type=str, help="Path to hf config.json of model to convert")
  lowerCamelCase_								=		parser.parse_args()

  convert_dalle_checkpoint(args.checkpoint_path, args.pytorch_dump_folder_path, args.config_path)