





import argparse
import json

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

from transformers import ConvNextConfig, SegformerImageProcessor, UperNetConfig, UperNetForSemanticSegmentation



def     _lowerCAmelCase		(  _lowerCAmelCase     ):
			'''simple docstring'''


			A_     :				Dict			      =	3_8_4
			if "tiny" in model_name:
						A_     :				Optional[Any]			      =	[3, 3, 9, 3]
						A_     :				str			      =	[9_6, 1_9_2, 3_8_4, 7_6_8]
			if "small" in model_name:
						A_     :				List[Any]			      =	[3, 3, 2_7, 3]
						A_     :				Union[str, Any]			      =	[9_6, 1_9_2, 3_8_4, 7_6_8]
			if "base" in model_name:
						A_     :				Any			      =	[3, 3, 2_7, 3]
						A_     :				Dict			      =	[1_2_8, 2_5_6, 5_1_2, 1_0_2_4]
						A_     :				Any			      =	5_1_2
			if "large" in model_name:
						A_     :				Optional[int]			      =	[3, 3, 2_7, 3]
						A_     :				int			      =	[1_9_2, 3_8_4, 7_6_8, 1_5_3_6]
						A_     :				Optional[Any]			      =	7_6_8
			if "xlarge" in model_name:
						A_     :				Union[str, Any]			      =	[3, 3, 2_7, 3]
						A_     :				Optional[Any]			      =	[2_5_6, 5_1_2, 1_0_2_4, 2_0_4_8]
						A_     :				int			      =	1_0_2_4

			# set label information
			A_     :				Tuple			      =	1_5_0
			A_     :				Optional[int]			      =	"""huggingface/label-files"""
			A_     :				Any			      =	"""ade20k-id2label.json"""
			A_     :				Dict			      =	json.load(open(hf_hub_download(_lowerCAmelCase  ,_lowerCAmelCase  ,repo_type="""dataset"""     )  ,"""r"""     )     )
			A_     :				List[str]			      =	{int(_lowerCAmelCase     ): v for k, v in idalabel.items()}
			A_     :				Any			      =	{v: k for k, v in idalabel.items()}

			A_     :				List[Any]			      =	ConvNextConfig(
			    depths=_lowerCAmelCase  ,hidden_sizes=_lowerCAmelCase  ,out_features=["""stage1""", """stage2""", """stage3""", """stage4"""]     )
			A_     :				Any			      =	UperNetConfig(
			    backbone_config=_lowerCAmelCase  ,auxiliary_in_channels=_lowerCAmelCase  ,num_labels=_lowerCAmelCase  ,idalabel=_lowerCAmelCase  ,labelaid=_lowerCAmelCase  ,)

			return config



