





import argparse

import gdown
import numpy as np
import torch
from huggingface_hub import hf_hub_download

from transformers import (
    CLIPTokenizer,
    CLIPTokenizerFast,
    VideoMAEImageProcessor,
    XCLIPConfig,
    XCLIPModel,
    XCLIPProcessor,
    XCLIPTextConfig,
    XCLIPVisionConfig,
)

def     snake_case							(UpperCamelCase			:     Optional[int]			,    UpperCamelCase			:     Union[str, Any]							):




	'''simple docstring'''







	lowerCamelCase__										=			XCLIPTextConfig()

	# derive patch size from model name
	lowerCamelCase__										=			model_name.find("""patch"""							)
	lowerCamelCase__										=			int(model_name[start_idx + len("""patch"""							) : start_idx + len("""patch"""							) + 2]							)
	lowerCamelCase__										=			XCLIPVisionConfig(patch_size=lowercase_			,    num_frames=lowercase_							)

	if "large" in model_name:
		lowerCamelCase__										=			768
		lowerCamelCase__										=			3072
		lowerCamelCase__										=			12

		lowerCamelCase__										=			1024
		lowerCamelCase__										=			4096
		lowerCamelCase__										=			16
		lowerCamelCase__										=			24
		lowerCamelCase__										=			768
		lowerCamelCase__										=			3072

	if model_name == "xclip-large-patch14-16-frames":
		lowerCamelCase__										=			336

	lowerCamelCase__										=			XCLIPConfig.from_text_vision_configs(lowercase_			,    lowercase_							)

	if "large" in model_name:
		lowerCamelCase__										=			768

	return config

def     snake_case							(UpperCamelCase			:     int							):




	'''simple docstring'''







	if name == "token_embedding.weight":
		lowerCamelCase__										=			name.replace("""token_embedding.weight"""			,    """text_model.embeddings.token_embedding.weight"""							)
	if name == "positional_embedding":
		lowerCamelCase__										=			name.replace("""positional_embedding"""			,    """text_model.embeddings.position_embedding.weight"""							)
	if "ln_1" in name:
		lowerCamelCase__										=			name.replace("""ln_1"""			,    """layer_norm1"""							)
	if "ln_2" in name:
		lowerCamelCase__										=			name.replace("""ln_2"""			,    """layer_norm2"""							)
	if "c_fc" in name:
		lowerCamelCase__										=			name.replace("""c_fc"""			,    """fc1"""							)
	if "c_proj" in name:
		lowerCamelCase__										=			name.replace("""c_proj"""			,    """fc2"""							)
	if name.startswith("""transformer.resblocks"""							):
		lowerCamelCase__										=			name.replace("""transformer.resblocks"""			,    """text_model.encoder.layers"""							)
	if "attn.out_proj" in name and "message" not in name:
		lowerCamelCase__										=			name.replace("""attn.out_proj"""			,    """self_attn.out_proj"""							)
	if "ln_final" in name:
		lowerCamelCase__										=			name.replace("""ln_final"""			,    """text_model.final_layer_norm"""							)
	# visual encoder
	if name == "visual.class_embedding":
		lowerCamelCase__										=			name.replace("""visual.class_embedding"""			,    """vision_model.embeddings.class_embedding"""							)
	if name == "visual.positional_embedding":
		lowerCamelCase__										=			name.replace("""visual.positional_embedding"""			,    """vision_model.embeddings.position_embedding.weight"""							)
	if name.startswith("""visual.transformer.resblocks"""							):
		lowerCamelCase__										=			name.replace("""visual.transformer.resblocks"""			,    """vision_model.encoder.layers"""							)
	if "visual.conv1" in name:
		lowerCamelCase__										=			name.replace("""visual.conv1"""			,    """vision_model.embeddings.patch_embedding"""							)
	if "visual.ln_pre" in name:
		lowerCamelCase__										=			name.replace("""visual.ln_pre"""			,    """vision_model.pre_layernorm"""							)
	if "visual.ln_post" in name:
		lowerCamelCase__										=			name.replace("""visual.ln_post"""			,    """vision_model.post_layernorm"""							)
	if "visual.proj" in name:
		lowerCamelCase__										=			name.replace("""visual.proj"""			,    """visual_projection.weight"""							)
	if "text_projection" in name:
		lowerCamelCase__										=			name.replace("""text_projection"""			,    """text_projection.weight"""							)
	# things on top
	if "prompts_visual_proj" in name:
		lowerCamelCase__										=			name.replace("""prompts_visual_proj"""			,    """prompts_visual_projection"""							)
	if "prompts_visual_ln" in name:
		lowerCamelCase__										=			name.replace("""prompts_visual_ln"""			,    """prompts_visual_layernorm"""							)
	# mit
	if name == "mit.positional_embedding":
		lowerCamelCase__										=			name.replace("""positional"""			,    """position"""							)
	if name.startswith("""mit.resblocks"""							):
		lowerCamelCase__										=			name.replace("""mit.resblocks"""			,    """mit.encoder.layers"""							)
	# prompts generator
	if name.startswith("""prompts_generator.norm"""							):
		lowerCamelCase__										=			name.replace("""prompts_generator.norm"""			,    """prompts_generator.layernorm"""							)

	return name

