






import argparse
import logging
import os
from datetime import datetime

import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader, RandomSampler, TensorDataset
from tqdm import tqdm

from transformers import GPTaLMHeadModel


__lowerCamelCase   :     Optional[Any] 		= logging.getLogger(__name__)







def 	SCREAMING_SNAKE_CASE							(					snake_case_   :				Dict	,    snake_case_   :				int  ):
	# save results
	if os.path.exists(__UpperCamelCase  ):
		if os.path.exists(os.path.join(__UpperCamelCase	,    "config.json"  )  ) and os.path.isfile(
		    os.path.join(__UpperCamelCase	,    "config.json"  )  ):
			os.remove(os.path.join(__UpperCamelCase	,    "config.json"  )  )
		if os.path.exists(os.path.join(__UpperCamelCase	,    "pytorch_model.bin"  )  ) and os.path.isfile(
		    os.path.join(__UpperCamelCase	,    "pytorch_model.bin"  )  ):
			os.remove(os.path.join(__UpperCamelCase	,    "pytorch_model.bin"  )  )
	else:
		os.makedirs(__UpperCamelCase  )
	model.save_pretrained(__UpperCamelCase  )







def 	SCREAMING_SNAKE_CASE							(					snake_case_   :				List[str]	,    snake_case_   :				Tuple=False  ):
	snake_case__					: Dict  =    2
	if unlogit:
		snake_case__					: Tuple  =    torch.pow(__UpperCamelCase	,    __UpperCamelCase  )
	snake_case__					: str  =    p * torch.log(__UpperCamelCase  )
	snake_case__					: List[str]  =    0
	return -plogp.sum(dim=-1  )







def 	SCREAMING_SNAKE_CASE							(					snake_case_   :				int  ):
	logger.info("lv, h >\t" + "\t".join(F'''{x + 1}''' for x in range(len(__UpperCamelCase  )  )  )  )
	for row in range(len(__UpperCamelCase  )  ):
		if tensor.dtype != torch.long:
			logger.info(F'''layer {row + 1}:\t''' + "\t".join(F'''{x:.5f}''' for x in tensor[row].cpu().data  )  )
		else:
			logger.info(F'''layer {row + 1}:\t''' + "\t".join(F'''{x:d}''' for x in tensor[row].cpu().data  )  )







def 	SCREAMING_SNAKE_CASE							(					snake_case_   :				str	,    snake_case_   :				List[Any]	,    snake_case_   :				Tuple	,    snake_case_   :				Optional[Any]=True	,    snake_case_   :				Dict=True	,    snake_case_   :				Tuple=None	,    snake_case_   :				Optional[int]=False  ):
	snake_case__					: Any  =    model.config.num_hidden_layers, model.config.num_attention_heads
	snake_case__					: int  =    torch.zeros(__UpperCamelCase	,    __UpperCamelCase  ).to(args.device  )
	snake_case__					: Dict  =    torch.zeros(__UpperCamelCase	,    __UpperCamelCase  ).to(args.device  )

	if head_mask is None:
		snake_case__					: Any  =    torch.ones(__UpperCamelCase	,    __UpperCamelCase  ).to(args.device  )

	head_mask.requires_grad_(requires_grad=__UpperCamelCase  )
	# If actually pruned attention multi-head, set head mask to None to avoid shape mismatch
	if actually_pruned:
		snake_case__					: List[str]  =    None

	snake_case__					: Optional[int]  =    0.0
	snake_case__					: Dict  =    0.0
	for step, inputs in enumerate(tqdm(__UpperCamelCase	,    desc="Iteration"	,    disable=args.local_rank not in [-1, 0]  )  ):
		snake_case__					: Tuple  =    tuple(t.to(args.device  ) for t in inputs  )
		(snake_case__   )					: Tuple  =    inputs

		# Do a forward pass (not with torch.no_grad() since we need gradients for importance score - see below)
		snake_case__					: Any  =    model(__UpperCamelCase	,    labels=__UpperCamelCase	,    head_mask=__UpperCamelCase  )
		#  (loss), lm_logits, presents, (all hidden_states), (attentions)
		snake_case__					: Dict  =    (
		    outputs[0],
		    outputs[1],
		    outputs[-1],
		)  # Loss and logits are the first, attention the last
		loss.backward()  # Backpropagate to populate the gradients in the head mask
		total_loss += loss.detach().cpu().numpy()
		if compute_entropy:
			for layer, attn in enumerate(__UpperCamelCase  ):
				snake_case__					: Tuple  =    entropy(attn.detach()	,    __UpperCamelCase  )
				attn_entropy[layer] += masked_entropy.sum(-1  ).sum(0  ).sum(0  ).detach()

		if compute_importance:
			head_importance += head_mask.grad.abs().detach()
		tot_tokens += torch.ones_like(__UpperCamelCase  ).float().detach().sum().data

	# Normalize
	attn_entropy /= tot_tokens
	head_importance /= tot_tokens
	# Layerwise importance normalization
	if not args.dont_normalize_importance_by_layer:
		snake_case__					: Tuple  =    2
		snake_case__					: int  =    torch.pow(torch.pow(__UpperCamelCase	,    __UpperCamelCase  ).sum(-1  )	,    1 / exponent  )
		head_importance /= norm_by_layer.unsqueeze(-1  ) + 1E-20

	if not args.dont_normalize_global_importance:
		snake_case__					: Tuple  =    (head_importance - head_importance.min()) / (head_importance.max() - head_importance.min())

	# Print matrices
	if compute_entropy:
		logger.info("Attention entropies"  )
		print_ad_tensor(__UpperCamelCase  )
	if compute_importance:
		logger.info("Head importance scores"  )
		print_ad_tensor(__UpperCamelCase  )
	logger.info("Head ranked by importance scores"  )
	snake_case__					: List[Any]  =    torch.zeros(head_importance.numel()	,    dtype=torch.long	,    device=args.device  )
	snake_case__					: Optional[Any]  =    torch.arange(
	    head_importance.numel()	,    device=args.device  )
	snake_case__					: str  =    head_ranks.view_as(__UpperCamelCase  )
	print_ad_tensor(__UpperCamelCase  )
	return attn_entropy, head_importance, total_loss







