





from __future__ import annotations

import random
import unittest

from transformers import TransfoXLConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

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 (
          TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
          TFTransfoXLForSequenceClassification,
          TFTransfoXLLMHeadModel,
          TFTransfoXLModel,
      )

class    _UpperCAmelCase							:
      def __init__(		self      ,       a__      ,       ):
         A_     :				str			      =	parent
         A_     :				Any			      =	13
         A_     :				List[Any]			      =	7
         A_     :				Optional[Any]			      =	30
         A_     :				Optional[Any]			      =	self.seq_length + self.mem_len
         A_     :				Dict			      =	15
         A_     :				Union[str, Any]			      =	True
         A_     :				int			      =	True
         A_     :				List[str]			      =	99
         A_     :				Dict			      =	[10, 50, 80]
         A_     :				Dict			      =	32
         A_     :				Optional[Any]			      =	32
         A_     :				List[Any]			      =	4
         A_     :				Optional[Any]			      =	8
         A_     :				Any			      =	128
         A_     :				Optional[int]			      =	2
         A_     :				Union[str, Any]			      =	2
         A_     :				Union[str, Any]			      =	None
         A_     :				Union[str, Any]			      =	1
         A_     :				Optional[Any]			      =	0
         A_     :				Dict			      =	3
         A_     :				List[Any]			      =	self.vocab_size - 1
         A_     :				int			      =	0.01
      def 		_lowerCamelCase      (		self       ):
         A_     :				Tuple			      =	ids_tensor([self.batch_size, self.seq_length]      ,       self.vocab_size       )
         A_     :				Any			      =	ids_tensor([self.batch_size, self.seq_length]      ,       self.vocab_size       )

         A_     :				List[str]			      =	None
         if self.use_labels:
            A_     :				Optional[Any]			      =	ids_tensor([self.batch_size, self.seq_length]      ,       self.vocab_size       )

         A_     :				Optional[int]			      =	TransfoXLConfig(
             vocab_size=self.vocab_size      ,       mem_len=self.mem_len      ,       clamp_len=self.clamp_len      ,       cutoffs=self.cutoffs      ,       d_model=self.hidden_size      ,       d_embed=self.d_embed      ,       n_head=self.num_attention_heads      ,       d_head=self.d_head      ,       d_inner=self.d_inner      ,       div_val=self.div_val      ,       n_layer=self.num_hidden_layers      ,       eos_token_id=self.eos_token_id      ,       pad_token_id=self.vocab_size - 1      ,       init_range=self.init_range      ,       num_labels=self.num_labels      ,       )

         return (config, input_ids_a, input_ids_a, lm_labels)
      def 		_lowerCamelCase      (		self       ):
         random.seed(self.seed       )
         tf.random.set_seed(self.seed       )
      def 		_lowerCamelCase      (		self      ,       a__      ,       a__      ,       a__      ,       a__       ):
         A_     :				Any			      =	TFTransfoXLModel(a__       )

         A_      ,				A_     :				str			      =	model(a__       ).to_tuple()

         A_     :				Union[str, Any]			      =	{"""input_ids""": input_ids_a, """mems""": mems_a}

         A_      ,				A_     :				Optional[int]			      =	model(a__       ).to_tuple()

         self.parent.assertEqual(hidden_states_a.shape      ,       (self.batch_size, self.seq_length, self.hidden_size)       )
         self.parent.assertEqual(hidden_states_a.shape      ,       (self.batch_size, self.seq_length, self.hidden_size)       )
         self.parent.assertListEqual(
             [mem.shape for mem in mems_a]      ,       [(self.mem_len, self.batch_size, self.hidden_size)] * self.num_hidden_layers      ,       )
         self.parent.assertListEqual(
             [mem.shape for mem in mems_a]      ,       [(self.mem_len, self.batch_size, self.hidden_size)] * self.num_hidden_layers      ,       )
      def 		_lowerCamelCase      (		self      ,       a__      ,       a__      ,       a__      ,       a__       ):
         A_     :				Dict			      =	TFTransfoXLLMHeadModel(a__       )

         A_      ,				A_     :				List[Any]			      =	model(a__       ).to_tuple()

         A_     :				int			      =	{"""input_ids""": input_ids_a, """labels""": lm_labels}
         A_      ,				A_     :				Any			      =	model(a__       ).to_tuple()

         A_      ,				A_     :				List[str]			      =	model([input_ids_a, mems_a]       ).to_tuple()

         A_     :				List[Any]			      =	{"""input_ids""": input_ids_a, """mems""": mems_a, """labels""": lm_labels}

         A_      ,				A_     :				str			      =	model(a__       ).to_tuple()

         self.parent.assertEqual(lm_logits_a.shape      ,       (self.batch_size, self.seq_length, self.vocab_size)       )
         self.parent.assertListEqual(
             [mem.shape for mem in mems_a]      ,       [(self.mem_len, self.batch_size, self.hidden_size)] * self.num_hidden_layers      ,       )

         self.parent.assertEqual(lm_logits_a.shape      ,       (self.batch_size, self.seq_length, self.vocab_size)       )
         self.parent.assertListEqual(
             [mem.shape for mem in mems_a]      ,       [(self.mem_len, self.batch_size, self.hidden_size)] * self.num_hidden_layers      ,       )
      def 		_lowerCamelCase      (		self      ,       a__      ,       a__      ,       a__      ,       a__       ):
         A_     :				List[Any]			      =	TFTransfoXLForSequenceClassification(a__       )
         A_     :				Optional[Any]			      =	model(a__       )
         self.parent.assertEqual(result.logits.shape      ,       (self.batch_size, self.num_labels)       )






      def 		_lowerCamelCase      (		self       ):
         A_     :				Tuple			      =	self.prepare_config_and_inputs()
         ((A_)      ,				(A_)      ,				(A_)      ,				(A_))     :				Tuple			      =	config_and_inputs
         A_     :				Optional[int]			      =	{"""input_ids""": input_ids_a}
         return config, inputs_dict






