
"""simple docstring"""




from __future__ import annotations

import unittest

from transformers import AutoTokenizer, MBartConfig, is_tf_available
from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
						import tensorflow as tf

						from transformers import TFAutoModelForSeqaSeqLM, TFMBartForConditionalGeneration, TFMBartModel





@require_tf
class 			snake_case  :
							UpperCAmelCase__      =   MBartConfig
							UpperCAmelCase__      =   {}
							UpperCAmelCase__      =   '''gelu'''

							def __init__(self							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_=13							,							SCREAMING_SNAKE_CASE_=7							,							SCREAMING_SNAKE_CASE_=True							,							SCREAMING_SNAKE_CASE_=False							,							SCREAMING_SNAKE_CASE_=99							,							SCREAMING_SNAKE_CASE_=32							,							SCREAMING_SNAKE_CASE_=2							,							SCREAMING_SNAKE_CASE_=4							,							SCREAMING_SNAKE_CASE_=37							,							SCREAMING_SNAKE_CASE_=0.1							,							SCREAMING_SNAKE_CASE_=0.1							,							SCREAMING_SNAKE_CASE_=20							,							SCREAMING_SNAKE_CASE_=2							,							SCREAMING_SNAKE_CASE_=1							,							SCREAMING_SNAKE_CASE_=0							,							):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			parent
									SCREAMING_SNAKE_CASE_    			=			batch_size
									SCREAMING_SNAKE_CASE_    			=			seq_length
									SCREAMING_SNAKE_CASE_    			=			is_training
									SCREAMING_SNAKE_CASE_    			=			use_labels
									SCREAMING_SNAKE_CASE_    			=			vocab_size
									SCREAMING_SNAKE_CASE_    			=			hidden_size
									SCREAMING_SNAKE_CASE_    			=			num_hidden_layers
									SCREAMING_SNAKE_CASE_    			=			num_attention_heads
									SCREAMING_SNAKE_CASE_    			=			intermediate_size
									SCREAMING_SNAKE_CASE_    			=			hidden_dropout_prob
									SCREAMING_SNAKE_CASE_    			=			attention_probs_dropout_prob
									SCREAMING_SNAKE_CASE_    			=			max_position_embeddings
									SCREAMING_SNAKE_CASE_    			=			eos_token_id
									SCREAMING_SNAKE_CASE_    			=			pad_token_id
									SCREAMING_SNAKE_CASE_    			=			bos_token_id

							def _lowercase       (self     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			ids_tensor([self.batch_size, self.seq_length - 1]							,							self.vocab_size     )
									SCREAMING_SNAKE_CASE_    			=			tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size     )							,							1     )
									SCREAMING_SNAKE_CASE_    			=			tf.concat([input_ids, eos_tensor]							,							axis=1     )

									SCREAMING_SNAKE_CASE_    			=			ids_tensor([self.batch_size, self.seq_length]							,							self.vocab_size     )

									SCREAMING_SNAKE_CASE_    			=			self.config_cls(
									    vocab_size=self.vocab_size							,							d_model=self.hidden_size							,							encoder_layers=self.num_hidden_layers							,							decoder_layers=self.num_hidden_layers							,							encoder_attention_heads=self.num_attention_heads							,							decoder_attention_heads=self.num_attention_heads							,							encoder_ffn_dim=self.intermediate_size							,							decoder_ffn_dim=self.intermediate_size							,							dropout=self.hidden_dropout_prob							,							attention_dropout=self.attention_probs_dropout_prob							,							max_position_embeddings=self.max_position_embeddings							,							eos_token_ids=[2]							,							bos_token_id=self.bos_token_id							,							pad_token_id=self.pad_token_id							,							decoder_start_token_id=self.pad_token_id							,							**self.config_updates							,							)
									SCREAMING_SNAKE_CASE_    			=			prepare_mbart_inputs_dict(SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_     )
									return config, inputs_dict







							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			TFMBartModel(config=SCREAMING_SNAKE_CASE_     ).get_decoder()
									SCREAMING_SNAKE_CASE_    			=			inputs_dict['''input_ids''']

									SCREAMING_SNAKE_CASE_    			=			input_ids[:1, :]
									SCREAMING_SNAKE_CASE_    			=			inputs_dict['''attention_mask'''][:1, :]
									SCREAMING_SNAKE_CASE_    			=			inputs_dict['''head_mask''']
									SCREAMING_SNAKE_CASE_    			=			1

									# first forward pass
									SCREAMING_SNAKE_CASE_    			=			model(SCREAMING_SNAKE_CASE_							,							attention_mask=SCREAMING_SNAKE_CASE_							,							head_mask=SCREAMING_SNAKE_CASE_							,							use_cache=SCREAMING_SNAKE_CASE_     )

									SCREAMING_SNAKE_CASE_			,SCREAMING_SNAKE_CASE_    			=			outputs.to_tuple()
									SCREAMING_SNAKE_CASE_    			=			past_key_values[1]







