




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


a       :						Any								=		logging.get_logger(__name__)

# General docstring
a       :						str								=		'''PoolFormerConfig'''

# Base docstring
a       :						str								=		'''sail/poolformer_s12'''
a       :						Tuple								=		[1, 512, 7, 7]

# Image classification docstring
a       :						List[str]								=		'''sail/poolformer_s12'''
a       :						Dict								=		'''tabby, tabby cat'''

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


def __UpperCAmelCase	(   _UpperCAmelCase		:   List[Any]		,				_UpperCAmelCase		:   float = 0.0		,				_UpperCAmelCase		:   bool = False						)      ->      Optional[int]:




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



class 		SCREAMING_SNAKE_CASE__  (			nn.Module	):





			def __init__(      self						:    Tuple      ,		a_						:    Optional[int] = None							):



						"""simple docstring"""




						super().__init__()
						__snake_case			=		drop_prob





			def 							A (      self						:    Optional[Any]      ,		a_						:    Dict							):



						"""simple docstring"""




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






			def 							A (      self						:    int							):



						"""simple docstring"""




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







class 		SCREAMING_SNAKE_CASE__  (			nn.Module	):





			def __init__(      self						:    str      ,		a_						:    str      ,		a_						:    Tuple      ,		a_						:    Optional[int]      ,		a_						:    Tuple      ,		a_						:    Any      ,		a_						:    Dict=None							):



						"""simple docstring"""




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

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






			def 							A (      self						:    List[str]      ,		a_						:    Optional[Any]							):



						"""simple docstring"""




						__snake_case			=		self.projection(UpperCamelCase__							)
						__snake_case			=		self.norm(UpperCamelCase__							)
						return embeddings







class 		SCREAMING_SNAKE_CASE__  (			nn.GroupNorm	):





			def __init__(      self						:    Optional[int]      ,		a_						:    List[Any]      ,		**a_						:    List[Any]							):



						"""simple docstring"""




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







class 		SCREAMING_SNAKE_CASE__  (			nn.Module	):





			def __init__(      self						:    Union[str, Any]      ,		a_						:    str							):



						"""simple docstring"""




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






			def 							A (      self						:    List[str]      ,		a_						:    int							):



						"""simple docstring"""




						return self.pool(UpperCamelCase__							) - hidden_states







class 		SCREAMING_SNAKE_CASE__  (			nn.Module	):





			def __init__(      self						:    int      ,		a_						:    List[Any]      ,		a_						:    List[str]      ,		a_						:    List[str]      ,		a_						:    Tuple							):



						"""simple docstring"""




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






			def 							A (      self						:    Optional[Any]      ,		a_						:    int							):



						"""simple docstring"""




						__snake_case			=		self.conva(UpperCamelCase__							)
						__snake_case			=		self.act_fn(UpperCamelCase__							)
						__snake_case			=		self.drop(UpperCamelCase__							)
						__snake_case			=		self.conva(UpperCamelCase__							)
						__snake_case			=		self.drop(UpperCamelCase__							)

						return hidden_states







class 		SCREAMING_SNAKE_CASE__  (			nn.Module	):





			def __init__(      self						:    List[str]      ,		a_						:    Tuple      ,		a_						:    Any      ,		a_						:    Optional[int]      ,		a_						:    int      ,		a_						:    Optional[Any]      ,		a_						:    List[Any]							):



						"""simple docstring"""




						super().__init__()
						__snake_case			=		PoolFormerPooling(UpperCamelCase__							)
						__snake_case			=		PoolFormerOutput(UpperCamelCase__      ,		UpperCamelCase__      ,		UpperCamelCase__      ,		UpperCamelCase__							)
						__snake_case			=		PoolFormerGroupNorm(UpperCamelCase__							)
						__snake_case			=		PoolFormerGroupNorm(UpperCamelCase__							)

						# Useful for training neural nets
						__snake_case			=		PoolFormerDropPath(UpperCamelCase__							) if drop_path > 0.0 else nn.Identity()
						__snake_case			=		config.use_layer_scale
						if config.use_layer_scale:
									__snake_case			=		nn.Parameter(
									    config.layer_scale_init_value * torch.ones((UpperCamelCase__)							)      ,		requires_grad=UpperCamelCase__							)
									__snake_case			=		nn.Parameter(
									    config.layer_scale_init_value * torch.ones((UpperCamelCase__)							)      ,		requires_grad=UpperCamelCase__							)






			def 							A (      self						:    int      ,		a_						:    Any							):



						"""simple docstring"""




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

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

									__snake_case			=		(output,) + outputs
									return outputs

						else:
									__snake_case			=		self.drop_path(self.pooling(self.before_norm(UpperCamelCase__							)							)							)
									# First residual connection
									__snake_case			=		pooling_output + hidden_states
									__snake_case			=		()

									# Second residual connection inside the PoolFormerOutput block
									__snake_case			=		self.drop_path(self.output(self.after_norm(UpperCamelCase__							)							)							)
									__snake_case			=		hidden_states + layer_output

									__snake_case			=		(output,) + outputs
									return outputs