def     snake_case							(UpperCamelCase			:     Optional[Any]			,    UpperCamelCase			:     Union[str, Any]							):




	'''simple docstring'''







	for key in orig_state_dict.copy().keys():
		lowerCamelCase__										=			orig_state_dict.pop(lowercase_							)

		if "attn.in_proj" in key:
			lowerCamelCase__										=			key.split("""."""							)
			if key.startswith("""visual"""							):
				lowerCamelCase__										=			key_split[3]
				lowerCamelCase__										=			config.vision_config.hidden_size
				if "message_attn" in key:
					if "weight" in key:
						lowerCamelCase__										=			val[
						    :dim, :
						]
						lowerCamelCase__										=			val[
						    dim : dim * 2, :
						]
						lowerCamelCase__										=			val[
						    -dim:, :
						]
					else:
						lowerCamelCase__										=			val[
						    :dim
						]
						lowerCamelCase__										=			val[
						    dim : dim * 2
						]
						lowerCamelCase__										=			val[
						    -dim:
						]
				else:
					if "weight" in key:
						lowerCamelCase__										=			val[
						    :dim, :
						]
						lowerCamelCase__										=			val[
						    dim : dim * 2, :
						]
						lowerCamelCase__										=			val[
						    -dim:, :
						]
					else:
						lowerCamelCase__										=			val[:dim]
						lowerCamelCase__										=			val[
						    dim : dim * 2
						]
						lowerCamelCase__										=			val[-dim:]
			elif key.startswith("""mit"""							):
				lowerCamelCase__										=			key_split[2]
				lowerCamelCase__										=			config.vision_config.mit_hidden_size
				if "weight" in key:
					lowerCamelCase__										=			val[:dim, :]
					lowerCamelCase__										=			val[dim : dim * 2, :]
					lowerCamelCase__										=			val[-dim:, :]
				else:
					lowerCamelCase__										=			val[:dim]
					lowerCamelCase__										=			val[dim : dim * 2]
					lowerCamelCase__										=			val[-dim:]
			else:
				lowerCamelCase__										=			key_split[2]
				lowerCamelCase__										=			config.text_config.hidden_size
				if "weight" in key:
					lowerCamelCase__										=			val[:dim, :]
					lowerCamelCase__										=			val[
					    dim : dim * 2, :
					]
					lowerCamelCase__										=			val[-dim:, :]
				else:
					lowerCamelCase__										=			val[:dim]
					lowerCamelCase__										=			val[
					    dim : dim * 2
					]
					lowerCamelCase__										=			val[-dim:]
		else:
			lowerCamelCase__										=			rename_key(lowercase_							)
			if new_key_name in ["visual_projection.weight", "text_projection.weight"]:
				lowerCamelCase__										=			val.T
			lowerCamelCase__										=			val

	return orig_state_dict