def 	SCREAMING_SNAKE_CASE							(					snake_case_   :				Optional[Any]	,    snake_case_   :				Optional[int]	,    snake_case_   :				List[Any]  ):
	snake_case__					: List[Any]  =    compute_heads_importance(__UpperCamelCase	,    __UpperCamelCase	,    __UpperCamelCase	,    compute_entropy=__UpperCamelCase  )
	snake_case__					: Optional[int]  =    1 / loss  # instead of downsteam score use the LM loss
	logger.info("Pruning: original score: %f, threshold: %f"	,    __UpperCamelCase	,    original_score * args.masking_threshold  )

	snake_case__					: int  =    torch.ones_like(__UpperCamelCase  )
	snake_case__					: Dict  =    max(1	,    int(new_head_mask.numel() * args.masking_amount  )  )

	snake_case__					: int  =    original_score
	while current_score >= original_score * args.masking_threshold:
		snake_case__					: Optional[Any]  =    new_head_mask.clone().detach()  # save current head mask
		# heads from least important to most - keep only not-masked heads
		snake_case__					: Any  =    float("Inf"  )
		snake_case__					: Tuple  =    head_importance.view(-1  ).sort()[1]

		if len(__UpperCamelCase  ) <= num_to_mask:
			print("BREAK BY num_to_mask"  )
			break

		# mask heads
		snake_case__					: int  =    current_heads_to_mask[:num_to_mask]
		logger.info("Heads to mask: %s"	,    str(current_heads_to_mask.tolist()  )  )
		snake_case__					: Optional[Any]  =    new_head_mask.view(-1  )
		snake_case__					: int  =    0.0
		snake_case__					: List[Any]  =    new_head_mask.view_as(__UpperCamelCase  )
		snake_case__					: Optional[Any]  =    new_head_mask.clone().detach()
		print_ad_tensor(__UpperCamelCase  )

		# Compute metric and head importance again
		snake_case__					: str  =    compute_heads_importance(
		    __UpperCamelCase	,    __UpperCamelCase	,    __UpperCamelCase	,    compute_entropy=__UpperCamelCase	,    head_mask=__UpperCamelCase  )
		snake_case__					: int  =    1 / loss
		logger.info(
		    "Masking: current score: %f, remaining heads %d (%.1f percents)"	,    __UpperCamelCase	,    new_head_mask.sum()	,    new_head_mask.sum() / new_head_mask.numel() * 100	,    )

	logger.info("Final head mask"  )
	print_ad_tensor(__UpperCamelCase  )
	np.save(os.path.join(args.output_dir	,    "head_mask.npy"  )	,    head_mask.detach().cpu().numpy()  )

	return head_mask