class 		SCREAMING_SNAKE_CASE__  (			nn.Module	):





			def __init__(      self						:    Union[str, Any]      ,		a_						:    Any							):



						"""simple docstring"""




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

						# patch embeddings
						__snake_case			=		[]
						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]      ,		)							)
						__snake_case			=		nn.ModuleList(UpperCamelCase__							)

						# Transformer blocks
						__snake_case			=		[]
						__snake_case			=		0
						for i in range(config.num_encoder_blocks							):
									# each block consists of layers
									__snake_case			=		[]
									if i != 0:
												cur += config.depths[i - 1]
									for j in range(config.depths[i]							):
												layers.append(
												    PoolFormerLayer(
												        UpperCamelCase__      ,		num_channels=config.hidden_sizes[i]      ,		pool_size=config.pool_size      ,		hidden_size=config.hidden_sizes[i]      ,		intermediate_size=int(config.hidden_sizes[i] * config.mlp_ratio							)      ,		drop_path=dpr[cur + j]      ,		)							)
									blocks.append(nn.ModuleList(UpperCamelCase__							)							)

						__snake_case			=		nn.ModuleList(UpperCamelCase__							)






			def 							A (      self						:    int      ,		a_						:    List[str]      ,		a_						:    Union[str, Any]=False      ,		a_						:    List[Any]=True							):



						"""simple docstring"""




						__snake_case			=		() if output_hidden_states else None

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

									if output_hidden_states:
												__snake_case			=		all_hidden_states + (hidden_states,)

						if not return_dict:
									return tuple(v for v in [hidden_states, all_hidden_states] if v is not None							)

						return BaseModelOutputWithNoAttention(last_hidden_state=UpperCamelCase__      ,		hidden_states=UpperCamelCase__							)







class 		SCREAMING_SNAKE_CASE__  (			__A	):
			__SCREAMING_SNAKE_CASE    =			PoolFormerConfig
			__SCREAMING_SNAKE_CASE    =			'''poolformer'''
			__SCREAMING_SNAKE_CASE    =			'''pixel_values'''
			__SCREAMING_SNAKE_CASE    =			True





			def 							A (      self						:    List[str]      ,		a_						:    int							):



						"""simple docstring"""




						if isinstance(UpperCamelCase__      ,		(nn.Linear, nn.Convad)							):
									module.weight.data.normal_(mean=0.0      ,		std=self.config.initializer_range							)
									if module.bias is not None:
												module.bias.data.zero_()
						elif isinstance(UpperCamelCase__      ,		nn.LayerNorm							):
									module.bias.data.zero_()
									module.weight.data.fill_(1.0							)






			def 							A (      self						:    Union[str, Any]      ,		a_						:    Union[str, Any]      ,		a_						:    Any=False							):



						"""simple docstring"""




						if isinstance(UpperCamelCase__      ,		UpperCamelCase__							):
									__snake_case			=		value







a       :						Optional[int]								=		r'''\n    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use\n    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and\n    behavior.\n\n    Parameters:\n        config ([`PoolFormerConfig`]): Model configuration class with all the parameters of the model.\n            Initializing with a config file does not load the weights associated with the model, only the\n            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.\n'''