@require_tf
class    _UpperCAmelCase							(		_lowerCamelCase			,       _lowerCamelCase			,       unittest.TestCase     ):
      a						       =							(
          (TFTransfoXLModel, TFTransfoXLLMHeadModel, TFTransfoXLForSequenceClassification) if is_tf_available() else ()
      )
      a						       =							() if is_tf_available() else ()
      a						       =							(
          {
              '''feature-extraction''': TFTransfoXLModel,
              '''text-classification''': TFTransfoXLForSequenceClassification,
              '''text-generation''': TFTransfoXLLMHeadModel,
              '''zero-shot''': TFTransfoXLForSequenceClassification,
          }
          if is_tf_available()
          else {}
      )
      # TODO: add this test when TFTransfoXLLMHead has a linear output layer implemented
      a						       =							False
      a						       =							False
      a						       =							False
      a						       =							False
      def 		_lowerCamelCase      (		self      ,       a__      ,       a__      ,       a__      ,       a__      ,       a__       ):
         if pipeline_test_casse_name == "TextGenerationPipelineTests":
            # Get `ValueError: AttributeError: 'NoneType' object has no attribute 'new_ones'` or `AssertionError`.
            # `TransfoXLConfig` was never used in pipeline tests: cannot create a simple
            # tokenizer.
            return True

         return False
      def 		_lowerCamelCase      (		self       ):
         A_     :				Optional[int]			      =	TFTransfoXLModelTester(self       )
         A_     :				List[Any]			      =	ConfigTester(self      ,       config_class=a__      ,       d_embed=37       )
      def 		_lowerCamelCase      (		self       ):
         self.config_tester.run_common_tests()
      def 		_lowerCamelCase      (		self       ):
         self.model_tester.set_seed()
         A_     :				Dict			      =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_transfo_xl_model(*a__       )
      def 		_lowerCamelCase      (		self       ):
         self.model_tester.set_seed()
         A_     :				Optional[Any]			      =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_transfo_xl_lm_head(*a__       )
      def 		_lowerCamelCase      (		self       ):
         A_     :				Union[str, Any]			      =	self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_transfo_xl_for_sequence_classification(*a__       )
      def 		_lowerCamelCase      (		self       ):
         A_      ,				A_     :				Any			      =	self.model_tester.prepare_config_and_inputs_for_common()
         A_     :				Dict			      =	[TFTransfoXLForSequenceClassification]

         for model_class in self.all_model_classes:
            A_     :				List[Any]			      =	model_class(a__       )
            assert isinstance(model.get_input_embeddings()      ,       tf.keras.layers.Layer       )
            if model_class in list_other_models_with_output_ebd:
               A_     :				Optional[int]			      =	model.get_output_embeddings()
               assert isinstance(a__      ,       tf.keras.layers.Layer       )
               A_     :				Union[str, Any]			      =	model.get_bias()
               assert name is None
            else:
               A_     :				Union[str, Any]			      =	model.get_output_embeddings()
               assert x is None
               A_     :				List[str]			      =	model.get_bias()
               assert name is None
      def 		_lowerCamelCase      (		self       ):
         # TODO JP: Make TransfoXL XLA compliant
         pass
      @slow
      def 		_lowerCamelCase      (		self       ):
         for model_name in TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            A_     :				Tuple			      =	TFTransfoXLModel.from_pretrained(a__       )
            self.assertIsNotNone(a__       )






      @unittest.skip(reason="""This model doesn't play well with fit() due to not returning a single loss."""       )
      def 		_lowerCamelCase      (		self       ):
         pass