def 	SCREAMING_SNAKE_CASE							(					snake_case_   :				Optional[Any]	,    snake_case_   :				Union[str, Any]	,    snake_case_   :				Dict	,    snake_case_   :				Optional[int]  ):
	snake_case__					: Tuple  =    datetime.now()
	snake_case__					: List[str]  =    compute_heads_importance(
	    __UpperCamelCase	,    __UpperCamelCase	,    __UpperCamelCase	,    compute_entropy=__UpperCamelCase	,    compute_importance=__UpperCamelCase	,    head_mask=__UpperCamelCase  )
	snake_case__					: int  =    1 / loss
	snake_case__					: Dict  =    datetime.now() - before_time

	snake_case__					: List[str]  =    sum(p.numel() for p in model.parameters()  )
	snake_case__					: Dict  =    {
	    layer: (1 - head_mask[layer].long()).nonzero().squeeze().tolist() for layer in range(len(__UpperCamelCase  )  )
	}

	for k, v in heads_to_prune.items():
		if isinstance(__UpperCamelCase	,    __UpperCamelCase  ):
			snake_case__					: List[Any]  =    [
			    v,
			]

	assert sum(len(__UpperCamelCase  ) for h in heads_to_prune.values()  ) == (1 - head_mask.long()).sum().item()
	model.prune_heads(__UpperCamelCase  )
	snake_case__					: Any  =    sum(p.numel() for p in model.parameters()  )

	snake_case__					: List[Any]  =    datetime.now()
	snake_case__					: Optional[int]  =    compute_heads_importance(
	    __UpperCamelCase	,    __UpperCamelCase	,    __UpperCamelCase	,    compute_entropy=__UpperCamelCase	,    compute_importance=__UpperCamelCase	,    head_mask=__UpperCamelCase	,    actually_pruned=__UpperCamelCase	,    )

	snake_case__					: List[str]  =    1 / loss
	snake_case__					: Union[str, Any]  =    datetime.now() - before_time

	logger.info(
	    "Pruning: original num of params: %.2e, after pruning %.2e (%.1f percents)"	,    __UpperCamelCase	,    __UpperCamelCase	,    pruned_num_params / original_num_params * 100	,    )
	logger.info("Pruning: score with masking: %f score with pruning: %f"	,    __UpperCamelCase	,    __UpperCamelCase  )
	logger.info("Pruning: speed ratio (original timing / new timing): %f percents"	,    original_time / new_time * 100  )
	save_model(__UpperCamelCase	,    args.output_dir  )