def     snake_case							(UpperCamelCase			:     Optional[int]							):




	'''simple docstring'''







	if num_frames == 8:
		lowerCamelCase__										=			"""eating_spaghetti_8_frames.npy"""
	elif num_frames == 16:
		lowerCamelCase__										=			"""eating_spaghetti.npy"""
	elif num_frames == 32:
		lowerCamelCase__										=			"""eating_spaghetti_32_frames.npy"""
	lowerCamelCase__										=			hf_hub_download(
	    repo_id="""hf-internal-testing/spaghetti-video"""			,    filename=lowercase_			,    repo_type="""dataset"""			,    )
	lowerCamelCase__										=			np.load(lowercase_							)
	return list(lowercase_							)

def     snake_case							(UpperCamelCase			:     Optional[Any]			,    UpperCamelCase			:     Any=None			,    UpperCamelCase			:     List[Any]=False							):




	'''simple docstring'''







	lowerCamelCase__										=			{
	    # fully supervised kinetics-400 checkpoints
	    """xclip-base-patch32""": """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_8.pth""",
	    """xclip-base-patch32-16-frames""": (
	        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_32_16.pth"""
	    ),
	    """xclip-base-patch16""": """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_8.pth""",
	    """xclip-base-patch16-16-frames""": (
	        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k400_16_16.pth"""
	    ),
	    """xclip-large-patch14""": """https://drive.google.com/u/0/uc?id=1NUOImq0o5DlQTST17iIP3vG7DgmHQuCx&amp;export=download&amp;confirm=t&amp;uuid=b26caedc-88e2-473e-830a-9d158b653cdb""",
	    """xclip-large-patch14-16-frames""": """https://drive.google.com/u/0/uc?id=1FOYgnJc097OJ4lGwtRCCydQyVPJEOH7d&amp;export=download&amp;confirm=t&amp;uuid=538fa810-e671-4050-b385-9a623f89804f""",
	    # fully supervised kinetics-600 checkpoints
	    """xclip-base-patch16-kinetics-600""": (
	        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_8.pth"""
	    ),
	    """xclip-base-patch16-kinetics-600-16-frames""": (
	        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/k600_16_16.pth"""
	    ),
	    """xclip-large-patch14-kinetics-600""": """https://drive.google.com/u/0/uc?id=1FV8C1INuM91sLAN4ImjzePLIlpMSihwV&amp;export=download&amp;confirm=t&amp;uuid=141d4977-4a65-44ae-864f-4b0c19f838be""",
	    # few shot
	    """xclip-base-patch16-hmdb-2-shot""": (
	        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_2.pth"""
	    ),
	    """xclip-base-patch16-hmdb-4-shot""": (
	        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_4.pth"""
	    ),
	    """xclip-base-patch16-hmdb-8-shot""": (
	        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_8.pth"""
	    ),
	    """xclip-base-patch16-hmdb-16-shot""": (
	        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_hmdb_16.pth"""
	    ),
	    """xclip-base-patch16-ucf-2-shot""": (
	        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_2.pth"""
	    ),
	    """xclip-base-patch16-ucf-4-shot""": (
	        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_4.pth"""
	    ),
	    """xclip-base-patch16-ucf-8-shot""": (
	        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_8.pth"""
	    ),
	    """xclip-base-patch16-ucf-16-shot""": (
	        """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/few_ucf_16.pth"""
	    ),
	    # zero shot
	    """xclip-base-patch16-zero-shot""": """https://github.com/nbl97/X-CLIP_Model_Zoo/releases/download/v1.0/zero.pth""",
	}

	lowerCamelCase__										=			model_to_url[model_name]
	lowerCamelCase__										=			8
	if "16-frames" in model_name:
		lowerCamelCase__										=			16
	elif "shot" in model_name:
		lowerCamelCase__										=			32

	lowerCamelCase__										=			get_xclip_config(lowercase_			,    lowercase_							)
	lowerCamelCase__										=			XCLIPModel(lowercase_							)
	model.eval()

	if "drive" in checkpoint_url:
		lowerCamelCase__										=			"""pytorch_model.bin"""
		gdown.cached_download(lowercase_			,    lowercase_			,    quiet=lowercase_							)
		lowerCamelCase__										=			torch.load(lowercase_			,    map_location="""cpu"""							)["""model"""]
	else:
		lowerCamelCase__										=			torch.hub.load_state_dict_from_url(lowercase_							)["""model"""]

	lowerCamelCase__										=			convert_state_dict(lowercase_			,    lowercase_							)

	lowerCamelCase__										=			XCLIPModel(lowercase_							)
	lowerCamelCase__							,							lowerCamelCase__										=			model.load_state_dict(lowercase_			,    strict=lowercase_							)
	assert missing_keys == ["text_model.embeddings.position_ids", "vision_model.embeddings.position_ids"]
	model.eval()

	lowerCamelCase__										=			336 if model_name == """xclip-large-patch14-16-frames""" else 224
	lowerCamelCase__										=			VideoMAEImageProcessor(size=lowercase_							)
	lowerCamelCase__										=			CLIPTokenizer.from_pretrained("""openai/clip-vit-base-patch32"""							)
	lowerCamelCase__										=			CLIPTokenizerFast.from_pretrained("""openai/clip-vit-base-patch32"""							)
	lowerCamelCase__										=			XCLIPProcessor(image_processor=lowercase_			,    tokenizer=lowercase_							)

	lowerCamelCase__										=			prepare_video(lowercase_							)
	lowerCamelCase__										=			processor(
	    text=["""playing sports""", """eating spaghetti""", """go shopping"""]			,    videos=lowercase_			,    return_tensors="""pt"""			,    padding=lowercase_							)

	print("""Shape of pixel values:"""			,    inputs.pixel_values.shape							)

	with torch.no_grad():
		lowerCamelCase__										=			model(**lowercase_							)

	# Verify outputs
	lowerCamelCase__										=			outputs.logits_per_video
	lowerCamelCase__										=			logits_per_video.softmax(dim=1							)
	print("""Probs:"""			,    lowercase_							)
	# kinetics-400
	if model_name == "xclip-base-patch32":
		lowerCamelCase__										=			torch.tensor([[0.0_0_1_9, 0.9_9_5_1, 0.0_0_3_0]]							)
	elif model_name == "xclip-base-patch32-16-frames":
		lowerCamelCase__										=			torch.tensor([[7.0999e-04, 9.9883e-01, 4.5580e-04]]							)
	elif model_name == "xclip-base-patch16":
		lowerCamelCase__										=			torch.tensor([[0.0_0_8_3, 0.9_6_8_1, 0.0_2_3_6]]							)
	elif model_name == "xclip-base-patch16-16-frames":
		lowerCamelCase__										=			torch.tensor([[7.6937e-04, 9.9728e-01, 1.9473e-03]]							)
	elif model_name == "xclip-large-patch14":
		lowerCamelCase__										=			torch.tensor([[0.0_0_6_2, 0.9_8_6_4, 0.0_0_7_5]]							)
	elif model_name == "xclip-large-patch14-16-frames":
		lowerCamelCase__										=			torch.tensor([[3.3877e-04, 9.9937e-01, 2.8888e-04]]							)
	# kinetics-600
	elif model_name == "xclip-base-patch16-kinetics-600":
		lowerCamelCase__										=			torch.tensor([[0.0_5_5_5, 0.8_9_1_4, 0.0_5_3_1]]							)
	elif model_name == "xclip-base-patch16-kinetics-600-16-frames":
		lowerCamelCase__										=			torch.tensor([[3.8554e-04, 9.9929e-01, 3.2754e-04]]							)
	elif model_name == "xclip-large-patch14-kinetics-600":
		lowerCamelCase__										=			torch.tensor([[0.0_0_3_6, 0.9_9_2_0, 0.0_0_4_5]]							)
	# few shot
	elif model_name == "xclip-base-patch16-hmdb-2-shot":
		lowerCamelCase__										=			torch.tensor([[7.1890e-06, 9.9994e-01, 5.6559e-05]]							)
	elif model_name == "xclip-base-patch16-hmdb-4-shot":
		lowerCamelCase__										=			torch.tensor([[1.0320e-05, 9.9993e-01, 6.2435e-05]]							)
	elif model_name == "xclip-base-patch16-hmdb-8-shot":
		lowerCamelCase__										=			torch.tensor([[4.1377e-06, 9.9990e-01, 9.8386e-05]]							)
	elif model_name == "xclip-base-patch16-hmdb-16-shot":
		lowerCamelCase__										=			torch.tensor([[4.1347e-05, 9.9962e-01, 3.3411e-04]]							)
	elif model_name == "xclip-base-patch16-ucf-2-shot":
		lowerCamelCase__										=			torch.tensor([[8.5857e-05, 9.9928e-01, 6.3291e-04]]							)
	elif model_name == "xclip-base-patch16-ucf-4-shot":
		lowerCamelCase__										=			torch.tensor([[8.5857e-05, 9.9928e-01, 6.3291e-04]]							)
	elif model_name == "xclip-base-patch16-ucf-8-shot":
		lowerCamelCase__										=			torch.tensor([[0.0_0_2_7, 0.9_9_0_4, 0.0_0_7_0]]							)
	elif model_name == "xclip-base-patch16-ucf-16-shot":
		lowerCamelCase__										=			torch.tensor([[9.8219e-04, 9.9593e-01, 3.0863e-03]]							)
	# zero shot
	elif model_name == "xclip-base-patch16-zero-shot":
		lowerCamelCase__										=			torch.tensor([[3.5082e-04, 9.9785e-01, 1.7966e-03]]							)
	else:
		raise ValueError(f'''Model name {model_name} not supported'''							)
	assert torch.allclose(lowercase_			,    lowercase_			,    atol=1e-3							)
	print("""Looks ok!"""							)

	if pytorch_dump_folder_path is not None:
		print(f'''Saving model {model_name} to {pytorch_dump_folder_path}'''							)
		model.save_pretrained(lowercase_							)

	if push_to_hub:
		print("""Pushing model, processor and slow tokenizer files to the hub..."""							)
		model.push_to_hub(lowercase_			,    organization="""nielsr"""							)
		processor.push_to_hub(lowercase_			,    organization="""nielsr"""							)
		slow_tokenizer.push_to_hub(lowercase_			,    organization="""nielsr"""							)


if __name__ == "__main__":
					a__    :		Tuple      						=     argparse.ArgumentParser()
					# Required parameters
					parser.add_argument(
					    """--model_name""",
					    default="""xclip-base-patch32""",
					    type=str,
					    help="""Name of the model.""",
					)
					parser.add_argument(
					    """--pytorch_dump_folder_path""", default=None, type=str, help="""Path to the output PyTorch model directory."""
					)
					parser.add_argument(
					    """--push_to_hub""", action="""store_true""", help="""Whether or not to push the converted model to the 🤗 hub."""
					)

					a__    :		Optional[int]      						=     parser.parse_args()
					convert_xclip_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)