def     _lowerCAmelCase		(  _lowerCAmelCase     ):
			'''simple docstring'''


			A_     :				int			      =	[]

			# fmt: off
			# stem
			rename_keys.append(("""backbone.downsample_layers.0.0.weight""", """backbone.embeddings.patch_embeddings.weight""")     )
			rename_keys.append(("""backbone.downsample_layers.0.0.bias""", """backbone.embeddings.patch_embeddings.bias""")     )
			rename_keys.append(("""backbone.downsample_layers.0.1.weight""", """backbone.embeddings.layernorm.weight""")     )
			rename_keys.append(("""backbone.downsample_layers.0.1.bias""", """backbone.embeddings.layernorm.bias""")     )
			# stages
			for i in range(len(config.backbone_config.depths     )     ):
						for j in range(config.backbone_config.depths[i]     ):
									rename_keys.append((f"""backbone.stages.{i}.{j}.gamma""", f"""backbone.encoder.stages.{i}.layers.{j}.layer_scale_parameter""")     )
									rename_keys.append((f"""backbone.stages.{i}.{j}.depthwise_conv.weight""", f"""backbone.encoder.stages.{i}.layers.{j}.dwconv.weight""")     )
									rename_keys.append((f"""backbone.stages.{i}.{j}.depthwise_conv.bias""", f"""backbone.encoder.stages.{i}.layers.{j}.dwconv.bias""")     )
									rename_keys.append((f"""backbone.stages.{i}.{j}.norm.weight""", f"""backbone.encoder.stages.{i}.layers.{j}.layernorm.weight""")     )
									rename_keys.append((f"""backbone.stages.{i}.{j}.norm.bias""", f"""backbone.encoder.stages.{i}.layers.{j}.layernorm.bias""")     )
									rename_keys.append((f"""backbone.stages.{i}.{j}.pointwise_conv1.weight""", f"""backbone.encoder.stages.{i}.layers.{j}.pwconv1.weight""")     )
									rename_keys.append((f"""backbone.stages.{i}.{j}.pointwise_conv1.bias""", f"""backbone.encoder.stages.{i}.layers.{j}.pwconv1.bias""")     )
									rename_keys.append((f"""backbone.stages.{i}.{j}.pointwise_conv2.weight""", f"""backbone.encoder.stages.{i}.layers.{j}.pwconv2.weight""")     )
									rename_keys.append((f"""backbone.stages.{i}.{j}.pointwise_conv2.bias""", f"""backbone.encoder.stages.{i}.layers.{j}.pwconv2.bias""")     )
						if i > 0:
									rename_keys.append((f"""backbone.downsample_layers.{i}.0.weight""", f"""backbone.encoder.stages.{i}.downsampling_layer.0.weight""")     )
									rename_keys.append((f"""backbone.downsample_layers.{i}.0.bias""", f"""backbone.encoder.stages.{i}.downsampling_layer.0.bias""")     )
									rename_keys.append((f"""backbone.downsample_layers.{i}.1.weight""", f"""backbone.encoder.stages.{i}.downsampling_layer.1.weight""")     )
									rename_keys.append((f"""backbone.downsample_layers.{i}.1.bias""", f"""backbone.encoder.stages.{i}.downsampling_layer.1.bias""")     )

						rename_keys.append((f"""backbone.norm{i}.weight""", f"""backbone.hidden_states_norms.stage{i+1}.weight""")     )
						rename_keys.append((f"""backbone.norm{i}.bias""", f"""backbone.hidden_states_norms.stage{i+1}.bias""")     )

			# decode head
			rename_keys.extend(
			    [
			        ("""decode_head.conv_seg.weight""", """decode_head.classifier.weight"""),
			        ("""decode_head.conv_seg.bias""", """decode_head.classifier.bias"""),
			        ("""auxiliary_head.conv_seg.weight""", """auxiliary_head.classifier.weight"""),
			        ("""auxiliary_head.conv_seg.bias""", """auxiliary_head.classifier.bias"""),
			    ]     )
			# fmt: on

			return rename_keys



def     _lowerCAmelCase		(  _lowerCAmelCase  ,_lowerCAmelCase  ,_lowerCAmelCase     ):
			'''simple docstring'''


			A_     :				str			      =	dct.pop(_lowerCAmelCase     )
			A_     :				Optional[int]			      =	val