def 	SCREAMING_SNAKE_CASE							(					):
	snake_case__					: Optional[Any]  =    argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    "--data_dir"	,    default=__UpperCamelCase	,    type=__UpperCamelCase	,    required=__UpperCamelCase	,    help="The input data dir. Should contain the .tsv files (or other data files) for the task."	,    )
	parser.add_argument(
	    "--model_name_or_path"	,    default=__UpperCamelCase	,    type=__UpperCamelCase	,    required=__UpperCamelCase	,    help="Path to pretrained model or model identifier from huggingface.co/models"	,    )
	parser.add_argument(
	    "--output_dir"	,    default=__UpperCamelCase	,    type=__UpperCamelCase	,    required=__UpperCamelCase	,    help="The output directory where the model predictions and checkpoints will be written."	,    )

	# Other parameters
	parser.add_argument(
	    "--config_name"	,    default=""	,    type=__UpperCamelCase	,    help="Pretrained config name or path if not the same as model_name_or_path"	,    )
	parser.add_argument(
	    "--tokenizer_name"	,    default=""	,    type=__UpperCamelCase	,    help="Pretrained tokenizer name or path if not the same as model_name_or_path"	,    )
	parser.add_argument(
	    "--cache_dir"	,    default=__UpperCamelCase	,    type=__UpperCamelCase	,    help="Where do you want to store the pre-trained models downloaded from s3"	,    )
	parser.add_argument(
	    "--data_subset"	,    type=__UpperCamelCase	,    default=-1	,    help="If > 0: limit the data to a subset of data_subset instances."  )
	parser.add_argument(
	    "--overwrite_output_dir"	,    action="store_true"	,    help="Whether to overwrite data in output directory"  )
	parser.add_argument(
	    "--overwrite_cache"	,    action="store_true"	,    help="Overwrite the cached training and evaluation sets"  )

	parser.add_argument(
	    "--dont_normalize_importance_by_layer"	,    action="store_true"	,    help="Don't normalize importance score by layers"  )
	parser.add_argument(
	    "--dont_normalize_global_importance"	,    action="store_true"	,    help="Don't normalize all importance scores between 0 and 1"	,    )

	parser.add_argument(
	    "--try_masking"	,    action="store_true"	,    help="Whether to try to mask head until a threshold of accuracy."  )
	parser.add_argument(
	    "--masking_threshold"	,    default=0.9	,    type=__UpperCamelCase	,    help="masking threshold in term of metrics (stop masking when metric < threshold * original metric value)."	,    )
	parser.add_argument(
	    "--masking_amount"	,    default=0.1	,    type=__UpperCamelCase	,    help="Amount to heads to masking at each masking step."  )
	parser.add_argument("--metric_name"	,    default="acc"	,    type=__UpperCamelCase	,    help="Metric to use for head masking."  )

	parser.add_argument(
	    "--max_seq_length"	,    default=128	,    type=__UpperCamelCase	,    help=(
	        "The maximum total input sequence length after WordPiece tokenization. \n"
	        "Sequences longer than this will be truncated, sequences shorter padded."
	    )	,    )
	parser.add_argument("--batch_size"	,    default=1	,    type=__UpperCamelCase	,    help="Batch size."  )

	parser.add_argument("--seed"	,    type=__UpperCamelCase	,    default=42  )
	parser.add_argument("--local_rank"	,    type=__UpperCamelCase	,    default=-1	,    help="local_rank for distributed training on gpus"  )
	parser.add_argument("--no_cuda"	,    action="store_true"	,    help="Whether not to use CUDA when available"  )
	parser.add_argument("--server_ip"	,    type=__UpperCamelCase	,    default=""	,    help="Can be used for distant debugging."  )
	parser.add_argument("--server_port"	,    type=__UpperCamelCase	,    default=""	,    help="Can be used for distant debugging."  )
	snake_case__					: Dict  =    parser.parse_args()

	if args.server_ip and args.server_port:
		# Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
		import ptvsd

		print("Waiting for debugger attach"  )
		ptvsd.enable_attach(address=(args.server_ip, args.server_port)	,    redirect_output=__UpperCamelCase  )
		ptvsd.wait_for_attach()

	# Setup devices and distributed training
	if args.local_rank == -1 or args.no_cuda:
		snake_case__					: Optional[int]  =    torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu"  )
		snake_case__					: Optional[int]  =    0 if args.no_cuda else torch.cuda.device_count()
	else:
		torch.cuda.set_device(args.local_rank  )
		snake_case__					: Union[str, Any]  =    torch.device("cuda"	,    args.local_rank  )
		snake_case__					: List[str]  =    1
		torch.distributed.init_process_group(backend="nccl"  )  # Initializes the distributed backend

	# Setup logging
	logging.basicConfig(level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN  )
	logger.info("device: {} n_gpu: {}, distributed: {}".format(args.device	,    args.n_gpu	,    bool(args.local_rank != -1  )  )  )

	snake_case__					: int  =    GPTaLMHeadModel.from_pretrained(args.model_name_or_path  )

	# Distributed and parallel training
	model.to(args.device  )
	if args.local_rank != -1:
		snake_case__					: Any  =    nn.parallel.DistributedDataParallel(
		    __UpperCamelCase	,    device_ids=[args.local_rank]	,    output_device=args.local_rank	,    find_unused_parameters=__UpperCamelCase  )
	elif args.n_gpu > 1:
		snake_case__					: Tuple  =    nn.DataParallel(__UpperCamelCase  )

	# Print/save training arguments
	os.makedirs(args.output_dir	,    exist_ok=__UpperCamelCase  )
	torch.save(__UpperCamelCase	,    os.path.join(args.output_dir	,    "run_args.bin"  )  )
	logger.info("Training/evaluation parameters %s"	,    __UpperCamelCase  )

	# Prepare dataset
	snake_case__					: Any  =    np.concatenate(
	    [
	        np.loadtxt(args.data_dir	,    dtype=np.intaa  ),
	    ]  )
	snake_case__					: str  =    (torch.from_numpy(__UpperCamelCase  ),)
	snake_case__					: Optional[Any]  =    TensorDataset(*__UpperCamelCase  )
	snake_case__					: Optional[int]  =    RandomSampler(__UpperCamelCase  )
	snake_case__					: Dict  =    DataLoader(__UpperCamelCase	,    sampler=__UpperCamelCase	,    batch_size=args.batch_size  )

	# Compute head entropy and importance score
	compute_heads_importance(__UpperCamelCase	,    __UpperCamelCase	,    __UpperCamelCase  )

	# Try head masking (set heads to zero until the score goes under a threshole)
	# and head pruning (remove masked heads and see the effect on the network)
	if args.try_masking and args.masking_threshold > 0.0 and args.masking_threshold < 1.0:
		snake_case__					: List[Any]  =    mask_heads(__UpperCamelCase	,    __UpperCamelCase	,    __UpperCamelCase  )
		prune_heads(__UpperCamelCase	,    __UpperCamelCase	,    __UpperCamelCase	,    __UpperCamelCase  )