a       :						Optional[int]								=		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						:    int      ,		a_						:    Dict							):



						"""simple docstring"""




						super().__init__(UpperCamelCase__							)
						__snake_case			=		config

						__snake_case			=		PoolFormerEncoder(UpperCamelCase__							)

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





			def 							A (      self						:    str							):



						"""simple docstring"""




						return self.embeddings.patch_embeddings






			@add_start_docstrings_to_model_forward(UpperCamelCase__							)
			@add_code_sample_docstrings(
			    checkpoint=_CHECKPOINT_FOR_DOC      ,		output_type=UpperCamelCase__      ,		config_class=_CONFIG_FOR_DOC      ,		modality="vision"      ,		expected_output=_EXPECTED_OUTPUT_SHAPE      ,		)
			def 							A (      self						:    int      ,		a_						:    List[str] = None      ,		a_						:    Dict = None      ,		a_						:    Tuple = None      ,		):



						"""simple docstring"""




						__snake_case			=		(
						    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
						)
						__snake_case			=		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"							)

						__snake_case			=		self.encoder(
						    UpperCamelCase__      ,		output_hidden_states=UpperCamelCase__      ,		return_dict=UpperCamelCase__      ,		)
						__snake_case			=		encoder_outputs[0]

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

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







class 		SCREAMING_SNAKE_CASE__  (			nn.Module	):





			def __init__(      self						:    Any      ,		a_						:    Optional[int]							):



						"""simple docstring"""




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






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



						"""simple docstring"""




						__snake_case			=		self.dense(UpperCamelCase__							)
						return output







@add_start_docstrings(
    """
    PoolFormer Model transformer with an image classification head on top
    """				,     __A				,     )
class 		SCREAMING_SNAKE_CASE__  (			__A	):





			def __init__(      self						:    Tuple      ,		a_						:    str							):



						"""simple docstring"""




						super().__init__(UpperCamelCase__							)
						__snake_case			=		config.num_labels
						__snake_case			=		PoolFormerModel(UpperCamelCase__							)

						# Final norm
						__snake_case			=		PoolFormerGroupNorm(config.hidden_sizes[-1]							)
						# Classifier head
						__snake_case			=		(
						    nn.Linear(config.hidden_sizes[-1]      ,		config.num_labels							) if config.num_labels > 0 else nn.Identity()
						)

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






			@add_start_docstrings_to_model_forward(UpperCamelCase__							)
			@add_code_sample_docstrings(
			    checkpoint=_IMAGE_CLASS_CHECKPOINT      ,		output_type=UpperCamelCase__      ,		config_class=_CONFIG_FOR_DOC      ,		expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT      ,		)
			def 							A (      self						:    List[Any]      ,		a_						:    str = None      ,		a_						:    Any = None      ,		a_						:    List[str] = None      ,		a_						:    List[Any] = None      ,		):



						"""simple docstring"""




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

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

						__snake_case			=		outputs[0]

						__snake_case			=		self.classifier(self.norm(UpperCamelCase__							).mean([-2, -1]							)							)

						__snake_case			=		None
						if labels is not None:
									if self.config.problem_type is None:
												if self.num_labels == 1:
															__snake_case			=		"regression"
												elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
															__snake_case			=		"single_label_classification"
												else:
															__snake_case			=		"multi_label_classification"

									if self.config.problem_type == "regression":
												__snake_case			=		MSELoss()
												if self.num_labels == 1:
															__snake_case			=		loss_fct(logits.squeeze()      ,		labels.squeeze()							)
												else:
															__snake_case			=		loss_fct(UpperCamelCase__      ,		UpperCamelCase__							)
									elif self.config.problem_type == "single_label_classification":
												__snake_case			=		CrossEntropyLoss()
												__snake_case			=		loss_fct(logits.view(-1      ,		self.num_labels							)      ,		labels.view(-1							)							)
									elif self.config.problem_type == "multi_label_classification":
												__snake_case			=		BCEWithLogitsLoss()
												__snake_case			=		loss_fct(UpperCamelCase__      ,		UpperCamelCase__							)

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

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













'''simple docstring'''







import argparse
from collections import OrderedDict
from pathlib import Path

import torch

from transformers import (
    VisualBertConfig,
    VisualBertForMultipleChoice,
    VisualBertForPreTraining,
    VisualBertForQuestionAnswering,
    VisualBertForVisualReasoning,
)
from transformers.utils import logging


logging.set_verbosity_info()
__lowerCAmelCase        =							logging.get_logger(__name__)

__lowerCAmelCase        =							[
    ("bert.bert", "visual_bert"),
    ("bert.cls", "cls"),
    ("bert.classifier", "cls"),
    ("token_type_embeddings_visual", "visual_token_type_embeddings"),
    ("position_embeddings_visual", "visual_position_embeddings"),
    ("projection", "visual_projection"),
]

