
'''simple docstring'''
import argparse
import json
from pathlib import Path

import requests
import torch
from huggingface_hub import hf_hub_download
from PIL import Image

from transformers import (
    BertTokenizer,
    ViltConfig,
    ViltForImageAndTextRetrieval,
    ViltForImagesAndTextClassification,
    ViltForMaskedLM,
    ViltForQuestionAnswering,
    ViltImageProcessor,
    ViltProcessor,
)
from transformers.utils import logging


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






def      UpperCamelCase_				(      A__  ,      A__=False  ,      A__=False  ,      A__=False						):
	a_      				=						[]
	for i in range(config.num_hidden_layers						):
		# encoder layers: output projection, 2 feedforward neural networks and 2 layernorms
		rename_keys.append((F'''transformer.blocks.{i}.norm1.weight''', F'''vilt.encoder.layer.{i}.layernorm_before.weight''')						)
		rename_keys.append((F'''transformer.blocks.{i}.norm1.bias''', F'''vilt.encoder.layer.{i}.layernorm_before.bias''')						)
		rename_keys.append(
		    (F'''transformer.blocks.{i}.attn.proj.weight''', F'''vilt.encoder.layer.{i}.attention.output.dense.weight''')						)
		rename_keys.append(
		    (F'''transformer.blocks.{i}.attn.proj.bias''', F'''vilt.encoder.layer.{i}.attention.output.dense.bias''')						)
		rename_keys.append((F'''transformer.blocks.{i}.norm2.weight''', F'''vilt.encoder.layer.{i}.layernorm_after.weight''')						)
		rename_keys.append((F'''transformer.blocks.{i}.norm2.bias''', F'''vilt.encoder.layer.{i}.layernorm_after.bias''')						)
		rename_keys.append(
		    (F'''transformer.blocks.{i}.mlp.fc1.weight''', F'''vilt.encoder.layer.{i}.intermediate.dense.weight''')						)
		rename_keys.append((F'''transformer.blocks.{i}.mlp.fc1.bias''', F'''vilt.encoder.layer.{i}.intermediate.dense.bias''')						)
		rename_keys.append((F'''transformer.blocks.{i}.mlp.fc2.weight''', F'''vilt.encoder.layer.{i}.output.dense.weight''')						)
		rename_keys.append((F'''transformer.blocks.{i}.mlp.fc2.bias''', F'''vilt.encoder.layer.{i}.output.dense.bias''')						)

	# embeddings
	rename_keys.extend(
	    [
	        # text embeddings
	        ("""text_embeddings.word_embeddings.weight""", """vilt.embeddings.text_embeddings.word_embeddings.weight"""),
	        (
	            """text_embeddings.position_embeddings.weight""",
	            """vilt.embeddings.text_embeddings.position_embeddings.weight""",
	        ),
	        ("""text_embeddings.position_ids""", """vilt.embeddings.text_embeddings.position_ids"""),
	        (
	            """text_embeddings.token_type_embeddings.weight""",
	            """vilt.embeddings.text_embeddings.token_type_embeddings.weight""",
	        ),
	        ("""text_embeddings.LayerNorm.weight""", """vilt.embeddings.text_embeddings.LayerNorm.weight"""),
	        ("""text_embeddings.LayerNorm.bias""", """vilt.embeddings.text_embeddings.LayerNorm.bias"""),
	        # patch embeddings
	        ("""transformer.cls_token""", """vilt.embeddings.cls_token"""),
	        ("""transformer.patch_embed.proj.weight""", """vilt.embeddings.patch_embeddings.projection.weight"""),
	        ("""transformer.patch_embed.proj.bias""", """vilt.embeddings.patch_embeddings.projection.bias"""),
	        ("""transformer.pos_embed""", """vilt.embeddings.position_embeddings"""),
	        # token type embeddings
	        ("""token_type_embeddings.weight""", """vilt.embeddings.token_type_embeddings.weight"""),
	    ]						)

	# final layernorm + pooler
	rename_keys.extend(
	    [
	        ("""transformer.norm.weight""", """vilt.layernorm.weight"""),
	        ("""transformer.norm.bias""", """vilt.layernorm.bias"""),
	        ("""pooler.dense.weight""", """vilt.pooler.dense.weight"""),
	        ("""pooler.dense.bias""", """vilt.pooler.dense.bias"""),
	    ]						)

	# classifier head(s)
	if vqa_model:
		# classification head
		rename_keys.extend(
		    [
		        ("""vqa_classifier.0.weight""", """classifier.0.weight"""),
		        ("""vqa_classifier.0.bias""", """classifier.0.bias"""),
		        ("""vqa_classifier.1.weight""", """classifier.1.weight"""),
		        ("""vqa_classifier.1.bias""", """classifier.1.bias"""),
		        ("""vqa_classifier.3.weight""", """classifier.3.weight"""),
		        ("""vqa_classifier.3.bias""", """classifier.3.bias"""),
		    ]						)
	elif nlvr_model:
		# classification head
		rename_keys.extend(
		    [
		        ("""nlvr2_classifier.0.weight""", """classifier.0.weight"""),
		        ("""nlvr2_classifier.0.bias""", """classifier.0.bias"""),
		        ("""nlvr2_classifier.1.weight""", """classifier.1.weight"""),
		        ("""nlvr2_classifier.1.bias""", """classifier.1.bias"""),
		        ("""nlvr2_classifier.3.weight""", """classifier.3.weight"""),
		        ("""nlvr2_classifier.3.bias""", """classifier.3.bias"""),
		    ]						)
	else:
		pass

	return rename_keys