def     _lowerCAmelCase		(  _lowerCAmelCase  ,_lowerCAmelCase  ,_lowerCAmelCase     ):
			'''simple docstring'''


			A_     :				int			      =	{
			    """upernet-convnext-tiny""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_tiny_fp16_512x512_160k_ade20k/upernet_convnext_tiny_fp16_512x512_160k_ade20k_20220227_124553-cad485de.pth""",
			    """upernet-convnext-small""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_small_fp16_512x512_160k_ade20k/upernet_convnext_small_fp16_512x512_160k_ade20k_20220227_131208-1b1e394f.pth""",
			    """upernet-convnext-base""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_base_fp16_512x512_160k_ade20k/upernet_convnext_base_fp16_512x512_160k_ade20k_20220227_181227-02a24fc6.pth""",
			    """upernet-convnext-large""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_large_fp16_640x640_160k_ade20k/upernet_convnext_large_fp16_640x640_160k_ade20k_20220226_040532-e57aa54d.pth""",
			    """upernet-convnext-xlarge""": """https://download.openmmlab.com/mmsegmentation/v0.5/convnext/upernet_convnext_xlarge_fp16_640x640_160k_ade20k/upernet_convnext_xlarge_fp16_640x640_160k_ade20k_20220226_080344-95fc38c2.pth""",
			}
			A_     :				str			      =	model_name_to_url[model_name]
			A_     :				str			      =	torch.hub.load_state_dict_from_url(_lowerCAmelCase  ,map_location="""cpu"""     )["""state_dict"""]

			A_     :				Any			      =	get_upernet_config(_lowerCAmelCase     )
			A_     :				List[str]			      =	UperNetForSemanticSegmentation(_lowerCAmelCase     )
			model.eval()

			# replace "bn" => "batch_norm"
			for key in state_dict.copy().keys():
						A_     :				List[Any]			      =	state_dict.pop(_lowerCAmelCase     )
						if "bn" in key:
									A_     :				Any			      =	key.replace("""bn"""  ,"""batch_norm"""     )
						A_     :				List[Any]			      =	val

			# rename keys
			A_     :				int			      =	create_rename_keys(_lowerCAmelCase     )
			for src, dest in rename_keys:
						rename_key(_lowerCAmelCase  ,_lowerCAmelCase  ,_lowerCAmelCase     )

			model.load_state_dict(_lowerCAmelCase     )

			# verify on image
			A_     :				List[Any]			      =	"""https://huggingface.co/datasets/hf-internal-testing/fixtures_ade20k/resolve/main/ADE_val_00000001.jpg"""
			A_     :				Union[str, Any]			      =	Image.open(requests.get(_lowerCAmelCase  ,stream=_lowerCAmelCase     ).raw     ).convert("""RGB"""     )

			A_     :				Optional[int]			      =	SegformerImageProcessor()
			A_     :				Union[str, Any]			      =	processor(_lowerCAmelCase  ,return_tensors="""pt"""     ).pixel_values

			with torch.no_grad():
						A_     :				List[Any]			      =	model(_lowerCAmelCase     )

			if model_name == "upernet-convnext-tiny":
						A_     :				int			      =	torch.tensor(
						    [[-8.8110, -8.8110, -8.6521], [-8.8110, -8.8110, -8.6521], [-8.7746, -8.7746, -8.6130]]     )
			elif model_name == "upernet-convnext-small":
						A_     :				Dict			      =	torch.tensor(
						    [[-8.8236, -8.8236, -8.6771], [-8.8236, -8.8236, -8.6771], [-8.7638, -8.7638, -8.6240]]     )
			elif model_name == "upernet-convnext-base":
						A_     :				str			      =	torch.tensor(
						    [[-8.8558, -8.8558, -8.6905], [-8.8558, -8.8558, -8.6905], [-8.7669, -8.7669, -8.6021]]     )
			elif model_name == "upernet-convnext-large":
						A_     :				str			      =	torch.tensor(
						    [[-8.6660, -8.6660, -8.6210], [-8.6660, -8.6660, -8.6210], [-8.6310, -8.6310, -8.5964]]     )
			elif model_name == "upernet-convnext-xlarge":
						A_     :				str			      =	torch.tensor(
						    [[-8.4980, -8.4980, -8.3977], [-8.4980, -8.4980, -8.3977], [-8.4379, -8.4379, -8.3412]]     )
			print("""Logits:"""  ,outputs.logits[0, 0, :3, :3]     )
			assert torch.allclose(outputs.logits[0, 0, :3, :3]  ,_lowerCAmelCase  ,atol=1e-4     )
			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(_lowerCAmelCase     )
						print(f"""Saving processor to {pytorch_dump_folder_path}"""     )
						processor.save_pretrained(_lowerCAmelCase     )

			if push_to_hub:
						print(f"""Pushing model and processor for {model_name} to hub"""     )
						model.push_to_hub(f"""openmmlab/{model_name}"""     )
						processor.push_to_hub(f"""openmmlab/{model_name}"""     )