@require_tf
class    _UpperCAmelCase							(		unittest.TestCase     ):
      @unittest.skip("""Skip test until #12651 is resolved."""       )
      @slow
      def 		_lowerCamelCase      (		self       ):
         A_     :				List[Any]			      =	TFTransfoXLLMHeadModel.from_pretrained("""transfo-xl-wt103"""       )
         # fmt: off
         A_     :				Optional[int]			      =	tf.convert_to_tensor([[33,1297,2,1,1009,4,1109,11739,4762,358,5,25,245,22,1706,17,20098,5,3215,21,37,1110,3,13,1041,4,24,603,490,2,71477,20098,104447,2,20961,1,2604,4,1,329,3,6224,831,16002,2,8,603,78967,29546,23,803,20,25,416,5,8,232,4,277,6,1855,4601,3,29546,54,8,3609,5,57211,49,4,1,277,18,8,1755,15691,3,341,25,416,693,42573,71,17,401,94,31,17919,2,29546,7873,18,1,435,23,11011,755,5,5167,3,7983,98,84,2,29546,3267,8,3609,4,1,4865,1075,2,6087,71,6,346,8,5854,3,29546,824,1400,1868,2,19,160,2,311,8,5496,2,20920,17,25,15097,3,24,24,0]]      ,       dtype=tf.intaa       )  # noqa: E231
         # fmt: on
         #  In 1991 , the remains of Russian Tsar Nicholas II and his family
         #  ( except for Alexei and Maria ) are discovered .
         #  The voice of Nicholas's young son , Tsarevich Alexei Nikolaevich , narrates the
         #  remainder of the story . 1883 Western Siberia ,
         #  a young Grigori Rasputin is asked by his father and a group of men to perform magic .
         #  Rasputin has a vision and denounces one of the men as a horse thief . Although his
         #  father initially slaps him for making such an accusation , Rasputin watches as the
         #  man is chased outside and beaten . Twenty years later , Rasputin sees a vision of
         #  the Virgin Mary , prompting him to become a priest . Rasputin quickly becomes famous ,
         #  with people , even a bishop , begging for his blessing . <eod> </s> <eos>

         # fmt: off
         A_     :				Any			      =	[33,1297,2,1,1009,4,1109,11739,4762,358,5,25,245,22,1706,17,20098,5,3215,21,37,1110,3,13,1041,4,24,603,490,2,71477,20098,104447,2,20961,1,2604,4,1,329,3,6224,831,16002,2,8,603,78967,29546,23,803,20,25,416,5,8,232,4,277,6,1855,4601,3,29546,54,8,3609,5,57211,49,4,1,277,18,8,1755,15691,3,341,25,416,693,42573,71,17,401,94,31,17919,2,29546,7873,18,1,435,23,11011,755,5,5167,3,7983,98,84,2,29546,3267,8,3609,4,1,4865,1075,2,6087,71,6,346,8,5854,3,29546,824,1400,1868,2,19,160,2,311,8,5496,2,20920,17,25,15097,3,24,24,0,33,1,1857,2,1,1009,4,1109,11739,4762,358,5,25,245,28,1110,3,13,1041,4,24,603,490,2,71477,20098,104447,2,20961,1,2604,4,1,329,3,0]  # noqa: E231
         # fmt: on
         #  In 1991, the remains of Russian Tsar Nicholas II and his family (
         #  except for Alexei and Maria ) are discovered. The voice of young son,
         #  Tsarevich Alexei Nikolaevich, narrates the remainder of the story.
         #  1883 Western Siberia, a young Grigori Rasputin is asked by his father
         #  and a group of men to perform magic. Rasputin has a vision and
         #  denounces one of the men as a horse thief. Although his father initially
         #  slaps him for making such an accusation, Rasputin watches as the man
         #  is chased outside and beaten. Twenty years later, Rasputin sees a vision
         #  of the Virgin Mary, prompting him to become a priest.
         #  Rasputin quickly becomes famous, with people, even a bishop, begging for
         #  his blessing. <unk> <unk> <eos> In the 1990s, the remains of Russian Tsar
         # Nicholas II and his family were discovered. The voice of <unk> young son,
         # Tsarevich Alexei Nikolaevich, narrates the remainder of the story.<eos>

         A_     :				List[Any]			      =	model.generate(a__      ,       max_length=200      ,       do_sample=a__       )
         self.assertListEqual(output_ids[0].numpy().tolist()      ,       a__       )













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

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

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