def      UpperCamelCase_				(      A__  ,      A__						):
	for i in range(config.num_hidden_layers						):
		a_      				=						"""vilt."""
		# read in weights + bias of input projection layer (in timm, this is a single matrix + bias)
		a_      				=						state_dict.pop(F'''transformer.blocks.{i}.attn.qkv.weight'''						)
		a_      				=						state_dict.pop(F'''transformer.blocks.{i}.attn.qkv.bias'''						)
		# next, add query, keys and values (in that order) to the state dict
		a_      				=						in_proj_weight[
		    : config.hidden_size, :
		]
		a_      				=						in_proj_bias[: config.hidden_size]
		a_      				=						in_proj_weight[
		    config.hidden_size : config.hidden_size * 2, :
		]
		a_      				=						in_proj_bias[
		    config.hidden_size : config.hidden_size * 2
		]
		a_      				=						in_proj_weight[
		    -config.hidden_size :, :
		]
		a_      				=						in_proj_bias[-config.hidden_size :]






def      UpperCamelCase_				(      A__						):
	a_      				=						["""head.weight""", """head.bias"""]
	for k in ignore_keys:
		state_dict.pop(A__  ,      A__						)






def      UpperCamelCase_				(      A__  ,      A__  ,      A__						):
	a_      				=						dct.pop(A__						)
	a_      				=						val






