

import argparse
import logging
import os
from pathlib import Path
from typing import Any, Dict

import pytorch_lightning as pl
from pytorch_lightning.utilities import rank_zero_info

from transformers import (
    AdamW,
    AutoConfig,
    AutoModel,
    AutoModelForPreTraining,
    AutoModelForQuestionAnswering,
    AutoModelForSeqaSeqLM,
    AutoModelForSequenceClassification,
    AutoModelForTokenClassification,
    AutoModelWithLMHead,
    AutoTokenizer,
    PretrainedConfig,
    PreTrainedTokenizer,
)
from transformers.optimization import (
    Adafactor,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.utils.versions import require_version


_a			=       logging.getLogger(__name__)

require_version('''pytorch_lightning>=1.0.4''')

_a			=       {
    '''base''': AutoModel,
    '''sequence-classification''': AutoModelForSequenceClassification,
    '''question-answering''': AutoModelForQuestionAnswering,
    '''pretraining''': AutoModelForPreTraining,
    '''token-classification''': AutoModelForTokenClassification,
    '''language-modeling''': AutoModelWithLMHead,
    '''summarization''': AutoModelForSeqaSeqLM,
    '''translation''': AutoModelForSeqaSeqLM,
}


# update this and the import above to support new schedulers from transformers.optimization
_a			=       {
    '''linear''': get_linear_schedule_with_warmup,
    '''cosine''': get_cosine_schedule_with_warmup,
    '''cosine_w_restarts''': get_cosine_with_hard_restarts_schedule_with_warmup,
    '''polynomial''': get_polynomial_decay_schedule_with_warmup,
    # '': get_constant_schedule,             # not supported for now
    # '': get_constant_schedule_with_warmup, # not supported for now
}
_a			=       sorted(arg_to_scheduler.keys())
_a			=       '''{''' + ''', '''.join(arg_to_scheduler_choices) + '''}'''






class 							A_			(		pl.LightningModule ):


     def __init__(							self      :		Union[str, Any]      ,  UpperCAmelCase      :		argparse.Namespace      ,  UpperCAmelCase      :		List[str]=None      ,  UpperCAmelCase      :		Union[str, Any]="base"      ,  UpperCAmelCase      :		Optional[int]=None      ,  UpperCAmelCase      :		List[str]=None      ,  UpperCAmelCase      :		List[str]=None      ,  **UpperCAmelCase      :		List[Any]      ,  )							->     Union[str, Any]:
      super().__init__()
      # TODO: move to self.save_hyperparameters()
      # self.save_hyperparameters()
      # can also expand arguments into trainer signature for easier reading

      self.save_hyperparameters(UpperCAmelCase )
      __lowerCAmelCase:       Any            =				0
      __lowerCAmelCase:       Tuple            =				Path(self.hparams.output_dir )
      __lowerCAmelCase:       int            =				self.hparams.cache_dir if self.hparams.cache_dir else None
      if config is None:
       __lowerCAmelCase:       int            =				AutoConfig.from_pretrained(
           self.hparams.config_name if self.hparams.config_name else self.hparams.model_name_or_path      ,  **({'num_labels': num_labels} if num_labels is not None else {})      ,  cache_dir=UpperCAmelCase      ,  **UpperCAmelCase      ,  )
      else:
       __lowerCAmelCase:       PretrainedConfig            =				config

      __lowerCAmelCase:       Any            =				('encoder_layerdrop', 'decoder_layerdrop', 'dropout', 'attention_dropout')
      for p in extra_model_params:
       if getattr(self.hparams      ,  UpperCAmelCase      ,  UpperCAmelCase ):
        assert hasattr(self.config      ,  UpperCAmelCase ), F'''model config doesn\'t have a `{p}` attribute'''
        setattr(self.config      ,  UpperCAmelCase      ,  getattr(self.hparams      ,  UpperCAmelCase ) )

      if tokenizer is None:
       __lowerCAmelCase:       Tuple            =				AutoTokenizer.from_pretrained(
           self.hparams.tokenizer_name if self.hparams.tokenizer_name else self.hparams.model_name_or_path      ,  cache_dir=UpperCAmelCase      ,  )
      else:
       __lowerCAmelCase:       PreTrainedTokenizer            =				tokenizer
      __lowerCAmelCase:       Tuple            =				MODEL_MODES[mode]
      if model is None:
       __lowerCAmelCase:       Union[str, Any]            =				self.model_type.from_pretrained(
           self.hparams.model_name_or_path      ,  from_tf=bool('.ckpt' in self.hparams.model_name_or_path )      ,  config=self.config      ,  cache_dir=UpperCAmelCase      ,  )
      else:
       __lowerCAmelCase:       int            =				model


     def   UpperCAmelCase  (							self      :		str      ,  *UpperCAmelCase      :		Any      ,  **UpperCAmelCase      :		Optional[int] )							->     Dict:
      __lowerCAmelCase:       Tuple            =				self.model_type.from_pretrained(*UpperCAmelCase      ,  **UpperCAmelCase )


     def   UpperCAmelCase  (							self      :		Union[str, Any] )							->     List[Any]:
      __lowerCAmelCase:       Any            =				arg_to_scheduler[self.hparams.lr_scheduler]
      __lowerCAmelCase:       Optional[int]            =				get_schedule_func(
          self.opt      ,  num_warmup_steps=self.hparams.warmup_steps      ,  num_training_steps=self.total_steps() )
      __lowerCAmelCase:       Any            =				{'scheduler': scheduler, 'interval': 'step', 'frequency': 1}
      return scheduler


     def   UpperCAmelCase  (							self      :		Any )							->     Optional[Any]:
      __lowerCAmelCase:       Optional[int]            =				self.model
      __lowerCAmelCase:       Union[str, Any]            =				['bias', 'LayerNorm.weight']
      __lowerCAmelCase:       Optional[Any]            =				[
          {
              'params': [p for n, p in model.named_parameters() if not any(nd in n for nd in no_decay )],
              'weight_decay': self.hparams.weight_decay,
          },
          {
              'params': [p for n, p in model.named_parameters() if any(nd in n for nd in no_decay )],
              'weight_decay': 0.0,
          },
      ]
      if self.hparams.adafactor:
       __lowerCAmelCase:       str            =				Adafactor(
           UpperCAmelCase      ,  lr=self.hparams.learning_rate      ,  scale_parameter=UpperCAmelCase      ,  relative_step=UpperCAmelCase )

      else:
       __lowerCAmelCase:       List[Any]            =				AdamW(
           UpperCAmelCase      ,  lr=self.hparams.learning_rate      ,  eps=self.hparams.adam_epsilon )
      __lowerCAmelCase:       Dict            =				optimizer

      __lowerCAmelCase:       int            =				self.get_lr_scheduler()

      return [optimizer], [scheduler]


     def   UpperCAmelCase  (							self      :		Union[str, Any]      ,  UpperCAmelCase      :		Optional[int]      ,  UpperCAmelCase      :		int )							->     List[Any]:
      return self.validation_step(UpperCAmelCase      ,  UpperCAmelCase )


     def   UpperCAmelCase  (							self      :		List[str]      ,  UpperCAmelCase      :		Any )							->     Optional[Any]:
      return self.validation_end(UpperCAmelCase )


     def   UpperCAmelCase  (							self      :		Optional[int] )							->     int:
      __lowerCAmelCase:       Any            =				max(1      ,  self.hparams.gpus )  # TODO: consider num_tpu_cores
      __lowerCAmelCase:       Optional[int]            =				self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
      return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs


     def   UpperCAmelCase  (							self      :		List[Any]      ,  UpperCAmelCase      :		Optional[Any] )							->     str:
      if stage == "test":
       __lowerCAmelCase:       Optional[int]            =				len(self.test_dataloader().dataset )
      else:
       __lowerCAmelCase:       Optional[Any]            =				self.get_dataloader('train'      ,  self.hparams.train_batch_size      ,  shuffle=UpperCAmelCase )
       __lowerCAmelCase:       int            =				len(self.train_dataloader().dataset )


     def   UpperCAmelCase  (							self      :		Optional[int]      ,  UpperCAmelCase      :		str      ,  UpperCAmelCase      :		int      ,  UpperCAmelCase      :		bool = False )							->     Optional[Any]:
      raise NotImplementedError('You must implement this for your task' )


     def   UpperCAmelCase  (							self      :		List[str] )							->     List[Any]:
      return self.train_loader


     def   UpperCAmelCase  (							self      :		Tuple )							->     Optional[Any]:
      return self.get_dataloader('dev'      ,  self.hparams.eval_batch_size      ,  shuffle=UpperCAmelCase )


     def   UpperCAmelCase  (							self      :		Any )							->     Optional[Any]:
      return self.get_dataloader('test'      ,  self.hparams.eval_batch_size      ,  shuffle=UpperCAmelCase )


     def   UpperCAmelCase  (							self      :		str      ,  UpperCAmelCase      :		Tuple )							->     List[Any]:
      return os.path.join(
          self.hparams.data_dir      ,  'cached_{}_{}_{}'.format(
              UpperCAmelCase      ,  list(filter(UpperCAmelCase      ,  self.hparams.model_name_or_path.split('/' ) ) ).pop()      ,  str(self.hparams.max_seq_length )      ,  )      ,  )


     @pl.utilities.rank_zero_only
     def   UpperCAmelCase  (							self      :		Tuple      ,  UpperCAmelCase      :		Dict[str, Any] )							->     None:
      __lowerCAmelCase:       Dict            =				self.output_dir.joinpath('best_tfmr' )
      __lowerCAmelCase:       List[Any]            =				self.step_count
      self.model.save_pretrained(UpperCAmelCase )
      self.tokenizer.save_pretrained(UpperCAmelCase )


     @staticmethod
     def   UpperCAmelCase  (							UpperCAmelCase      :		str      ,  UpperCAmelCase      :		int )							->     Optional[Any]:
      parser.add_argument(
          '--model_name_or_path'      ,  default=UpperCAmelCase      ,  type=UpperCAmelCase      ,  required=UpperCAmelCase      ,  help='Path to pretrained model or model identifier from huggingface.co/models'      ,  )
      parser.add_argument(
          '--config_name'      ,  default=''      ,  type=UpperCAmelCase      ,  help='Pretrained config name or path if not the same as model_name' )
      parser.add_argument(
          '--tokenizer_name'      ,  default=UpperCAmelCase      ,  type=UpperCAmelCase      ,  help='Pretrained tokenizer name or path if not the same as model_name'      ,  )
      parser.add_argument(
          '--cache_dir'      ,  default=''      ,  type=UpperCAmelCase      ,  help='Where do you want to store the pre-trained models downloaded from huggingface.co'      ,  )
      parser.add_argument(
          '--encoder_layerdrop'      ,  type=UpperCAmelCase      ,  help='Encoder layer dropout probability (Optional). Goes into model.config'      ,  )
      parser.add_argument(
          '--decoder_layerdrop'      ,  type=UpperCAmelCase      ,  help='Decoder layer dropout probability (Optional). Goes into model.config'      ,  )
      parser.add_argument(
          '--dropout'      ,  type=UpperCAmelCase      ,  help='Dropout probability (Optional). Goes into model.config'      ,  )
      parser.add_argument(
          '--attention_dropout'      ,  type=UpperCAmelCase      ,  help='Attention dropout probability (Optional). Goes into model.config'      ,  )
      parser.add_argument('--learning_rate'      ,  default=5E-5      ,  type=UpperCAmelCase      ,  help='The initial learning rate for Adam.' )
      parser.add_argument(
          '--lr_scheduler'      ,  default='linear'      ,  choices=UpperCAmelCase      ,  metavar=UpperCAmelCase      ,  type=UpperCAmelCase      ,  help='Learning rate scheduler'      ,  )
      parser.add_argument('--weight_decay'      ,  default=0.0      ,  type=UpperCAmelCase      ,  help='Weight decay if we apply some.' )
      parser.add_argument('--adam_epsilon'      ,  default=1E-8      ,  type=UpperCAmelCase      ,  help='Epsilon for Adam optimizer.' )
      parser.add_argument('--warmup_steps'      ,  default=0      ,  type=UpperCAmelCase      ,  help='Linear warmup over warmup_steps.' )
      parser.add_argument('--num_workers'      ,  default=4      ,  type=UpperCAmelCase      ,  help='kwarg passed to DataLoader' )
      parser.add_argument('--num_train_epochs'      ,  dest='max_epochs'      ,  default=3      ,  type=UpperCAmelCase )
      parser.add_argument('--train_batch_size'      ,  default=3_2      ,  type=UpperCAmelCase )
      parser.add_argument('--eval_batch_size'      ,  default=3_2      ,  type=UpperCAmelCase )
      parser.add_argument('--adafactor'      ,  action='store_true' )







class 							A_			(		pl.Callback ):


     def   UpperCAmelCase  (							self      :		Optional[int]      ,  UpperCAmelCase      :		Any      ,  UpperCAmelCase      :		Union[str, Any] )							->     List[str]:
      if (
          trainer.is_global_zero and trainer.global_rank == 0
      ):  # we initialize the retriever only on master worker with RAY. In new pytorch-lightning accelorators are removed.
       pl_module.model.rag.retriever.init_retrieval()  # better to use hook functions.







class 							A_			(		pl.Callback ):


     def   UpperCAmelCase  (							self      :		str      ,  UpperCAmelCase      :		Dict      ,  UpperCAmelCase      :		List[Any] )							->     int:
      __lowerCAmelCase:       List[str]            =				trainer.lr_schedulers[0]['scheduler']
      __lowerCAmelCase:       str            =				{F'''lr_group_{i}''': lr for i, lr in enumerate(lr_scheduler.get_lr() )}
      pl_module.logger.log_metrics(UpperCAmelCase )


     def   UpperCAmelCase  (							self      :		Union[str, Any]      ,  UpperCAmelCase      :		pl.Trainer      ,  UpperCAmelCase      :		pl.LightningModule )							->     Union[str, Any]:
      rank_zero_info('***** Validation results *****' )
      __lowerCAmelCase:       Tuple            =				trainer.callback_metrics
      # Log results
      for key in sorted(UpperCAmelCase ):
       if key not in ["log", "progress_bar"]:
        rank_zero_info('{} = {}\n'.format(UpperCAmelCase      ,  str(metrics[key] ) ) )


     def   UpperCAmelCase  (							self      :		Union[str, Any]      ,  UpperCAmelCase      :		pl.Trainer      ,  UpperCAmelCase      :		pl.LightningModule )							->     Any:
      rank_zero_info('***** Test results *****' )
      __lowerCAmelCase:       List[str]            =				trainer.callback_metrics
      # Log and save results to file
      __lowerCAmelCase:       Tuple            =				os.path.join(pl_module.hparams.output_dir      ,  'test_results.txt' )
      with open(UpperCAmelCase      ,  'w' ) as writer:
       for key in sorted(UpperCAmelCase ):
        if key not in ["log", "progress_bar"]:
         rank_zero_info('{} = {}\n'.format(UpperCAmelCase      ,  str(metrics[key] ) ) )
         writer.write('{} = {}\n'.format(UpperCAmelCase      ,  str(metrics[key] ) ) )







def _a			(	SCREAMING_SNAKE_CASE		:      int   , SCREAMING_SNAKE_CASE		:      Dict		)   ->   None:

 """simple docstring"""






 parser.add_argument(
     '--output_dir'   , default=SCREAMING_SNAKE_CASE   , type=SCREAMING_SNAKE_CASE   , required=SCREAMING_SNAKE_CASE   , help='The output directory where the model predictions and checkpoints will be written.'   , )
 parser.add_argument(
     '--fp16'   , action='store_true'   , help='Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit'   , )

 parser.add_argument(
     '--fp16_opt_level'   , type=SCREAMING_SNAKE_CASE   , default='O2'   , help=(
         'For fp16: Apex AMP optimization level selected in [\'O0\', \'O1\', \'O2\', and \'O3\'].'
         'See details at https://nvidia.github.io/apex/amp.html'
     )   , )
 parser.add_argument('--n_tpu_cores'   , dest='tpu_cores'   , type=SCREAMING_SNAKE_CASE		)
 parser.add_argument('--max_grad_norm'   , dest='gradient_clip_val'   , default=1.0   , type=SCREAMING_SNAKE_CASE   , help='Max gradient norm'		)
 parser.add_argument('--do_train'   , action='store_true'   , help='Whether to run training.'		)
 parser.add_argument('--do_predict'   , action='store_true'   , help='Whether to run predictions on the test set.'		)
 parser.add_argument(
     '--gradient_accumulation_steps'   , dest='accumulate_grad_batches'   , type=SCREAMING_SNAKE_CASE   , default=1   , help='Number of updates steps to accumulate before performing a backward/update pass.'   , )
 parser.add_argument('--seed'   , type=SCREAMING_SNAKE_CASE   , default=42   , help='random seed for initialization'		)
 parser.add_argument(
     '--data_dir'   , default=SCREAMING_SNAKE_CASE   , type=SCREAMING_SNAKE_CASE   , required=SCREAMING_SNAKE_CASE   , help='The input data dir. Should contain the training files for the CoNLL-2003 NER task.'   , )







def _a			(	SCREAMING_SNAKE_CASE		:      BaseTransformer   , SCREAMING_SNAKE_CASE		:      argparse.Namespace   , SCREAMING_SNAKE_CASE		:      Tuple=None   , SCREAMING_SNAKE_CASE		:      Tuple=True   , SCREAMING_SNAKE_CASE		:      List[Any]=None   , SCREAMING_SNAKE_CASE		:      List[Any]=[]   , SCREAMING_SNAKE_CASE		:      Any=None   , SCREAMING_SNAKE_CASE		:      str=None   , **SCREAMING_SNAKE_CASE		:      List[str]   , )   ->   Optional[Any]:

 """simple docstring"""






 pl.seed_everything(args.seed		)

 # init model
 __lowerCAmelCase:       Dict            =				Path(model.hparams.output_dir		)
 odir.mkdir(exist_ok=SCREAMING_SNAKE_CASE		)

 # add custom checkpoints
 if checkpoint_callback is None:
  __lowerCAmelCase:       int            =				pl.callbacks.ModelCheckpoint(
      filepath=args.output_dir   , prefix='checkpoint'   , monitor='val_loss'   , mode='min'   , save_top_k=1		)
 if early_stopping_callback:
  extra_callbacks.append(SCREAMING_SNAKE_CASE		)
 if logging_callback is None:
  __lowerCAmelCase:       Union[str, Any]            =				LoggingCallback()

 __lowerCAmelCase:       Optional[Any]            =				{}

 # TODO: remove with PyTorch 1.6 since pl uses native amp
 if args.fpaa:
  __lowerCAmelCase:       str            =				16
  # train_params["amp_level"] = args.fp16_opt_level

 if args.gpus > 1:
  __lowerCAmelCase:       Any            =				'auto'  # "ddp"
  __lowerCAmelCase:       List[Any]            =				'ddp'

 __lowerCAmelCase:       Optional[int]            =				args.accumulate_grad_batches
 __lowerCAmelCase:       Union[str, Any]            =				None  # extra_train_kwargs.get("profiler", None) #get unwanted logs
 __lowerCAmelCase:       str            =				'auto'

 __lowerCAmelCase:       Optional[int]            =				pl.Trainer.from_argparse_args(
     SCREAMING_SNAKE_CASE   , weights_summary=SCREAMING_SNAKE_CASE   , callbacks=[logging_callback] + extra_callbacks + [checkpoint_callback] + [InitCallback()]   , logger=SCREAMING_SNAKE_CASE   , **SCREAMING_SNAKE_CASE   , )

 if args.do_train:
  trainer.fit(SCREAMING_SNAKE_CASE		)

 return trainer








import inspect
import unittest

from transformers import is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, slow, torch_device

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch

  from transformers import MODEL_MAPPING, PoolFormerConfig, PoolFormerForImageClassification, PoolFormerModel
  from transformers.models.poolformer.modeling_poolformer import POOLFORMER_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  from PIL import Image

  from transformers import PoolFormerImageProcessor






class 							A_			(		snake_case__ ):


     def   UpperCAmelCase  (							self      :		int )							->     Tuple:
      __lowerCAmelCase:       Optional[Any]            =				self.config_class(**self.inputs_dict )
      self.parent.assertTrue(hasattr(UpperCAmelCase      ,  'hidden_sizes' ) )
      self.parent.assertTrue(hasattr(UpperCAmelCase      ,  'num_encoder_blocks' ) )







class 							A_			:


     def __init__(							self      :		Optional[int]      ,  UpperCAmelCase      :		Any      ,  UpperCAmelCase      :		List[str]=1_3      ,  UpperCAmelCase      :		Optional[Any]=6_4      ,  UpperCAmelCase      :		Union[str, Any]=3      ,  UpperCAmelCase      :		Optional[Any]=4      ,  UpperCAmelCase      :		Any=[2, 2, 2, 2]      ,  UpperCAmelCase      :		Optional[Any]=[8, 4, 2, 1]      ,  UpperCAmelCase      :		str=[1_6, 3_2, 6_4, 1_2_8]      ,  UpperCAmelCase      :		str=[1, 4, 8, 1_6]      ,  UpperCAmelCase      :		Optional[Any]=False      ,  UpperCAmelCase      :		str=True      ,  UpperCAmelCase      :		Optional[int]="gelu"      ,  UpperCAmelCase      :		Union[str, Any]=0.1      ,  UpperCAmelCase      :		Optional[Any]=0.02      ,  UpperCAmelCase      :		List[Any]=3      ,  UpperCAmelCase      :		Optional[int]=None      ,  )							->     Any:
      __lowerCAmelCase:       List[str]            =				parent
      __lowerCAmelCase:       Union[str, Any]            =				batch_size
      __lowerCAmelCase:       Union[str, Any]            =				image_size
      __lowerCAmelCase:       str            =				num_channels
      __lowerCAmelCase:       Tuple            =				num_encoder_blocks
      __lowerCAmelCase:       Optional[int]            =				sr_ratios
      __lowerCAmelCase:       Union[str, Any]            =				depths
      __lowerCAmelCase:       int            =				hidden_sizes
      __lowerCAmelCase:       List[Any]            =				downsampling_rates
      __lowerCAmelCase:       List[Any]            =				is_training
      __lowerCAmelCase:       Union[str, Any]            =				use_labels
      __lowerCAmelCase:       Union[str, Any]            =				hidden_act
      __lowerCAmelCase:       Any            =				hidden_dropout_prob
      __lowerCAmelCase:       List[Any]            =				initializer_range
      __lowerCAmelCase:       Dict            =				num_labels
      __lowerCAmelCase:       List[Any]            =				scope


     def   UpperCAmelCase  (							self      :		Optional[int] )							->     List[str]:
      __lowerCAmelCase:       Dict            =				floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

      __lowerCAmelCase:       str            =				None
      if self.use_labels:
       __lowerCAmelCase:       int            =				ids_tensor([self.batch_size, self.image_size, self.image_size]      ,  self.num_labels )

      __lowerCAmelCase:       Dict            =				PoolFormerConfig(
          image_size=self.image_size      ,  num_channels=self.num_channels      ,  num_encoder_blocks=self.num_encoder_blocks      ,  depths=self.depths      ,  hidden_sizes=self.hidden_sizes      ,  hidden_act=self.hidden_act      ,  hidden_dropout_prob=self.hidden_dropout_prob      ,  initializer_range=self.initializer_range      ,  )

      return config, pixel_values, labels


     def   UpperCAmelCase  (							self      :		str      ,  UpperCAmelCase      :		Tuple      ,  UpperCAmelCase      :		int      ,  UpperCAmelCase      :		List[Any] )							->     str:
      __lowerCAmelCase:       Any            =				PoolFormerModel(config=UpperCAmelCase )
      model.to(UpperCAmelCase )
      model.eval()
      __lowerCAmelCase:       List[str]            =				model(UpperCAmelCase )
      __lowerCAmelCase:       Union[str, Any]            =				self.image_size // 32.0
      self.parent.assertEqual(
          result.last_hidden_state.shape      ,  (self.batch_size, self.hidden_sizes[-1], expected_height, expected_width) )


     def   UpperCAmelCase  (							self      :		Any )							->     int:
      __lowerCAmelCase:       List[str]            =				self.prepare_config_and_inputs()
      __lowerCAmelCase     ,							__lowerCAmelCase     ,							__lowerCAmelCase:       Union[str, Any]            =				config_and_inputs
      __lowerCAmelCase:       List[str]            =				{'pixel_values': pixel_values}
      return config, inputs_dict







@require_torch
class 							A_			(		snake_case__							,     snake_case__							,     unittest.TestCase ):
     _lowercase      :					Optional[Any]       =     (PoolFormerModel, PoolFormerForImageClassification) if is_torch_available() else ()
     _lowercase      :					List[Any]       =     (
         {'feature-extraction': PoolFormerModel, 'image-classification': PoolFormerForImageClassification}
         if is_torch_available()
         else {}
     )

     _lowercase      :					Tuple       =     False
     _lowercase      :					Tuple       =     False
     _lowercase      :					List[Any]       =     False
     _lowercase      :					Any       =     False
     _lowercase      :					List[Any]       =     False


     def   UpperCAmelCase  (							self      :		str )							->     Any:
      __lowerCAmelCase:       int            =				PoolFormerModelTester(self )
      __lowerCAmelCase:       int            =				PoolFormerConfigTester(self      ,  config_class=UpperCAmelCase )


     def   UpperCAmelCase  (							self      :		Optional[Any] )							->     int:
      self.config_tester.run_common_tests()


     def   UpperCAmelCase  (							self      :		int )							->     int:
      __lowerCAmelCase:       Union[str, Any]            =				self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*UpperCAmelCase )


     @unittest.skip('PoolFormer does not use inputs_embeds' )
     def   UpperCAmelCase  (							self      :		Optional[int] )							->     Union[str, Any]:
      pass


     @unittest.skip('PoolFormer does not have get_input_embeddings method and get_output_embeddings methods' )
     def   UpperCAmelCase  (							self      :		int )							->     List[str]:
      pass


     def   UpperCAmelCase  (							self      :		List[str] )							->     List[str]:
      def check_hidden_states_output(UpperCAmelCase      :		List[str]      ,  UpperCAmelCase      :		int      ,  UpperCAmelCase      :		List[Any] ):
       __lowerCAmelCase:       Union[str, Any]            =				model_class(UpperCAmelCase )
       model.to(UpperCAmelCase )
       model.eval()

       with torch.no_grad():
        __lowerCAmelCase:       Dict            =				model(**self._prepare_for_class(UpperCAmelCase      ,  UpperCAmelCase ) )

       __lowerCAmelCase:       int            =				outputs.hidden_states

       __lowerCAmelCase:       Optional[int]            =				self.model_tester.num_encoder_blocks
       self.assertEqual(len(UpperCAmelCase )      ,  UpperCAmelCase )

       # verify the first hidden states (first block)
       self.assertListEqual(
           list(hidden_states[0].shape[-3:] )      ,  [
               self.model_tester.hidden_sizes[0],
               self.model_tester.image_size // 4,
               self.model_tester.image_size // 4,
           ]      ,  )

      __lowerCAmelCase     ,							__lowerCAmelCase:       Dict            =				self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
       __lowerCAmelCase:       Optional[int]            =				True
       check_hidden_states_output(UpperCAmelCase      ,  UpperCAmelCase      ,  UpperCAmelCase )

       # check that output_hidden_states also work using config
       del inputs_dict["output_hidden_states"]
       __lowerCAmelCase:       Optional[Any]            =				True

       check_hidden_states_output(UpperCAmelCase      ,  UpperCAmelCase      ,  UpperCAmelCase )


     def   UpperCAmelCase  (							self      :		List[Any] )							->     Optional[int]:
      if not self.model_tester.is_training:
       return

      __lowerCAmelCase     ,							__lowerCAmelCase:       Union[str, Any]            =				self.model_tester.prepare_config_and_inputs_for_common()
      __lowerCAmelCase:       Optional[int]            =				True

      for model_class in self.all_model_classes:
       if model_class in get_values(UpperCAmelCase ):
        continue
       __lowerCAmelCase:       str            =				model_class(UpperCAmelCase )
       model.to(UpperCAmelCase )
       model.train()
       __lowerCAmelCase:       Any            =				self._prepare_for_class(UpperCAmelCase      ,  UpperCAmelCase      ,  return_labels=UpperCAmelCase )
       __lowerCAmelCase:       str            =				model(**UpperCAmelCase ).loss
       loss.backward()


     def   UpperCAmelCase  (							self      :		Dict )							->     Tuple:
      __lowerCAmelCase     ,							__lowerCAmelCase:       Union[str, Any]            =				self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
       __lowerCAmelCase:       Union[str, Any]            =				model_class(UpperCAmelCase )
       __lowerCAmelCase:       Dict            =				inspect.signature(model.forward )
       # signature.parameters is an OrderedDict => so arg_names order is deterministic
       __lowerCAmelCase:       List[str]            =				[*signature.parameters.keys()]

       __lowerCAmelCase:       List[Any]            =				['pixel_values']
       self.assertListEqual(arg_names[:1]      ,  UpperCAmelCase )


     @slow
     def   UpperCAmelCase  (							self      :		int )							->     str:
      for model_name in POOLFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
       __lowerCAmelCase:       List[str]            =				PoolFormerModel.from_pretrained(UpperCAmelCase )
       self.assertIsNotNone(UpperCAmelCase )







