




from __future__ import annotations

import unittest

from transformers import AutoTokenizer, PegasusConfig, 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, TFPegasusForConditionalGeneration, TFPegasusModel





@require_tf
class 			__lowerCamelCase	:




    """simple docstring"""

    snake_case__ =				PegasusConfig
    snake_case__ =				{}
    snake_case__ =				"gelu"


    def __init__(	self					:					List[Any]       ,     SCREAMING_SNAKE_CASE__					:					str       ,     SCREAMING_SNAKE_CASE__					:					Any=13       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]=7       ,     SCREAMING_SNAKE_CASE__					:					List[Any]=True       ,     SCREAMING_SNAKE_CASE__					:					List[Any]=False       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]=99       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]=32       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]=2       ,     SCREAMING_SNAKE_CASE__					:					Dict=4       ,     SCREAMING_SNAKE_CASE__					:					List[str]=37       ,     SCREAMING_SNAKE_CASE__					:					Any=0.1       ,     SCREAMING_SNAKE_CASE__					:					int=0.1       ,     SCREAMING_SNAKE_CASE__					:					Dict=40       ,     SCREAMING_SNAKE_CASE__					:					str=2       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]=1       ,     SCREAMING_SNAKE_CASE__					:					Any=0       ,     )				->	Optional[Any]:
         lowerCAmelCase__          =			parent
         lowerCAmelCase__          =			batch_size
         lowerCAmelCase__          =			seq_length
         lowerCAmelCase__          =			is_training
         lowerCAmelCase__          =			use_labels
         lowerCAmelCase__          =			vocab_size
         lowerCAmelCase__          =			hidden_size
         lowerCAmelCase__          =			num_hidden_layers
         lowerCAmelCase__          =			num_attention_heads
         lowerCAmelCase__          =			intermediate_size

         lowerCAmelCase__          =			hidden_dropout_prob
         lowerCAmelCase__          =			attention_probs_dropout_prob
         lowerCAmelCase__          =			max_position_embeddings
         lowerCAmelCase__          =			eos_token_id
         lowerCAmelCase__          =			pad_token_id
         lowerCAmelCase__          =			bos_token_id


    def        a				(	self					:					str       )				->	Any:
         lowerCAmelCase__          =			ids_tensor([self.batch_size, self.seq_length - 1]       ,     self.vocab_size       )
         lowerCAmelCase__          =			tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size       )       ,     1       )
         lowerCAmelCase__          =			tf.concat([input_ids, eos_tensor]       ,     axis=1       )

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

         lowerCAmelCase__          =			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       ,     )
         lowerCAmelCase__          =			prepare_pegasus_inputs_dict(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )
         return config, inputs_dict


    def        a				(	self					:					Optional[Any]       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]       )				->	Optional[Any]:
         lowerCAmelCase__          =			TFPegasusModel(config=SCREAMING_SNAKE_CASE__       ).get_decoder()
         lowerCAmelCase__          =			inputs_dict["input_ids"]

         lowerCAmelCase__          =			input_ids[:1, :]
         lowerCAmelCase__          =			inputs_dict["attention_mask"][:1, :]
         lowerCAmelCase__          =			inputs_dict["head_mask"]
         lowerCAmelCase__          =			1

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

         lowerCAmelCase__     ,		lowerCAmelCase__          =			outputs.to_tuple()

         # create hypothetical next token and extent to next_input_ids
         lowerCAmelCase__          =			ids_tensor((self.batch_size, 3)       ,     config.vocab_size       )
         lowerCAmelCase__          =			tf.cast(ids_tensor((self.batch_size, 3)       ,     2       )       ,     tf.inta       )

         # append to next input_ids and
         lowerCAmelCase__          =			tf.concat([input_ids, next_tokens]       ,     axis=-1       )
         lowerCAmelCase__          =			tf.concat([attention_mask, next_attn_mask]       ,     axis=-1       )

         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       ,     attention_mask=SCREAMING_SNAKE_CASE__       )[0]
         lowerCAmelCase__          =			model(SCREAMING_SNAKE_CASE__       ,     attention_mask=SCREAMING_SNAKE_CASE__       ,     past_key_values=SCREAMING_SNAKE_CASE__       )[0]

         self.parent.assertEqual(next_tokens.shape[1]       ,     output_from_past.shape[1]       )

         # select random slice
         lowerCAmelCase__          =			int(ids_tensor((1,)       ,     output_from_past.shape[-1]       )       )
         lowerCAmelCase__          =			output_from_no_past[:, -3:, random_slice_idx]
         lowerCAmelCase__          =			output_from_past[:, :, random_slice_idx]

         # test that outputs are equal for slice
         tf.debugging.assert_near(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     rtol=1e-3       )