def 							_lowerCamelCase			(			__a,				__a,				__a,				__a=None,				__a=None,				__a=None,				__a=None,				__a=None,				):
		if attention_mask is None:
				SCREAMING_SNAKE_CASE_    			=			tf.cast(tf.math.not_equal(__a,				config.pad_token_id     ),				tf.inta     )
		if decoder_attention_mask is None:
				SCREAMING_SNAKE_CASE_    			=			tf.concat(
				    [
				        tf.ones(decoder_input_ids[:, :1].shape,				dtype=tf.inta     ),
				        tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:],				config.pad_token_id     ),				tf.inta     ),
				    ],				axis=-1,				)
		if head_mask is None:
				SCREAMING_SNAKE_CASE_    			=			tf.ones((config.encoder_layers, config.encoder_attention_heads)     )
		if decoder_head_mask is None:
				SCREAMING_SNAKE_CASE_    			=			tf.ones((config.decoder_layers, config.decoder_attention_heads)     )
		if cross_attn_head_mask is None:
				SCREAMING_SNAKE_CASE_    			=			tf.ones((config.decoder_layers, config.decoder_attention_heads)     )
		return {
		    "input_ids": input_ids,
		    "decoder_input_ids": decoder_input_ids,
		    "attention_mask": attention_mask,
		    "decoder_attention_mask": decoder_attention_mask,
		    "head_mask": head_mask,
		    "decoder_head_mask": decoder_head_mask,
		    "cross_attn_head_mask": cross_attn_head_mask,
		}