_lowerCAmelCase 				=     logging.get_logger(__name__)

# General docstring
_lowerCAmelCase 				=     """PoolFormerConfig"""

# Base docstring
_lowerCAmelCase 				=     """sail/poolformer_s12"""
_lowerCAmelCase 				=     [1, 512, 7, 7]

# Image classification docstring
_lowerCAmelCase 				=     """sail/poolformer_s12"""
_lowerCAmelCase 				=     """tabby, tabby cat"""

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



def     _lowerCAmelCase		(  _lowerCAmelCase  ,_lowerCAmelCase = 0.0  ,_lowerCAmelCase = False     ):
   '''simple docstring'''


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

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






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






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

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






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






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






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






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






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






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

         return hidden_states






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

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






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

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

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

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

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

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






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

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

         # Transformer blocks
         A_     :				Union[str, Any]			      =	[]
         A_     :				Union[str, Any]			      =	0
         for i in range(config.num_encoder_blocks       ):
            # each block consists of layers
            A_     :				Dict			      =	[]
            if i != 0:
               cur += config.depths[i - 1]
            for j in range(config.depths[i]       ):
               layers.append(
                   PoolFormerLayer(
                       a__      ,       num_channels=config.hidden_sizes[i]      ,       pool_size=config.pool_size      ,       hidden_size=config.hidden_sizes[i]      ,       intermediate_size=int(config.hidden_sizes[i] * config.mlp_ratio       )      ,       drop_path=dpr[cur + j]      ,       )       )
            blocks.append(nn.ModuleList(a__       )       )

         A_     :				str			      =	nn.ModuleList(a__       )






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

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

            if output_hidden_states:
               A_     :				Dict			      =	all_hidden_states + (hidden_states,)

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

         return BaseModelOutputWithNoAttention(last_hidden_state=a__      ,       hidden_states=a__       )






class    _UpperCAmelCase							(		_lowerCamelCase     ):
      a						       =							PoolFormerConfig
      a						       =							'''poolformer'''
      a						       =							'''pixel_values'''
      a						       =							True
      def 		_lowerCamelCase      (		self      ,       a__       ):
         if isinstance(a__      ,       (nn.Linear, nn.Convad)       ):
            module.weight.data.normal_(mean=0.0      ,       std=self.config.initializer_range       )
            if module.bias is not None:
               module.bias.data.zero_()
         elif isinstance(a__      ,       nn.LayerNorm       ):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0       )






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






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

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

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

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

         A_     :				Optional[int]			      =	PoolFormerEncoder(a__       )

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






      @add_start_docstrings_to_model_forward(a__       )
      @add_code_sample_docstrings(
          checkpoint=_CHECKPOINT_FOR_DOC      ,       output_type=a__      ,       config_class=_CONFIG_FOR_DOC      ,       modality="""vision"""      ,       expected_output=_EXPECTED_OUTPUT_SHAPE      ,       )
      def 		_lowerCamelCase      (		self      ,       a__ = None      ,       a__ = None      ,       a__ = None      ,       ):
         A_     :				Union[str, Any]			      =	(
             output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
         )
         A_     :				str			      =	return_dict if return_dict is not None else self.config.use_return_dict

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

         A_     :				Optional[Any]			      =	self.encoder(
             a__      ,       output_hidden_states=a__      ,       return_dict=a__      ,       )
         A_     :				Optional[int]			      =	encoder_outputs[0]

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

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






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






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






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

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

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






      @add_start_docstrings_to_model_forward(a__       )
      @add_code_sample_docstrings(
          checkpoint=_IMAGE_CLASS_CHECKPOINT      ,       output_type=a__      ,       config_class=_CONFIG_FOR_DOC      ,       expected_output=_IMAGE_CLASS_EXPECTED_OUTPUT      ,       )
      def 		_lowerCamelCase      (		self      ,       a__ = None      ,       a__ = None      ,       a__ = None      ,       a__ = None      ,       ):
         A_     :				List[Any]			      =	return_dict if return_dict is not None else self.config.use_return_dict

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

         A_     :				List[str]			      =	outputs[0]

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

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

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

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

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