__lowerCAmelCase        =							[
    "nlvr2_coco_pre_trained.th",
    "nlvr2_fine_tuned.th",
    "nlvr2_pre_trained.th",
    "vcr_coco_pre_train.th",
    "vcr_fine_tune.th",
    "vcr_pre_train.th",
    "vqa_coco_pre_trained.th",
    "vqa_fine_tuned.th",
    "vqa_pre_trained.th",
]


def       __UpperCamelCase					(      lowercase_       :  int   ):
       """simple docstring"""




       a_						     =			torch.load(lowercase_   ,       map_location='cpu'   )
       return sd


def       __UpperCamelCase					(      lowercase_       :  List[str]   ,       lowercase_       :  Tuple   ,       lowercase_       :  Union[str, Any]=rename_keys_prefix   ):
       """simple docstring"""




       a_						     =			OrderedDict()
       a_						     =			torch.arange(config.max_position_embeddings   ).expand((1, -1)   )
       # detector_d = OrderedDict()
       for key in d:
              if "detector" in key:
                     # detector_d[key.replace('detector.','')] = d[key]
                     continue
              a_						     =			key
              for name_pair in rename_keys_prefix:
                     a_						     =			new_key.replace(name_pair[0]   ,       name_pair[1]   )
              a_						     =			d[key]
              if key == "bert.cls.predictions.decoder.weight":
                     # Old bert code didn't have `decoder.bias`, but was added separately
                     a_						     =			new_d['cls.predictions.bias']
       return new_d


@torch.no_grad()
def       __UpperCamelCase					(      lowercase_       :  Union[str, Any]   ,       lowercase_       :  Union[str, Any]   ):
       """simple docstring"""




       assert (
           checkpoint_path.split('/'   )[-1] in ACCEPTABLE_CHECKPOINTS
       ), F'The checkpoint provided must be in {ACCEPTABLE_CHECKPOINTS}.'

       # Get Config
       if "pre" in checkpoint_path:
              a_						     =			'pretraining'
              if "vcr" in checkpoint_path:
                     a_						     =			{'visual_embedding_dim': 512}
              elif "vqa_advanced" in checkpoint_path:
                     a_						     =			{'visual_embedding_dim': 2_048}
              elif "vqa" in checkpoint_path:
                     a_						     =			{'visual_embedding_dim': 2_048}
              elif "nlvr" in checkpoint_path:
                     a_						     =			{'visual_embedding_dim': 1_024}
              else:
                     raise NotImplementedError(F'No implementation found for `{checkpoint_path}`.'   )
       else:
              if "vcr" in checkpoint_path:
                     a_						     =			{'visual_embedding_dim': 512}
                     a_						     =			'multichoice'
              elif "vqa_advanced" in checkpoint_path:
                     a_						     =			{'visual_embedding_dim': 2_048}
                     a_						     =			'vqa_advanced'
              elif "vqa" in checkpoint_path:
                     a_						     =			{'visual_embedding_dim': 2_048, 'num_labels': 3_129}
                     a_						     =			'vqa'
              elif "nlvr" in checkpoint_path:
                     a_						     =			{
                         'visual_embedding_dim': 1_024,
                         'num_labels': 2,
                     }
                     a_						     =			'nlvr'

       a_						     =			VisualBertConfig(**lowercase_   )

       # Load State Dict
       a_						     =			load_state_dict(lowercase_   )

       a_						     =			get_new_dict(lowercase_   ,       lowercase_   )

       if model_type == "pretraining":
              a_						     =			VisualBertForPreTraining(lowercase_   )
       elif model_type == "vqa":
              a_						     =			VisualBertForQuestionAnswering(lowercase_   )
       elif model_type == "nlvr":
              a_						     =			VisualBertForVisualReasoning(lowercase_   )
       elif model_type == "multichoice":
              a_						     =			VisualBertForMultipleChoice(lowercase_   )

       model.load_state_dict(lowercase_   )
       # Save Checkpoints
       Path(lowercase_   ).mkdir(exist_ok=lowercase_   )
       model.save_pretrained(lowercase_   )


if __name__ == "__main__":
  __lowerCAmelCase        =							argparse.ArgumentParser()
  # Required parameters
  parser.add_argument("orig_checkpoint_path", type=str, help="A path to .th on local filesystem.")
  parser.add_argument("pytorch_dump_folder_path", type=str, help="Path to the output PyTorch model.")
  __lowerCAmelCase        =							parser.parse_args()
  convert_visual_bert_checkpoint(args.orig_checkpoint_path, args.pytorch_dump_folder_path)