@torch.no_grad()
def      UpperCamelCase_				(      A__  ,      A__						):
	a_      				=						ViltConfig(image_size=3_84  ,      patch_size=32  ,      tie_word_embeddings=A__						)
	a_      				=						False
	a_      				=						False
	a_      				=						False
	a_      				=						False
	if "vqa" in checkpoint_url:
		a_      				=						True
		a_      				=						31_29
		a_      				=						"""huggingface/label-files"""
		a_      				=						"""vqa2-id2label.json"""
		a_      				=						json.load(open(hf_hub_download(A__  ,      A__  ,      repo_type="""dataset"""						)  ,      """r"""						)						)
		a_      				=						{int(A__						): v for k, v in idalabel.items()}
		a_      				=						idalabel
		a_      				=						{v: k for k, v in idalabel.items()}
		a_      				=						ViltForQuestionAnswering(A__						)
	elif "nlvr" in checkpoint_url:
		a_      				=						True
		a_      				=						2
		a_      				=						{0: """False""", 1: """True"""}
		a_      				=						{v: k for k, v in config.idalabel.items()}
		a_      				=						3
		a_      				=						ViltForImagesAndTextClassification(A__						)
	elif "irtr" in checkpoint_url:
		a_      				=						True
		a_      				=						ViltForImageAndTextRetrieval(A__						)
	elif "mlm_itm" in checkpoint_url:
		a_      				=						True
		a_      				=						ViltForMaskedLM(A__						)
	else:
		raise ValueError("""Unknown model type"""						)

	# load state_dict of original model, remove and rename some keys
	a_      				=						torch.hub.load_state_dict_from_url(A__  ,      map_location="""cpu"""						)["""state_dict"""]
	a_      				=						create_rename_keys(A__  ,      A__  ,      A__  ,      A__						)
	for src, dest in rename_keys:
		rename_key(A__  ,      A__  ,      A__						)
	read_in_q_k_v(A__  ,      A__						)
	if mlm_model or irtr_model:
		a_      				=						["""itm_score.fc.weight""", """itm_score.fc.bias"""]
		for k in ignore_keys:
			state_dict.pop(A__  ,      A__						)

    # load state dict into HuggingFace model
	model.eval()
	if mlm_model:
		a_				,  a_      				=						model.load_state_dict(A__  ,      strict=A__						)
		assert missing_keys == ["mlm_score.decoder.bias"]
	else:
		model.load_state_dict(A__						)

	# Define processor
	a_      				=						ViltImageProcessor(size=3_84						)
	a_      				=						BertTokenizer.from_pretrained("""bert-base-uncased"""						)
	a_      				=						ViltProcessor(A__  ,      A__						)

	# Forward pass on example inputs (image + text)
	if nlvr_model:
		a_      				=						Image.open(requests.get("""https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg"""  ,      stream=A__						).raw						)
		a_      				=						Image.open(requests.get("""https://lil.nlp.cornell.edu/nlvr/exs/ex0_0.jpg"""  ,      stream=A__						).raw						)
		a_      				=						(
		    """The left image contains twice the number of dogs as the right image, and at least two dogs in total are"""
		    """ standing."""
		)
		a_      				=						processor(A__  ,      A__  ,      return_tensors="""pt"""						)
		a_      				=						processor(A__  ,      A__  ,      return_tensors="""pt"""						)
		a_      				=						model(
		    input_ids=encoding_a.input_ids  ,      pixel_values=encoding_a.pixel_values  ,      pixel_values_a=encoding_a.pixel_values  ,      )
	else:
		a_      				=						Image.open(requests.get("""http://images.cocodataset.org/val2017/000000039769.jpg"""  ,      stream=A__						).raw						)
		if mlm_model:
			a_      				=						"""a bunch of [MASK] laying on a [MASK]."""
		else:
			a_      				=						"""How many cats are there?"""
		a_      				=						processor(A__  ,      A__  ,      return_tensors="""pt"""						)
		a_      				=						model(**A__						)

	# Verify outputs
	if mlm_model:
		a_      				=						torch.Size([1, 11, 3_05_22]						)
		a_      				=						torch.tensor([-12.5061, -12.5123, -12.5174]						)
		assert outputs.logits.shape == expected_shape
		assert torch.allclose(outputs.logits[0, 0, :3]  ,      A__  ,      atol=1e-4						)

		# verify masked token prediction equals "cats"
		a_      				=						outputs.logits[0, 4, :].argmax(-1						).item()
		assert tokenizer.decode([predicted_id]						) == "cats"
	elif vqa_model:
		a_      				=						torch.Size([1, 31_29]						)
		a_      				=						torch.tensor([-15.9495, -18.1472, -10.3041]						)
		assert torch.allclose(outputs.logits[0, :3]  ,      A__  ,      atol=1e-4						)
		assert outputs.logits.shape == expected_shape
		assert torch.allclose(outputs.logits[0, 0, :3]  ,      A__  ,      atol=1e-4						)

		# verify vqa prediction equals "2"
		a_      				=						outputs.logits.argmax(-1						).item()
		assert model.config.idalabel[predicted_idx] == "2"
	elif nlvr_model:
		a_      				=						torch.Size([1, 2]						)
		a_      				=						torch.tensor([-2.8721, 2.1291]						)
		assert torch.allclose(outputs.logits[0, :3]  ,      A__  ,      atol=1e-4						)
		assert outputs.logits.shape == expected_shape

	Path(A__						).mkdir(exist_ok=A__						)
	print(F'''Saving model and processor to {pytorch_dump_folder_path}'''						)
	model.save_pretrained(A__						)
	processor.save_pretrained(A__						)


if __name__ == "__main__":
		lowercase__											=argparse.ArgumentParser()
		# Required parameters
		parser.add_argument(
		    '--checkpoint_url',
		    default='https://github.com/dandelin/ViLT/releases/download/200k/vilt_200k_mlm_itm.ckpt',
		    type=str,
		    help='URL of the checkpoint you\'d like to convert.',
		)
		parser.add_argument(
		    '--pytorch_dump_folder_path', default=None, type=str, help='Path to the output PyTorch model directory.'
		)

		lowercase__											=parser.parse_args()
		convert_vilt_checkpoint(args.checkpoint_url, args.pytorch_dump_folder_path)


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


lowercase__											=logging.get_logger(__name__)

# General docstring
lowercase__											='PoolFormerConfig'