'''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_12, 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       __UpperCamelCase					(      lowercase_       :  List[Any]   ,       lowercase_       :  float = 0.0   ,       lowercase_       :  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(lowercase_   ,       dtype=input.dtype   ,       device=input.device   )
       random_tensor.floor_()  # binarize
       a_						     =			input.div(lowercase_   ) * random_tensor
       return output




class   __SCREAMING_SNAKE_CASE       (nn.Module       ):


  """simple docstring"""


  def __init__(			self   ,	UpperCamelCase__ = None	):






         """simple docstring"""


         super().__init__()
         a_						     =			drop_prob







  def    _a       (			self   ,	UpperCamelCase__	):






         """simple docstring"""


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







  def    _a       (			self	):






         """simple docstring"""


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





class   __SCREAMING_SNAKE_CASE       (nn.Module       ):


  """simple docstring"""


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






         """simple docstring"""


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

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







  def    _a       (			self   ,	UpperCamelCase__	):






         """simple docstring"""


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





class   __SCREAMING_SNAKE_CASE       (nn.GroupNorm       ):


  """simple docstring"""


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






         """simple docstring"""


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





class   __SCREAMING_SNAKE_CASE       (nn.Module       ):


  """simple docstring"""


  def __init__(			self   ,	UpperCamelCase__	):






         """simple docstring"""


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







  def    _a       (			self   ,	UpperCamelCase__	):






         """simple docstring"""


         return self.pool(UpperCamelCase__	) - hidden_states





class   __SCREAMING_SNAKE_CASE       (nn.Module       ):


  """simple docstring"""


  def __init__(			self   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__	):






         """simple docstring"""


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







  def    _a       (			self   ,	UpperCamelCase__	):






         """simple docstring"""


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

         return hidden_states





class   __SCREAMING_SNAKE_CASE       (nn.Module       ):


  """simple docstring"""


  def __init__(			self   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__   ,	UpperCamelCase__	):






         """simple docstring"""


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

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







  def    _a       (			self   ,	UpperCamelCase__	):






         """simple docstring"""


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

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

                a_						     =			(output,) + outputs
                return outputs

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

                a_						     =			(output,) + outputs
                return outputs





class   __SCREAMING_SNAKE_CASE       (nn.Module       ):


  """simple docstring"""


  def __init__(			self   ,	UpperCamelCase__	):






         """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(UpperCamelCase__	)

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

         a_						     =			nn.ModuleList(UpperCamelCase__	)







  def    _a       (			self   ,	UpperCamelCase__   ,	UpperCamelCase__=False   ,	UpperCamelCase__=True	):






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





class   __SCREAMING_SNAKE_CASE       (__A       ):


  """simple docstring"""


  _a  : int      			=     PoolFormerConfig
  _a  : Optional[Any]      			=     '''poolformer'''
  _a  : Union[str, Any]      			=     '''pixel_values'''
  _a  : Optional[int]      			=     True







  def    _a       (			self   ,	UpperCamelCase__	):






         """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   ,	UpperCamelCase__   ,	UpperCamelCase__=False	):






         """simple docstring"""


         if isinstance(UpperCamelCase__   ,	UpperCamelCase__	):
                a_						     =			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       ):


  """simple docstring"""


  def __init__(			self   ,	UpperCamelCase__	):






         """simple docstring"""


         super().__init__(UpperCamelCase__	)
         a_						     =			config

         a_						     =			PoolFormerEncoder(UpperCamelCase__	)

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







  def    _a       (			self	):






         """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   ,	UpperCamelCase__ = None   ,	UpperCamelCase__ = None   ,	UpperCamelCase__ = None   ,	):






         """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(
             UpperCamelCase__   ,	output_hidden_states=UpperCamelCase__   ,	return_dict=UpperCamelCase__   ,	)
         a_						     =			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       ):


  """simple docstring"""


  def __init__(			self   ,	UpperCamelCase__	):






         """simple docstring"""


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







  def    _a       (			self   ,	UpperCamelCase__	):






         """simple docstring"""


         a_						     =			self.dense(UpperCamelCase__	)
         return output





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


  """simple docstring"""


  def __init__(			self   ,	UpperCamelCase__	):






         """simple docstring"""


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

         # 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(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   ,	UpperCamelCase__ = None   ,	UpperCamelCase__ = None   ,	UpperCamelCase__ = None   ,	UpperCamelCase__ = None   ,	):






         """simple docstring"""


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

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

         a_						     =			outputs[0]

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

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

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