if __name__ == "__main__":
						_lowerCAmelCase 				=     argparse.ArgumentParser()
						# Required parameters
						parser.add_argument(
						    """--model_name""",
						    default="""upernet-convnext-tiny""",
						    type=str,
						    choices=[f'''upernet-convnext-{size}''' for size in ["""tiny""", """small""", """base""", """large""", """xlarge"""]],
						    help="""Name of the ConvNext UperNet model you'd like to convert.""",
						)
						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."""
						)

						_lowerCAmelCase 				=     parser.parse_args()
						convert_upernet_checkpoint(args.model_name, args.pytorch_dump_folder_path, args.push_to_hub)













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, 512, 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		(  _lowerCAmelCase  ,_lowerCAmelCase = 0.0  ,_lowerCAmelCase = False     ):
   '''simple docstring'''


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

class    _UpperCAmelCase							(		nn.Module     ):
      def __init__(		self      ,       a__ = None       ):
         super().__init__()
         A_     :				int			      =	drop_prob
      def 		_lowerCamelCase      (		self      ,       a__       ):
         return drop_path(a__      ,       self.drop_prob      ,       self.training       )






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






class    _UpperCAmelCase							(		nn.Module     ):
      def __init__(		self      ,       a__      ,       a__      ,       a__      ,       a__      ,       a__      ,       a__=None       ):
         super().__init__()
         A_     :				Dict			      =	patch_size if isinstance(a__      ,       collections.abc.Iterable       ) else (patch_size, patch_size)
         A_     :				Optional[Any]			      =	stride if isinstance(a__      ,       collections.abc.Iterable       ) else (stride, stride)
         A_     :				int			      =	padding if isinstance(a__      ,       collections.abc.Iterable       ) else (padding, padding)

         A_     :				Dict			      =	nn.Convad(a__      ,       a__      ,       kernel_size=a__      ,       stride=a__      ,       padding=a__       )
         A_     :				Optional[Any]			      =	norm_layer(a__       ) if norm_layer else nn.Identity()






      def 		_lowerCamelCase      (		self      ,       a__       ):
         A_     :				Dict			      =	self.projection(a__       )
         A_     :				List[Any]			      =	self.norm(a__       )
         return embeddings






class    _UpperCAmelCase							(		nn.GroupNorm     ):
      def __init__(		self      ,       a__      ,       **a__       ):
         super().__init__(1      ,       a__      ,       **a__       )






class    _UpperCAmelCase							(		nn.Module     ):
      def __init__(		self      ,       a__       ):
         super().__init__()
         A_     :				Optional[int]			      =	nn.AvgPoolad(a__      ,       stride=1      ,       padding=pool_size // 2      ,       count_include_pad=a__       )






      def 		_lowerCamelCase      (		self      ,       a__       ):
         return self.pool(a__       ) - hidden_states






class    _UpperCAmelCase							(		nn.Module     ):
      def __init__(		self      ,       a__      ,       a__      ,       a__      ,       a__       ):
         super().__init__()
         A_     :				str			      =	nn.Convad(a__      ,       a__      ,       1       )
         A_     :				List[str]			      =	nn.Convad(a__      ,       a__      ,       1       )
         A_     :				Optional[Any]			      =	PoolFormerDropPath(a__       )
         if isinstance(config.hidden_act      ,       a__       ):
            A_     :				int			      =	ACTaFN[config.hidden_act]
         else:
            A_     :				Optional[int]			      =	config.hidden_act






      def 		_lowerCamelCase      (		self      ,       a__       ):
         A_     :				Tuple			      =	self.conva(a__       )
         A_     :				str			      =	self.act_fn(a__       )
         A_     :				Union[str, Any]			      =	self.drop(a__       )
         A_     :				int			      =	self.conva(a__       )
         A_     :				Any			      =	self.drop(a__       )

         return hidden_states






class    _UpperCAmelCase							(		nn.Module     ):
      def __init__(		self      ,       a__      ,       a__      ,       a__      ,       a__      ,       a__      ,       a__       ):
         super().__init__()
         A_     :				Dict			      =	PoolFormerPooling(a__       )
         A_     :				Dict			      =	PoolFormerOutput(a__      ,       a__      ,       a__      ,       a__       )
         A_     :				Dict			      =	PoolFormerGroupNorm(a__       )
         A_     :				Dict			      =	PoolFormerGroupNorm(a__       )

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






      def 		_lowerCamelCase      (		self      ,       a__       ):
         if self.use_layer_scale:
            A_     :				List[str]			      =	self.pooling(self.before_norm(a__       )       )
            A_     :				Optional[Any]			      =	self.layer_scale_a.unsqueeze(-1       ).unsqueeze(-1       ) * pooling_output
            # First residual connection
            A_     :				List[Any]			      =	hidden_states + self.drop_path(a__       )
            A_     :				Tuple			      =	()

            A_     :				List[Any]			      =	self.output(self.after_norm(a__       )       )
            A_     :				str			      =	self.layer_scale_a.unsqueeze(-1       ).unsqueeze(-1       ) * layer_output
            # Second residual connection
            A_     :				Optional[Any]			      =	hidden_states + self.drop_path(a__       )

            A_     :				int			      =	(output,) + outputs
            return outputs

         else:
            A_     :				Union[str, Any]			      =	self.drop_path(self.pooling(self.before_norm(a__       )       )       )
            # First residual connection
            A_     :				Any			      =	pooling_output + hidden_states
            A_     :				Tuple			      =	()

            # Second residual connection inside the PoolFormerOutput block
            A_     :				Dict			      =	self.drop_path(self.output(self.after_norm(a__       )       )       )
            A_     :				int			      =	hidden_states + layer_output

            A_     :				str			      =	(output,) + outputs
            return outputs






class    _UpperCAmelCase							(		nn.Module     ):
      def __init__(		self      ,       a__       ):
         super().__init__()
         A_     :				Union[str, Any]			      =	config
         # stochastic depth decay rule
         A_     :				Union[str, Any]			      =	[x.item() for x in torch.linspace(0      ,       config.drop_path_rate      ,       sum(config.depths       )       )]

         # patch embeddings
         A_     :				Union[str, 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]      ,       )       )
         A_     :				Optional[int]			      =	nn.ModuleList(a__       )

         # Transformer blocks
         A_     :				Union[str, Any]			      =	[]
         A_     :				Union[str, Any]			      =	0
         for i in range(config.num_encoder_blocks       ):
            # each block consists of layers
            A_     :				Dict			      =	[]
            if i != 0:
               cur += config.depths[i - 1]
            for j in range(config.depths[i]       ):
               layers.append(
                   PoolFormerLayer(
                       a__      ,       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(a__       )       )

         A_     :				str			      =	nn.ModuleList(a__       )






      def 		_lowerCamelCase      (		self      ,       a__      ,       a__=False      ,       a__=True       ):
         A_     :				Optional[int]			      =	() if output_hidden_states else None

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

            if output_hidden_states:
               A_     :				Dict			      =	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=a__      ,       hidden_states=a__       )






class    _UpperCAmelCase							(		_lowerCamelCase     ):
      a						       =							PoolFormerConfig
      a						       =							'''poolformer'''
      a						       =							'''pixel_values'''
      a						       =							True
      def 		_lowerCamelCase      (		self      ,       a__       ):
         if isinstance(a__      ,       (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(a__      ,       nn.LayerNorm       ):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0       )






      def 		_lowerCamelCase      (		self      ,       a__      ,       a__=False       ):
         if isinstance(a__      ,       a__       ):
            A_     :				List[Any]			      =	value






_lowerCAmelCase 				=     R"""
    This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use
    it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and
    behavior.

    Parameters:
        config ([`PoolFormerConfig`]): Model configuration class with all the parameters of the model.
            Initializing with a config file does not load the weights associated with the model, only the
            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