# Base docstring
lowercase__											='sail/poolformer_s12'
lowercase__											=[1, 5_12, 7, 7]

# Image classification docstring
lowercase__											='sail/poolformer_s12'
lowercase__											='tabby, tabby cat'

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






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







class   a_ (      nn.Module					):




			def __init__(       self					,   UpperCAmelCase = None	):
				super().__init__()
				a_      				=						drop_prob




			def 					lowerCAmelCase__					(       self					,   UpperCAmelCase	):
				return drop_path(UpperCAmelCase					,   self.drop_prob					,   self.training	)






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







class   a_ (      nn.Module					):




			def __init__(       self					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase=None	):
				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 					lowerCAmelCase__					(       self					,   UpperCAmelCase	):
				a_      				=						self.projection(UpperCAmelCase	)
				a_      				=						self.norm(UpperCAmelCase	)
				return embeddings







class   a_ (      nn.GroupNorm					):




			def __init__(       self					,   UpperCAmelCase					,   **UpperCAmelCase	):
				super().__init__(1					,   UpperCAmelCase					,   **UpperCAmelCase	)







class   a_ (      nn.Module					):




			def __init__(       self					,   UpperCAmelCase	):
				super().__init__()
				a_      				=						nn.AvgPoolad(UpperCAmelCase					,   stride=1					,   padding=pool_size // 2					,   count_include_pad=UpperCAmelCase	)






			def 					lowerCAmelCase__					(       self					,   UpperCAmelCase	):
				return self.pool(UpperCAmelCase	) - hidden_states







class   a_ (      nn.Module					):




			def __init__(       self					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase	):
				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 					lowerCAmelCase__					(       self					,   UpperCAmelCase	):
				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   a_ (      nn.Module					):




			def __init__(       self					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase					,   UpperCAmelCase	):
				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 					lowerCAmelCase__					(       self					,   UpperCAmelCase	):
				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   a_ (      nn.Module					):




			def __init__(       self					,   UpperCAmelCase	):
				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 					lowerCAmelCase__					(       self					,   UpperCAmelCase					,   UpperCAmelCase=False					,   UpperCAmelCase=True	):
				a_      				=						() if output_hidden_states else None

				a_      				=						pixel_values
				for idx, layers in enumerate(zip(self.patch_embeddings					,   self.block	)	):
					a_				,  a_      				=						layers
					# Get patch embeddings from hidden_states
					a_      				=						embedding_layer(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   a_ (      UpperCamelCase__					):
			lowerCamelCase__	:			Union[str, Any]       =    PoolFormerConfig
			lowerCamelCase__	:			Optional[Any]       =    'poolformer'
			lowerCamelCase__	:			List[Any]       =    'pixel_values'
			lowerCamelCase__	:			int       =    True




			def 					lowerCAmelCase__					(       self					,   UpperCAmelCase	):
				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 					lowerCAmelCase__					(       self					,   UpperCAmelCase					,   UpperCAmelCase=False	):
				if isinstance(UpperCAmelCase					,   UpperCAmelCase	):
					a_      				=						value







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

lowercase__											=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.'		,  UpperCamelCase__		,  )
class   a_ (      UpperCamelCase__					):




			def __init__(       self					,   UpperCAmelCase	):
				super().__init__(UpperCAmelCase	)
				a_      				=						config

				a_      				=						PoolFormerEncoder(UpperCAmelCase	)

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




			def 					lowerCAmelCase__					(       self	):
				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 					lowerCAmelCase__					(       self					,   UpperCAmelCase = None					,   UpperCAmelCase = None					,   UpperCAmelCase = None					,   ):
				a_      				=						(
				    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
				)
				a_      				=						return_dict if return_dict is not None else self.config.use_return_dict

				if pixel_values is None:
					raise ValueError("""You have to specify pixel_values"""	)

				a_      				=						self.encoder(
				    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   a_ (      nn.Module					):




			def __init__(       self					,   UpperCAmelCase	):
				super().__init__()
				a_      				=						nn.Linear(config.hidden_size					,   config.hidden_size	)






			def 					lowerCAmelCase__					(       self					,   UpperCAmelCase	):
				a_      				=						self.dense(UpperCAmelCase	)
				return output







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




			def __init__(       self					,   UpperCAmelCase	):
				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 					lowerCAmelCase__					(       self					,   UpperCAmelCase = None					,   UpperCAmelCase = None					,   UpperCAmelCase = None					,   UpperCAmelCase = None					,   ):
				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	)