if __name__ == "__main__":
	main()








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     SCREAMING_SNAKE_CASE(      __UpperCamelCase	,				__UpperCamelCase = 0.0	,				__UpperCamelCase = False )	->						Dict:
				if drop_prob == 0.0 or not training:
								return input
				a__	:							Tuple      						= 1 - drop_prob
				a__	:							Any      						= (input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
				a__	:							Dict      						= keep_prob + torch.rand(__UpperCamelCase	,				dtype=input.dtype	,				device=input.device )
				random_tensor.floor_()  # binarize
				a__	:							Optional[int]      						= input.div(__UpperCamelCase ) * random_tensor
				return output


class 				_a						(       nn.Module			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase = None				):







											"""simple docstring"""



											super().__init__()
											a__	:							Optional[Any]      						= 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 				_a						(       nn.Module			):





							'''simple docstring'''

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







											"""simple docstring"""



											super().__init__()
											a__	:							Optional[Any]      						= patch_size if isinstance(__UpperCAmelCase ,					collections.abc.Iterable				) else (patch_size, patch_size)
											a__	:							List[str]      						= stride if isinstance(__UpperCAmelCase ,					collections.abc.Iterable				) else (stride, stride)
											a__	:							Union[str, Any]      						= padding if isinstance(__UpperCAmelCase ,					collections.abc.Iterable				) else (padding, padding)

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




							def 		_A	(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											a__	:							Dict      						= self.projection(__UpperCAmelCase				)
											a__	:							Union[str, Any]      						= self.norm(__UpperCAmelCase				)
											return embeddings





class 				_a						(       nn.GroupNorm			):





							'''simple docstring'''

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







											"""simple docstring"""



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





class 				_a						(       nn.Module			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											super().__init__()
											a__	:							List[str]      						= 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 				_a						(       nn.Module			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase				):







											"""simple docstring"""



											super().__init__()
											a__	:							str      						= nn.Convad(__UpperCAmelCase ,					__UpperCAmelCase ,					1				)
											a__	:							int      						= nn.Convad(__UpperCAmelCase ,					__UpperCAmelCase ,					1				)
											a__	:							Any      						= PoolFormerDropPath(__UpperCAmelCase				)
											if isinstance(config.hidden_act ,					__UpperCAmelCase				):
															a__	:							List[Any]      						= ACTaFN[config.hidden_act]
											else:
															a__	:							Union[str, Any]      						= config.hidden_act




							def 		_A	(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											a__	:							Union[str, Any]      						= self.conva(__UpperCAmelCase				)
											a__	:							Union[str, Any]      						= self.act_fn(__UpperCAmelCase				)
											a__	:							Dict      						= self.drop(__UpperCAmelCase				)
											a__	:							Union[str, Any]      						= self.conva(__UpperCAmelCase				)
											a__	:							int      						= self.drop(__UpperCAmelCase				)

											return hidden_states





class 				_a						(       nn.Module			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase				):







											"""simple docstring"""



											super().__init__()
											a__	:							Any      						= PoolFormerPooling(__UpperCAmelCase				)
											a__	:							Any      						= PoolFormerOutput(__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase ,					__UpperCAmelCase				)
											a__	:							Any      						= PoolFormerGroupNorm(__UpperCAmelCase				)
											a__	:							Dict      						= PoolFormerGroupNorm(__UpperCAmelCase				)

											# Useful for training neural nets
											a__	:							List[Any]      						= PoolFormerDropPath(__UpperCAmelCase				) if drop_path > 0.0 else nn.Identity()
											a__	:							List[str]      						= config.use_layer_scale
											if config.use_layer_scale:
															a__	:							str      						= nn.Parameter(
															    config.layer_scale_init_value * torch.ones((__UpperCAmelCase)				) ,					requires_grad=__UpperCAmelCase				)
															a__	:							int      						= 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__	:							Any      						= self.pooling(self.before_norm(__UpperCAmelCase				)				)
															a__	:							int      						= self.layer_scale_a.unsqueeze(-1				).unsqueeze(-1				) * pooling_output
															# First residual connection
															a__	:							Optional[int]      						= hidden_states + self.drop_path(__UpperCAmelCase				)
															a__	:							Dict      						= ()

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

															a__	:							Optional[int]      						= (output,) + outputs
															return outputs

											else:
															a__	:							Optional[int]      						= self.drop_path(self.pooling(self.before_norm(__UpperCAmelCase				)				)				)
															# First residual connection
															a__	:							Tuple      						= pooling_output + hidden_states
															a__	:							Tuple      						= ()

															# Second residual connection inside the PoolFormerOutput block
															a__	:							Optional[int]      						= self.drop_path(self.output(self.after_norm(__UpperCAmelCase				)				)				)
															a__	:							str      						= hidden_states + layer_output

															a__	:							Optional[Any]      						= (output,) + outputs
															return outputs





class 				_a						(       nn.Module			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



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

											# patch embeddings
											a__	:							List[Any]      						= []
											for i in range(config.num_encoder_blocks				):
															embeddings.append(
															    PoolFormerEmbeddings(
															        patch_size=config.patch_sizes[i] ,					stride=config.strides[i] ,					padding=config.padding[i] ,					num_channels=config.num_channels if i == 0 else config.hidden_sizes[i - 1] ,					hidden_size=config.hidden_sizes[i] ,					)				)
											a__	:							Any      						= nn.ModuleList(__UpperCAmelCase				)

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




							def 		_A	(							self ,					__UpperCAmelCase ,					__UpperCAmelCase=False ,					__UpperCAmelCase=True				):







											"""simple docstring"""



											a__	:							int      						= () if output_hidden_states else None

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

															if output_hidden_states:
																			a__	:							Optional[Any]      						= 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						(       SCREAMING_SNAKE_CASE			):





							'''simple docstring'''

							A				:Optional[int] =       PoolFormerConfig
							A				:List[str] =       "poolformer"
							A				:Tuple =       "pixel_values"
							A				:List[str] =       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__	:							Dict      						= 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."       ,  SCREAMING_SNAKE_CASE       ,  )
class 				_a						(       SCREAMING_SNAKE_CASE			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											super().__init__(__UpperCAmelCase				)
											a__	:							Optional[Any]      						= config

											a__	:							int      						= 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__	:							int      						= (
											    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__	:							List[Any]      						= self.encoder(
											    __UpperCAmelCase ,					output_hidden_states=__UpperCAmelCase ,					return_dict=__UpperCAmelCase ,					)
											a__	:							Any      						= 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			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											super().__init__()
											a__	:							Any      						= nn.Linear(config.hidden_size ,					config.hidden_size				)




							def 		_A	(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											a__	:							Any      						= self.dense(__UpperCAmelCase				)
											return output





@add_start_docstrings(
    "\n    PoolFormer Model transformer with an image classification head on top\n    "       ,  SCREAMING_SNAKE_CASE       ,  )
class 				_a						(       SCREAMING_SNAKE_CASE			):





							'''simple docstring'''

							def __init__(							self ,					__UpperCAmelCase				):







											"""simple docstring"""



											super().__init__(__UpperCAmelCase				)
											a__	:							Optional[Any]      						= config.num_labels
											a__	:							int      						= PoolFormerModel(__UpperCAmelCase				)

											# Final norm
											a__	:							Any      						= PoolFormerGroupNorm(config.hidden_sizes[-1]				)
											# Classifier head
											a__	:							Dict      						= (
											    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__	:							str      						= return_dict if return_dict is not None else self.config.use_return_dict

											a__	:							str      						= self.poolformer(
											    __UpperCAmelCase ,					output_hidden_states=__UpperCAmelCase ,					return_dict=__UpperCAmelCase ,					)

											a__	:							Optional[Any]      						= outputs[0]

											a__	:							Union[str, Any]      						= self.classifier(self.norm(__UpperCAmelCase				).mean([-2, -1]				)				)

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

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

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

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