def _a			(	)   ->   Any:

 """simple docstring"""






 __lowerCAmelCase:       Optional[Any]            =				Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'		)
 return image






@require_torch
class 							A_			(		unittest.TestCase ):


     @slow
     def   UpperCAmelCase  (							self      :		Optional[Any] )							->     int:
      __lowerCAmelCase:       Optional[int]            =				PoolFormerImageProcessor()
      __lowerCAmelCase:       Any            =				PoolFormerForImageClassification.from_pretrained('sail/poolformer_s12' ).to(UpperCAmelCase )

      __lowerCAmelCase:       int            =				image_processor(images=prepare_img()      ,  return_tensors='pt' ).to(UpperCAmelCase )

      # forward pass
      with torch.no_grad():
       __lowerCAmelCase:       List[Any]            =				model(**UpperCAmelCase )

      # verify the logits
      __lowerCAmelCase:       List[Any]            =				torch.Size((1, 1_0_0_0) )
      self.assertEqual(outputs.logits.shape      ,  UpperCAmelCase )

      __lowerCAmelCase:       Union[str, Any]            =				torch.tensor([-0.6113, 0.1685, -0.0492] ).to(UpperCAmelCase )
      self.assertTrue(torch.allclose(outputs.logits[0, :3]      ,  UpperCAmelCase      ,  atol=1E-4 ) )