"""

_lowerCAmelCase 				=     R"""
    Args:
        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See
            [`PoolFormerImageProcessor.__call__`] for details.
"""

@add_start_docstrings(
    '''The bare PoolFormer Model transformer outputting raw hidden-states without any specific head on top.'''			,       _lowerCamelCase			,       )
class    _UpperCAmelCase							(		_lowerCamelCase     ):
      def __init__(		self      ,       a__       ):
         super().__init__(a__       )
         A_     :				List[Any]			      =	config

         A_     :				Optional[int]			      =	PoolFormerEncoder(a__       )

         # Initialize weights and apply final processing
         self.post_init()
      def 		_lowerCamelCase      (		self       ):
         return self.embeddings.patch_embeddings






      @add_start_docstrings_to_model_forward(a__       )
      @add_code_sample_docstrings(
          checkpoint=_CHECKPOINT_FOR_DOC      ,       output_type=a__      ,       config_class=_CONFIG_FOR_DOC      ,       modality="""vision"""      ,       expected_output=_EXPECTED_OUTPUT_SHAPE      ,       )
      def 		_lowerCamelCase      (		self      ,       a__ = None      ,       a__ = None      ,       a__ = None      ,       ):
         A_     :				Union[str, Any]			      =	(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         A_     :				str			      =	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_     :				Optional[Any]			      =	self.encoder(
             a__      ,       output_hidden_states=a__      ,       return_dict=a__      ,       )
         A_     :				Optional[int]			      =	encoder_outputs[0]

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

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






class    _UpperCAmelCase							(		nn.Module     ):
      def __init__(		self      ,       a__       ):
         super().__init__()
         A_     :				int			      =	nn.Linear(config.hidden_size      ,       config.hidden_size       )






      def 		_lowerCamelCase      (		self      ,       a__       ):
         A_     :				Any			      =	self.dense(a__       )
         return output






@add_start_docstrings(
    '''
    PoolFormer Model transformer with an image classification head on top
    '''			,       _lowerCamelCase			,       )
class    _UpperCAmelCase							(		_lowerCamelCase     ):
      def __init__(		self      ,       a__       ):
         super().__init__(a__       )
         A_     :				str			      =	config.num_labels
         A_     :				Union[str, Any]			      =	PoolFormerModel(a__       )

         # Final norm
         A_     :				str			      =	PoolFormerGroupNorm(config.hidden_sizes[-1]       )
         # Classifier head
         A_     :				Optional[int]			      =	(
             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(a__       )
      @add_code_sample_docstrings(
          checkpoint=_IMAGE_CLASS_CHECKPOINT      ,       output_type=a__      ,       config_class=_CONFIG_FOR_DOC      ,       expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT      ,       )
      def 		_lowerCamelCase      (		self      ,       a__ = None      ,       a__ = None      ,       a__ = None      ,       a__ = None      ,       ):
         A_     :				List[Any]			      =	return_dict if return_dict is not None else self.config.use_return_dict

         A_     :				Optional[int]			      =	self.poolformer(
             a__      ,       output_hidden_states=a__      ,       return_dict=a__      ,       )

         A_     :				List[str]			      =	outputs[0]

         A_     :				int			      =	self.classifier(self.norm(a__       ).mean([-2, -1]       )       )

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

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

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

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