@require_tf
class 			snake_case  (				__lowercase							,      __lowercase							,      unittest.TestCase ):
							UpperCAmelCase__      =   (TFMBartForConditionalGeneration, TFMBartModel) if is_tf_available() else ()
							UpperCAmelCase__      =   (TFMBartForConditionalGeneration,) if is_tf_available() else ()
							UpperCAmelCase__      =   (
							    {
							        '''conversational''': TFMBartForConditionalGeneration,
							        '''feature-extraction''': TFMBartModel,
							        '''summarization''': TFMBartForConditionalGeneration,
							        '''text2text-generation''': TFMBartForConditionalGeneration,
							        '''translation''': TFMBartForConditionalGeneration,
							    }
							    if is_tf_available()
							    else {}
							)
							UpperCAmelCase__      =   True
							UpperCAmelCase__      =   False
							UpperCAmelCase__      =   False

							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







									if pipeline_test_casse_name != "FeatureExtractionPipelineTests":
											# Exception encountered when calling layer '...'
											return True

									return False

							def _lowercase       (self     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			TFMBartModelTester(self     )
									SCREAMING_SNAKE_CASE_    			=			ConfigTester(self							,							config_class=SCREAMING_SNAKE_CASE_     )

							def _lowercase       (self     ):

									"""simple docstring"""







									self.config_tester.run_common_tests()







							def _lowercase       (self     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			self.model_tester.prepare_config_and_inputs_for_common()
									self.model_tester.check_decoder_model_past_large_inputs(*SCREAMING_SNAKE_CASE_     )







@require_sentencepiece
@require_tokenizers
@require_tf
class 			snake_case  (				unittest.TestCase ):
							UpperCAmelCase__      =   [
							    ''' UN Chief Says There Is No Military Solution in Syria''',
							]
							UpperCAmelCase__      =   [
							    '''Şeful ONU declară că nu există o soluţie militară în Siria''',
							]
							UpperCAmelCase__      =   '''facebook/mbart-large-en-ro'''

							@cached_property
							def _lowercase       (self     ):

									"""simple docstring"""







									return AutoTokenizer.from_pretrained(self.model_name     )

							@cached_property
							def _lowercase       (self     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name     )
									return model

							def _lowercase       (self							,							**SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			self.translate_src_text(**SCREAMING_SNAKE_CASE_     )
									self.assertListEqual(self.expected_text							,							SCREAMING_SNAKE_CASE_     )

							def _lowercase       (self							,							**SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			self.tokenizer(self.src_text							,							**SCREAMING_SNAKE_CASE_							,							return_tensors='''tf'''     )
									SCREAMING_SNAKE_CASE_    			=			self.model.generate(
									    model_inputs.input_ids							,							attention_mask=model_inputs.attention_mask							,							num_beams=2     )
									SCREAMING_SNAKE_CASE_    			=			self.tokenizer.batch_decode(SCREAMING_SNAKE_CASE_							,							skip_special_tokens=SCREAMING_SNAKE_CASE_     )
									return generated_words







							@slow
							def _lowercase       (self     ):

									"""simple docstring"""







									self._assert_generated_batch_equal_expected()

"""simple docstring"""




import collections.abc
from typing import Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import nn
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from ...activations import ACTaFN
from ...modeling_outputs import BaseModelOutputWithNoAttention, ImageClassifierOutputWithNoAttention
from ...modeling_utils import PreTrainedModel
from ...utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward, logging
from .configuration_poolformer import PoolFormerConfig


lowerCAmelCase__				=					logging.get_logger(__name__)

# General docstring
lowerCAmelCase__				=					'PoolFormerConfig'

# Base docstring
lowerCAmelCase__				=					'sail/poolformer_s12'
lowerCAmelCase__				=					[1, 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			(			__a,				__a = 0.0,				__a = False     ):
		if drop_prob == 0.0 or not training:
				return input
		SCREAMING_SNAKE_CASE_    			=			1 - drop_prob
		SCREAMING_SNAKE_CASE_    			=			(input.shape[0],) + (1,) * (input.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
		SCREAMING_SNAKE_CASE_    			=			keep_prob + torch.rand(__a,				dtype=input.dtype,				device=input.device     )
		random_tensor.floor_()  # binarize
		SCREAMING_SNAKE_CASE_    			=			input.div(__a     ) * random_tensor
		return output





class 			snake_case  (				nn.Module ):

							def __init__(self							,							SCREAMING_SNAKE_CASE_ = None     ):

									"""simple docstring"""







									super().__init__()
									SCREAMING_SNAKE_CASE_    			=			drop_prob

							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







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







							def _lowercase       (self     ):

									"""simple docstring"""







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







class 			snake_case  (				nn.Module ):

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

									"""simple docstring"""







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

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







							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			self.projection(SCREAMING_SNAKE_CASE_     )
									SCREAMING_SNAKE_CASE_    			=			self.norm(SCREAMING_SNAKE_CASE_     )
									return embeddings







class 			snake_case  (				nn.GroupNorm ):

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

									"""simple docstring"""







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







class 			snake_case  (				nn.Module ):

							def __init__(self							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







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







							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







									return self.pool(SCREAMING_SNAKE_CASE_     ) - hidden_states







class 			snake_case  (				nn.Module ):

							def __init__(self							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







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







							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			self.conva(SCREAMING_SNAKE_CASE_     )
									SCREAMING_SNAKE_CASE_    			=			self.act_fn(SCREAMING_SNAKE_CASE_     )
									SCREAMING_SNAKE_CASE_    			=			self.drop(SCREAMING_SNAKE_CASE_     )
									SCREAMING_SNAKE_CASE_    			=			self.conva(SCREAMING_SNAKE_CASE_     )
									SCREAMING_SNAKE_CASE_    			=			self.drop(SCREAMING_SNAKE_CASE_     )

									return hidden_states







class 			snake_case  (				nn.Module ):

							def __init__(self							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







									super().__init__()
									SCREAMING_SNAKE_CASE_    			=			PoolFormerPooling(SCREAMING_SNAKE_CASE_     )
									SCREAMING_SNAKE_CASE_    			=			PoolFormerOutput(SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_     )
									SCREAMING_SNAKE_CASE_    			=			PoolFormerGroupNorm(SCREAMING_SNAKE_CASE_     )
									SCREAMING_SNAKE_CASE_    			=			PoolFormerGroupNorm(SCREAMING_SNAKE_CASE_     )

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







							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







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

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

											SCREAMING_SNAKE_CASE_    			=			(output,) + outputs
											return outputs

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

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

											SCREAMING_SNAKE_CASE_    			=			(output,) + outputs
											return outputs







class 			snake_case  (				nn.Module ):

							def __init__(self							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







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

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

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

									SCREAMING_SNAKE_CASE_    			=			nn.ModuleList(SCREAMING_SNAKE_CASE_     )







							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_=False							,							SCREAMING_SNAKE_CASE_=True     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			() if output_hidden_states else None

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

											if output_hidden_states:
													SCREAMING_SNAKE_CASE_    			=			all_hidden_states + (hidden_states,)

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

									return BaseModelOutputWithNoAttention(last_hidden_state=SCREAMING_SNAKE_CASE_							,							hidden_states=SCREAMING_SNAKE_CASE_     )







class 			snake_case  (				__lowercase ):
							UpperCAmelCase__      =   PoolFormerConfig
							UpperCAmelCase__      =   '''poolformer'''
							UpperCAmelCase__      =   '''pixel_values'''
							UpperCAmelCase__      =   True

							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







									if isinstance(SCREAMING_SNAKE_CASE_							,							(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(SCREAMING_SNAKE_CASE_							,							nn.LayerNorm     ):
											module.bias.data.zero_()
											module.weight.data.fill_(1.0     )







							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_=False     ):

									"""simple docstring"""







									if isinstance(SCREAMING_SNAKE_CASE_							,							SCREAMING_SNAKE_CASE_     ):
											SCREAMING_SNAKE_CASE_    			=			value







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

lowerCAmelCase__				=					r'\n    Args:\n        pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):\n            Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See\n            [`PoolFormerImageProcessor.__call__`] for details.\n'





@add_start_docstrings(
    '''The bare PoolFormer Model transformer outputting raw hidden-states without any specific head on top.'''							,      __lowercase							,      )
class 			snake_case  (				__lowercase ):

							def __init__(self							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







									super().__init__(SCREAMING_SNAKE_CASE_     )
									SCREAMING_SNAKE_CASE_    			=			config

									SCREAMING_SNAKE_CASE_    			=			PoolFormerEncoder(SCREAMING_SNAKE_CASE_     )

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

							def _lowercase       (self     ):

									"""simple docstring"""







									return self.embeddings.patch_embeddings







							@add_start_docstrings_to_model_forward(SCREAMING_SNAKE_CASE_     )
							@add_code_sample_docstrings(
							    checkpoint=_CHECKPOINT_FOR_DOC							,							output_type=SCREAMING_SNAKE_CASE_							,							config_class=_CONFIG_FOR_DOC							,							modality='''vision'''							,							expected_output=_EXPECTED_OUTPUT_SHAPE							,							)
							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_ = None							,							SCREAMING_SNAKE_CASE_ = None							,							SCREAMING_SNAKE_CASE_ = None							,							):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			(
									    output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
									)
									SCREAMING_SNAKE_CASE_    			=			return_dict if return_dict is not None else self.config.use_return_dict

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

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

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

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







class 			snake_case  (				nn.Module ):

							def __init__(self							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







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







							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







									SCREAMING_SNAKE_CASE_    			=			self.dense(SCREAMING_SNAKE_CASE_     )
									return output







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

							def __init__(self							,							SCREAMING_SNAKE_CASE_     ):

									"""simple docstring"""







									super().__init__(SCREAMING_SNAKE_CASE_     )
									SCREAMING_SNAKE_CASE_    			=			config.num_labels
									SCREAMING_SNAKE_CASE_    			=			PoolFormerModel(SCREAMING_SNAKE_CASE_     )

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

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







							@add_start_docstrings_to_model_forward(SCREAMING_SNAKE_CASE_     )
							@add_code_sample_docstrings(
							    checkpoint=_IMAGE_CLASS_CHECKPOINT							,							output_type=SCREAMING_SNAKE_CASE_							,							config_class=_CONFIG_FOR_DOC							,							expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT							,							)
							def _lowercase       (self							,							SCREAMING_SNAKE_CASE_ = None							,							SCREAMING_SNAKE_CASE_ = None							,							SCREAMING_SNAKE_CASE_ = None							,							SCREAMING_SNAKE_CASE_ = None							,							):

									"""simple docstring"""







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

									SCREAMING_SNAKE_CASE_    			=			self.poolformer(
									    SCREAMING_SNAKE_CASE_							,							output_hidden_states=SCREAMING_SNAKE_CASE_							,							return_dict=SCREAMING_SNAKE_CASE_							,							)

									SCREAMING_SNAKE_CASE_    			=			outputs[0]

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

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

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

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

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