def 							_A       (						lowerCAmelCase_					: Optional[int]       ,    lowerCAmelCase_					: Optional[Any]       ,    lowerCAmelCase_					: Tuple       ,    lowerCAmelCase_					: str=None       ,    lowerCAmelCase_					: List[str]=None       ,    lowerCAmelCase_					: Tuple=None       ,    lowerCAmelCase_					: Tuple=None       ,    lowerCAmelCase_					: List[str]=None       ,    ):
     """simple docstring"""

     if attention_mask is None:
          lowerCAmelCase__          =			tf.cast(tf.math.not_equal(lowerCAmelCase_       ,    config.pad_token_id      )       ,    tf.inta      )
     if decoder_attention_mask is None:
          lowerCAmelCase__          =			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:
          lowerCAmelCase__          =			tf.ones((config.encoder_layers, config.encoder_attention_heads)      )
     if decoder_head_mask is None:
          lowerCAmelCase__          =			tf.ones((config.decoder_layers, config.decoder_attention_heads)      )
     if cross_attn_head_mask is None:
          lowerCAmelCase__          =			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 			__lowerCamelCase	(   UpperCamelCase__     ,    UpperCamelCase__     ,    unittest.TestCase   ):




    """simple docstring"""

    snake_case__ =				(TFPegasusForConditionalGeneration, TFPegasusModel) if is_tf_available() else ()
    snake_case__ =				(TFPegasusForConditionalGeneration,) if is_tf_available() else ()
    snake_case__ =				(
        {
            "conversational": TFPegasusForConditionalGeneration,
            "feature-extraction": TFPegasusModel,
            "summarization": TFPegasusForConditionalGeneration,
            "text2text-generation": TFPegasusForConditionalGeneration,
            "translation": TFPegasusForConditionalGeneration,
        }
        if is_tf_available()
        else {}
    )
    snake_case__ =				True
    snake_case__ =				False
    snake_case__ =				False


    def        a				(	self					:					Optional[int]       )				->	List[str]:
         lowerCAmelCase__          =			TFPegasusModelTester(self       )
         lowerCAmelCase__          =			ConfigTester(self       ,     config_class=SCREAMING_SNAKE_CASE__       )


    def        a				(	self					:					Tuple       )				->	Tuple:
         self.config_tester.run_common_tests()


    def        a				(	self					:					List[Any]       )				->	Optional[Any]:
         lowerCAmelCase__          =			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 			__lowerCamelCase	(   unittest.TestCase   ):




    """simple docstring"""

    snake_case__ =				[
        " PG&E stated it scheduled the blackouts in response to forecasts for high winds amid dry conditions. The aim is to reduce the risk of wildfires. Nearly 800 thousand customers were scheduled to be affected by the shutoffs which were expected to last through at least midday tomorrow.",
        " The London trio are up for best UK act and best album, as well as getting two nominations in the best song category.\"We got told like this morning 'Oh I think you're nominated'\", said Dappy.\"And I was like 'Oh yeah, which one?' And now we've got nominated for four awards. I mean, wow!\"Bandmate Fazer added: \"We thought it's best of us to come down and mingle with everyone and say hello to the cameras. And now we find we've got four nominations.\"The band have two shots at the best song prize, getting the nod for their Tynchy Stryder collaboration Number One, and single Strong Again.Their album Uncle B will also go up against records by the likes of Beyonce and Kanye West.N-Dubz picked up the best newcomer Mobo in 2007, but female member Tulisa said they wouldn't be too disappointed if they didn't win this time around.\"At the end of the day we're grateful to be where we are in our careers.\"If it don't happen then it don't happen - live to fight another day and keep on making albums and hits for the fans.\"Dappy also revealed they could be performing live several times on the night.The group will be doing Number One and also a possible rendition of the War Child single, I Got Soul.The charity song is a  re-working of The Killers' All These Things That I've Done and is set to feature artists like Chipmunk, Ironik and Pixie Lott.This year's Mobos will be held outside of London for the first time, in Glasgow on 30 September.N-Dubz said they were looking forward to performing for their Scottish fans and boasted about their recent shows north of the border.\"We just done Edinburgh the other day,\" said Dappy.\"We smashed up an N-Dubz show over there. We done Aberdeen about three or four months ago - we smashed up that show over there! Everywhere we go we smash it up!\" ",
    ]
    snake_case__ =				[
        "California's largest electricity provider has cut power to hundreds of thousands of customers in an effort to"
        " reduce the risk of wildfires.",
        "N-Dubz have revealed they\'re \"grateful\" to have been nominated for four Mobo Awards.",
    ]  # differs slightly from pytorch, likely due to numerical differences in linear layers
    snake_case__ =				"google/pegasus-xsum"


    @cached_property
    def        a				(	self					:					Union[str, Any]       )				->	List[str]:
         return AutoTokenizer.from_pretrained(self.model_name       )


    @cached_property
    def        a				(	self					:					int       )				->	Optional[Any]:
         lowerCAmelCase__          =			TFAutoModelForSeqaSeqLM.from_pretrained(self.model_name       )
         return model


    def        a				(	self					:					int       ,     **SCREAMING_SNAKE_CASE__					:					Optional[Any]       )				->	Tuple:
         lowerCAmelCase__          =			self.translate_src_text(**SCREAMING_SNAKE_CASE__       )
         assert self.expected_text == generated_words


    def        a				(	self					:					List[Any]       ,     **SCREAMING_SNAKE_CASE__					:					str       )				->	Dict:
         lowerCAmelCase__          =			self.tokenizer(self.src_text       ,     **SCREAMING_SNAKE_CASE__       ,     padding=SCREAMING_SNAKE_CASE__       ,     return_tensors="tf"       )
         lowerCAmelCase__          =			self.model.generate(
             model_inputs.input_ids       ,     attention_mask=model_inputs.attention_mask       ,     num_beams=2       ,     use_cache=SCREAMING_SNAKE_CASE__       ,     )
         lowerCAmelCase__          =			self.tokenizer.batch_decode(generated_ids.numpy()       ,     skip_special_tokens=SCREAMING_SNAKE_CASE__       )
         return generated_words


    @slow
    def        a				(	self					:					List[Any]       )				->	Tuple:
         self._assert_generated_batch_equal_expected()






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


UpperCamelCase								=					logging.get_logger(__name__)

# General docstring
UpperCamelCase								=					'PoolFormerConfig'

# Base docstring
UpperCamelCase								=					'sail/poolformer_s12'
UpperCamelCase								=					[1, 512, 7, 7]

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

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



def 							_A       (						lowerCAmelCase_					: Tuple       ,    lowerCAmelCase_					: float = 0.0       ,    lowerCAmelCase_					: bool = False      ):
     """simple docstring"""

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





class 			__lowerCamelCase	(   nn.Module   ):




    """simple docstring"""

    def __init__(	self					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					Optional[float] = None       )				->	None:
         super().__init__()
         lowerCAmelCase__          =			drop_prob


    def        a				(	self					:					str       ,     SCREAMING_SNAKE_CASE__					:					torch.Tensor       )				->	torch.Tensor:
         return drop_path(SCREAMING_SNAKE_CASE__       ,     self.drop_prob       ,     self.training       )


    def        a				(	self					:					Optional[Any]       )				->	str:
         return "p={}".format(self.drop_prob       )






class 			__lowerCamelCase	(   nn.Module   ):




    """simple docstring"""

    def __init__(	self					:					List[str]       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]       ,     SCREAMING_SNAKE_CASE__					:					Dict       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]       ,     SCREAMING_SNAKE_CASE__					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					List[str]       ,     SCREAMING_SNAKE_CASE__					:					str=None       )				->	Optional[Any]:
         super().__init__()
         lowerCAmelCase__          =			patch_size if isinstance(SCREAMING_SNAKE_CASE__       ,     collections.abc.Iterable       ) else (patch_size, patch_size)
         lowerCAmelCase__          =			stride if isinstance(SCREAMING_SNAKE_CASE__       ,     collections.abc.Iterable       ) else (stride, stride)
         lowerCAmelCase__          =			padding if isinstance(SCREAMING_SNAKE_CASE__       ,     collections.abc.Iterable       ) else (padding, padding)

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


    def        a				(	self					:					int       ,     SCREAMING_SNAKE_CASE__					:					List[Any]       )				->	Union[str, Any]:
         lowerCAmelCase__          =			self.projection(SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			self.norm(SCREAMING_SNAKE_CASE__       )
         return embeddings






class 			__lowerCamelCase	(   nn.GroupNorm   ):




    """simple docstring"""

    def __init__(	self					:					Optional[Any]       ,     SCREAMING_SNAKE_CASE__					:					List[str]       ,     **SCREAMING_SNAKE_CASE__					:					Any       )				->	Dict:
         super().__init__(1       ,     SCREAMING_SNAKE_CASE__       ,     **SCREAMING_SNAKE_CASE__       )






class 			__lowerCamelCase	(   nn.Module   ):




    """simple docstring"""

    def __init__(	self					:					Optional[int]       ,     SCREAMING_SNAKE_CASE__					:					int       )				->	List[Any]:
         super().__init__()
         lowerCAmelCase__          =			nn.AvgPoolad(SCREAMING_SNAKE_CASE__       ,     stride=1       ,     padding=pool_size // 2       ,     count_include_pad=SCREAMING_SNAKE_CASE__       )


    def        a				(	self					:					int       ,     SCREAMING_SNAKE_CASE__					:					List[Any]       )				->	Optional[Any]:
         return self.pool(SCREAMING_SNAKE_CASE__       ) - hidden_states






class 			__lowerCamelCase	(   nn.Module   ):




    """simple docstring"""

    def __init__(	self					:					Optional[Any]       ,     SCREAMING_SNAKE_CASE__					:					List[Any]       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]       ,     SCREAMING_SNAKE_CASE__					:					str       ,     SCREAMING_SNAKE_CASE__					:					Tuple       )				->	Dict:
         super().__init__()
         lowerCAmelCase__          =			nn.Convad(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     1       )
         lowerCAmelCase__          =			nn.Convad(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     1       )
         lowerCAmelCase__          =			PoolFormerDropPath(SCREAMING_SNAKE_CASE__       )
         if isinstance(config.hidden_act       ,     SCREAMING_SNAKE_CASE__       ):
              lowerCAmelCase__          =			ACTaFN[config.hidden_act]
         else:
              lowerCAmelCase__          =			config.hidden_act


    def        a				(	self					:					Union[str, Any]       ,     SCREAMING_SNAKE_CASE__					:					Optional[Any]       )				->	List[Any]:
         lowerCAmelCase__          =			self.conva(SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			self.act_fn(SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			self.drop(SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			self.conva(SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			self.drop(SCREAMING_SNAKE_CASE__       )

         return hidden_states






class 			__lowerCamelCase	(   nn.Module   ):




    """simple docstring"""

    def __init__(	self					:					List[str]       ,     SCREAMING_SNAKE_CASE__					:					int       ,     SCREAMING_SNAKE_CASE__					:					Any       ,     SCREAMING_SNAKE_CASE__					:					List[str]       ,     SCREAMING_SNAKE_CASE__					:					int       ,     SCREAMING_SNAKE_CASE__					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					int       )				->	Optional[Any]:
         super().__init__()
         lowerCAmelCase__          =			PoolFormerPooling(SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			PoolFormerOutput(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			PoolFormerGroupNorm(SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			PoolFormerGroupNorm(SCREAMING_SNAKE_CASE__       )

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


    def        a				(	self					:					Dict       ,     SCREAMING_SNAKE_CASE__					:					str       )				->	int:
         if self.use_layer_scale:
              lowerCAmelCase__          =			self.pooling(self.before_norm(SCREAMING_SNAKE_CASE__       )       )
              lowerCAmelCase__          =			self.layer_scale_a.unsqueeze(-1       ).unsqueeze(-1       ) * pooling_output
              # First residual connection
              lowerCAmelCase__          =			hidden_states + self.drop_path(SCREAMING_SNAKE_CASE__       )
              lowerCAmelCase__          =			()

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

              lowerCAmelCase__          =			(output,) + outputs
              return outputs

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

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

              lowerCAmelCase__          =			(output,) + outputs
              return outputs






class 			__lowerCamelCase	(   nn.Module   ):




    """simple docstring"""

    def __init__(	self					:					Optional[Any]       ,     SCREAMING_SNAKE_CASE__					:					Dict       )				->	Any:
         super().__init__()
         lowerCAmelCase__          =			config
         # stochastic depth decay rule
         lowerCAmelCase__          =			[x.item() for x in torch.linspace(0       ,     config.drop_path_rate       ,     sum(config.depths       )       )]

         # patch embeddings
         lowerCAmelCase__          =			[]
         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]       ,     )       )
         lowerCAmelCase__          =			nn.ModuleList(SCREAMING_SNAKE_CASE__       )

         # Transformer blocks
         lowerCAmelCase__          =			[]
         lowerCAmelCase__          =			0
         for i in range(config.num_encoder_blocks       ):
              # each block consists of layers
              lowerCAmelCase__          =			[]
              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__       )       )

         lowerCAmelCase__          =			nn.ModuleList(SCREAMING_SNAKE_CASE__       )


    def        a				(	self					:					Optional[Any]       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]       ,     SCREAMING_SNAKE_CASE__					:					int=False       ,     SCREAMING_SNAKE_CASE__					:					List[str]=True       )				->	Dict:
         lowerCAmelCase__          =			() if output_hidden_states else None

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

              if output_hidden_states:
                   lowerCAmelCase__          =			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 			__lowerCamelCase	(   UpperCamelCase__   ):




    """simple docstring"""

    snake_case__ =				PoolFormerConfig
    snake_case__ =				"poolformer"
    snake_case__ =				"pixel_values"
    snake_case__ =				True


    def        a				(	self					:					Optional[Any]       ,     SCREAMING_SNAKE_CASE__					:					int       )				->	List[Any]:
         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        a				(	self					:					int       ,     SCREAMING_SNAKE_CASE__					:					int       ,     SCREAMING_SNAKE_CASE__					:					Union[str, Any]=False       )				->	Tuple:
         if isinstance(SCREAMING_SNAKE_CASE__       ,     SCREAMING_SNAKE_CASE__       ):
              lowerCAmelCase__          =			value






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

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





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




    """simple docstring"""

    def __init__(	self					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					int       )				->	Optional[Any]:
         super().__init__(SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			config

         lowerCAmelCase__          =			PoolFormerEncoder(SCREAMING_SNAKE_CASE__       )

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


    def        a				(	self					:					Optional[int]       )				->	Optional[Any]:
         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        a				(	self					:					Optional[int]       ,     SCREAMING_SNAKE_CASE__					:					Optional[torch.FloatTensor] = None       ,     SCREAMING_SNAKE_CASE__					:					Optional[bool] = None       ,     SCREAMING_SNAKE_CASE__					:					Optional[bool] = None       ,     )				->	Union[Tuple, BaseModelOutputWithNoAttention]:
         lowerCAmelCase__          =			(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         lowerCAmelCase__          =			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"       )

         lowerCAmelCase__          =			self.encoder(
             SCREAMING_SNAKE_CASE__       ,     output_hidden_states=SCREAMING_SNAKE_CASE__       ,     return_dict=SCREAMING_SNAKE_CASE__       ,     )
         lowerCAmelCase__          =			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 			__lowerCamelCase	(   nn.Module   ):




    """simple docstring"""

    def __init__(	self					:					Optional[int]       ,     SCREAMING_SNAKE_CASE__					:					List[Any]       )				->	List[str]:
         super().__init__()
         lowerCAmelCase__          =			nn.Linear(config.hidden_size       ,     config.hidden_size       )


    def        a				(	self					:					Any       ,     SCREAMING_SNAKE_CASE__					:					Optional[int]       )				->	int:
         lowerCAmelCase__          =			self.dense(SCREAMING_SNAKE_CASE__       )
         return output






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




    """simple docstring"""

    def __init__(	self					:					List[str]       ,     SCREAMING_SNAKE_CASE__					:					List[Any]       )				->	Optional[int]:
         super().__init__(SCREAMING_SNAKE_CASE__       )
         lowerCAmelCase__          =			config.num_labels
         lowerCAmelCase__          =			PoolFormerModel(SCREAMING_SNAKE_CASE__       )

         # Final norm
         lowerCAmelCase__          =			PoolFormerGroupNorm(config.hidden_sizes[-1]       )
         # Classifier head
         lowerCAmelCase__          =			(
             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        a				(	self					:					Tuple       ,     SCREAMING_SNAKE_CASE__					:					Optional[torch.FloatTensor] = None       ,     SCREAMING_SNAKE_CASE__					:					Optional[torch.LongTensor] = None       ,     SCREAMING_SNAKE_CASE__					:					Optional[bool] = None       ,     SCREAMING_SNAKE_CASE__					:					Optional[bool] = None       ,     )				->	Union[Tuple, ImageClassifierOutputWithNoAttention]:
         lowerCAmelCase__          =			return_dict if return_dict is not None else self.config.use_return_dict

         lowerCAmelCase__          =			self.poolformer(
             SCREAMING_SNAKE_CASE__       ,     output_hidden_states=SCREAMING_SNAKE_CASE__       ,     return_dict=SCREAMING_SNAKE_CASE__       ,     )

         lowerCAmelCase__          =			outputs[0]

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

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

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

         if not return_dict:
              lowerCAmelCase__          =			(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       )
