




'''simple docstring'''



import json
import os
from dataclasses import dataclass
from functools import partial
from typing import Callable

import flax.linen as nn
import jax
import jax.numpy as jnp
import joblib
import optax
import wandb
from flax import jax_utils, struct, traverse_util
from flax.serialization import from_bytes, to_bytes
from flax.training import train_state
from flax.training.common_utils import shard
from tqdm.auto import tqdm

from transformers import BigBirdConfig, FlaxBigBirdForQuestionAnswering
from transformers.models.big_bird.modeling_flax_big_bird import FlaxBigBirdForQuestionAnsweringModule





class      lowerCAmelCase (			__UpperCamelCase):
       __lowercase	:   BigBirdConfig
       __lowercase	:   jnp.dtype										=      jnp.floataa
       __lowercase	:   bool										=      True


       def    lowerCAmelCase   ( self	)   ->	List[Any]:



            '''simple docstring'''




            super().setup()
            __snake_case            =	nn.Dense(5     ,	dtype=self.dtype	)




       def __call__( self     ,	*__SCREAMING_SNAKE_CASE     ,	**__SCREAMING_SNAKE_CASE	)   ->	str:



            '''simple docstring'''




            __snake_case            =	super().__call__(*__SCREAMING_SNAKE_CASE     ,	**__SCREAMING_SNAKE_CASE	)
            __snake_case            =	self.cls(outputs[2]	)
            return outputs[:2] + (cls_out,)





class      lowerCAmelCase (			__UpperCamelCase):
       __lowercase	:   List[str]										=      FlaxBigBirdForNaturalQuestionsModule







def       _UpperCamelCase			(_lowerCamelCase				:  Tuple ,						_lowerCamelCase				:  Union[str, Any] ,						_lowerCamelCase				:  Tuple ,						_lowerCamelCase				:  Dict ,						_lowerCamelCase				:  Dict ,						_lowerCamelCase				:  int						)-> Optional[Any]:




     '''simple docstring'''





     def cross_entropy(_lowerCamelCase				:  int ,						_lowerCamelCase				:  Optional[Any] ,						_lowerCamelCase				:  Union[str, Any]=None						):
          __snake_case            =	logits.shape[-1]
          __snake_case            =	(labels[..., None] == jnp.arange(__lowerCAmelCase						)[None]).astype('''f4'''						)
          __snake_case            =	jax.nn.log_softmax(__lowerCAmelCase ,						axis=-1						)
          __snake_case            =	-jnp.sum(labels * logits ,						axis=-1						)
          if reduction is not None:
               __snake_case            =	reduction(__lowerCAmelCase						)
          return loss

     __snake_case            =	partial(__lowerCAmelCase ,						reduction=jnp.mean						)
     __snake_case            =	cross_entropy(__lowerCAmelCase ,						__lowerCAmelCase						)
     __snake_case            =	cross_entropy(__lowerCAmelCase ,						__lowerCAmelCase						)
     __snake_case            =	cross_entropy(__lowerCAmelCase ,						__lowerCAmelCase						)
     return (start_loss + end_loss + pooled_loss) / 3





@dataclass
class      lowerCAmelCase :
       __lowercase	:   str										=      "google/bigbird-roberta-base"
       __lowercase	:   int										=      3000
       __lowercase	:   int										=      1_0500

       __lowercase	:   int										=      128
       __lowercase	:   int										=      3

       __lowercase	:   int										=      1
       __lowercase	:   int										=      5

       # tx_args
       __lowercase	:   float										=      3E-5
       __lowercase	:   float										=      0.0
       __lowercase	:   int										=      2_0000
       __lowercase	:   float										=      0.0095

       __lowercase	:   str										=      "bigbird-roberta-natural-questions"
       __lowercase	:   str										=      "training-expt"
       __lowercase	:   str										=      "data/nq-training.jsonl"
       __lowercase	:   str										=      "data/nq-validation.jsonl"


       def    lowerCAmelCase   ( self	)   ->	Tuple:



            '''simple docstring'''




            os.makedirs(self.base_dir     ,	exist_ok=__SCREAMING_SNAKE_CASE	)
            __snake_case            =	os.path.join(self.base_dir     ,	self.save_dir	)
            __snake_case            =	self.batch_size_per_device * jax.device_count()





@dataclass
class      lowerCAmelCase :
       __lowercase	:   int
       __lowercase	:   int										=      4096  # no dynamic padding on TPUs


       def __call__( self     ,	__SCREAMING_SNAKE_CASE	)   ->	Dict:



            '''simple docstring'''




            __snake_case            =	self.collate_fn(__SCREAMING_SNAKE_CASE	)
            __snake_case            =	jax.tree_util.tree_map(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)
            return batch


       def    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE	)   ->	Union[str, Any]:



            '''simple docstring'''




            __snake_case       ,     __snake_case            =	self.fetch_inputs(features['''input_ids''']	)
            __snake_case            =	{
                '''input_ids''': jnp.array(__SCREAMING_SNAKE_CASE     ,	dtype=jnp.intaa	),
                '''attention_mask''': jnp.array(__SCREAMING_SNAKE_CASE     ,	dtype=jnp.intaa	),
                '''start_labels''': jnp.array(features['''start_token''']     ,	dtype=jnp.intaa	),
                '''end_labels''': jnp.array(features['''end_token''']     ,	dtype=jnp.intaa	),
                '''pooled_labels''': jnp.array(features['''category''']     ,	dtype=jnp.intaa	),
            }
            return batch


       def    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE	)   ->	str:



            '''simple docstring'''




            __snake_case            =	[self._fetch_inputs(__SCREAMING_SNAKE_CASE	) for ids in input_ids]
            return zip(*__SCREAMING_SNAKE_CASE	)




       def    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE	)   ->	int:



            '''simple docstring'''




            __snake_case            =	[1 for _ in range(len(__SCREAMING_SNAKE_CASE	)	)]
            while len(__SCREAMING_SNAKE_CASE	) < self.max_length:
                 input_ids.append(self.pad_id	)
                 attention_mask.append(0	)
            return input_ids, attention_mask







def       _UpperCamelCase			(_lowerCamelCase				:  List[Any] ,						_lowerCamelCase				:  List[Any] ,						_lowerCamelCase				:  List[str]=None						)-> Any:




     '''simple docstring'''





     if seed is not None:
          __snake_case            =	dataset.shuffle(seed=__lowerCAmelCase						)
     for i in range(len(__lowerCAmelCase						) // batch_size						):
          __snake_case            =	dataset[i * batch_size : (i + 1) * batch_size]
          yield dict(__lowerCAmelCase						)







@partial(jax.pmap ,						axis_name='''batch'''						)
def       _UpperCamelCase			(_lowerCamelCase				:  Tuple ,						_lowerCamelCase				:  str ,						**_lowerCamelCase				:  List[Any]						)-> str:




     '''simple docstring'''





     def loss_fn(_lowerCamelCase				:  List[Any]						):
          __snake_case            =	model_inputs.pop('''start_labels'''						)
          __snake_case            =	model_inputs.pop('''end_labels'''						)
          __snake_case            =	model_inputs.pop('''pooled_labels'''						)

          __snake_case            =	state.apply_fn(**__lowerCAmelCase ,						params=__lowerCAmelCase ,						dropout_rng=__lowerCAmelCase ,						train=__lowerCAmelCase						)
          __snake_case       ,     __snake_case       ,     __snake_case            =	outputs

          return state.loss_fn(
              __lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						)

     __snake_case       ,     __snake_case            =	jax.random.split(__lowerCAmelCase						)
     __snake_case            =	jax.value_and_grad(__lowerCAmelCase						)
     __snake_case       ,     __snake_case            =	grad_fn(state.params						)
     __snake_case            =	jax.lax.pmean({'''loss''': loss} ,						axis_name='''batch'''						)
     __snake_case            =	jax.lax.pmean(__lowerCAmelCase ,						'''batch'''						)

     __snake_case            =	state.apply_gradients(grads=__lowerCAmelCase						)
     return state, metrics, new_drp_rng







@partial(jax.pmap ,						axis_name='''batch'''						)
def       _UpperCamelCase			(_lowerCamelCase				:  Optional[int] ,						**_lowerCamelCase				:  Tuple						)-> List[Any]:




     '''simple docstring'''





     __snake_case            =	model_inputs.pop('''start_labels'''						)
     __snake_case            =	model_inputs.pop('''end_labels'''						)
     __snake_case            =	model_inputs.pop('''pooled_labels'''						)

     __snake_case            =	state.apply_fn(**__lowerCAmelCase ,						params=state.params ,						train=__lowerCAmelCase						)
     __snake_case       ,     __snake_case       ,     __snake_case            =	outputs

     __snake_case            =	state.loss_fn(__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase						)
     __snake_case            =	jax.lax.pmean({'''loss''': loss} ,						axis_name='''batch'''						)
     return metrics





class      lowerCAmelCase (			train_state.TrainState):
       __lowercase	:   Callable										=      struct.field(pytree_node=__UpperCamelCase)





@dataclass
class      lowerCAmelCase :
       __lowercase	:   Args
       __lowercase	:   Callable
       __lowercase	:   Callable
       __lowercase	:   Callable
       __lowercase	:   Callable
       __lowercase	:   wandb
       __lowercase	:   Callable										=      None


       def    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE=None	)   ->	Optional[Any]:



            '''simple docstring'''




            __snake_case            =	model.params
            __snake_case            =	TrainState.create(
                apply_fn=model.__call__     ,	params=__SCREAMING_SNAKE_CASE     ,	tx=__SCREAMING_SNAKE_CASE     ,	loss_fn=__SCREAMING_SNAKE_CASE     ,	)
            if ckpt_dir is not None:
                 __snake_case       ,     __snake_case       ,     __snake_case       ,     __snake_case       ,     __snake_case            =	restore_checkpoint(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)
                 __snake_case            =	{
                     '''lr''': args.lr,
                     '''init_lr''': args.init_lr,
                     '''warmup_steps''': args.warmup_steps,
                     '''num_train_steps''': num_train_steps,
                     '''weight_decay''': args.weight_decay,
                 }
                 __snake_case       ,     __snake_case            =	build_tx(**__SCREAMING_SNAKE_CASE	)
                 __snake_case            =	train_state.TrainState(
                     step=__SCREAMING_SNAKE_CASE     ,	apply_fn=model.__call__     ,	params=__SCREAMING_SNAKE_CASE     ,	tx=__SCREAMING_SNAKE_CASE     ,	opt_state=__SCREAMING_SNAKE_CASE     ,	)
                 __snake_case            =	args
                 __snake_case            =	data_collator
                 __snake_case            =	lr
                 __snake_case            =	params
            __snake_case            =	jax_utils.replicate(__SCREAMING_SNAKE_CASE	)
            return state


       def    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)   ->	Dict:



            '''simple docstring'''




            __snake_case            =	self.args
            __snake_case            =	len(__SCREAMING_SNAKE_CASE	) // args.batch_size

            __snake_case            =	jax.random.PRNGKey(0	)
            __snake_case            =	jax.random.split(__SCREAMING_SNAKE_CASE     ,	jax.device_count()	)
            for epoch in range(args.max_epochs	):
                 __snake_case            =	jnp.array(0     ,	dtype=jnp.floataa	)
                 __snake_case            =	get_batched_dataset(__SCREAMING_SNAKE_CASE     ,	args.batch_size     ,	seed=__SCREAMING_SNAKE_CASE	)
                 __snake_case            =	0
                 for batch in tqdm(__SCREAMING_SNAKE_CASE     ,	total=__SCREAMING_SNAKE_CASE     ,	desc=F'''Running EPOCH-{epoch}'''	):
                      __snake_case            =	self.data_collator(__SCREAMING_SNAKE_CASE	)
                      __snake_case       ,     __snake_case       ,     __snake_case            =	self.train_step_fn(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	**__SCREAMING_SNAKE_CASE	)
                      running_loss += jax_utils.unreplicate(metrics['''loss''']	)
                      i += 1
                      if i % args.logging_steps == 0:
                           __snake_case            =	jax_utils.unreplicate(state.step	)
                           __snake_case            =	running_loss.item() / i
                           __snake_case            =	self.scheduler_fn(state_step - 1	)

                           __snake_case            =	self.evaluate(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)
                           __snake_case            =	{
                               '''step''': state_step.item(),
                               '''eval_loss''': eval_loss.item(),
                               '''tr_loss''': tr_loss,
                               '''lr''': lr.item(),
                           }
                           tqdm.write(str(__SCREAMING_SNAKE_CASE	)	)
                           self.logger.log(__SCREAMING_SNAKE_CASE     ,	commit=__SCREAMING_SNAKE_CASE	)

                      if i % args.save_steps == 0:
                           self.save_checkpoint(args.save_dir + F'''-e{epoch}-s{i}'''     ,	state=__SCREAMING_SNAKE_CASE	)


       def    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)   ->	int:



            '''simple docstring'''




            __snake_case            =	get_batched_dataset(__SCREAMING_SNAKE_CASE     ,	self.args.batch_size	)
            __snake_case            =	len(__SCREAMING_SNAKE_CASE	) // self.args.batch_size
            __snake_case            =	jnp.array(0     ,	dtype=jnp.floataa	)
            __snake_case            =	0
            for batch in tqdm(__SCREAMING_SNAKE_CASE     ,	total=__SCREAMING_SNAKE_CASE     ,	desc='''Evaluating ... '''	):
                 __snake_case            =	self.data_collator(__SCREAMING_SNAKE_CASE	)
                 __snake_case            =	self.val_step_fn(__SCREAMING_SNAKE_CASE     ,	**__SCREAMING_SNAKE_CASE	)
                 running_loss += jax_utils.unreplicate(metrics['''loss''']	)
                 i += 1
            return running_loss / i




       def    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)   ->	Optional[int]:



            '''simple docstring'''




            __snake_case            =	jax_utils.unreplicate(__SCREAMING_SNAKE_CASE	)
            print(F'''SAVING CHECKPOINT IN {save_dir}'''     ,	end=''' ... '''	)
            self.model_save_fn(__SCREAMING_SNAKE_CASE     ,	params=state.params	)
            with open(os.path.join(__SCREAMING_SNAKE_CASE     ,	'''opt_state.msgpack'''	)     ,	'''wb'''	) as f:
                 f.write(to_bytes(state.opt_state	)	)
            joblib.dump(self.args     ,	os.path.join(__SCREAMING_SNAKE_CASE     ,	'''args.joblib'''	)	)
            joblib.dump(self.data_collator     ,	os.path.join(__SCREAMING_SNAKE_CASE     ,	'''data_collator.joblib'''	)	)
            with open(os.path.join(__SCREAMING_SNAKE_CASE     ,	'''training_state.json'''	)     ,	'''w'''	) as f:
                 json.dump({'''step''': state.step.item()}     ,	__SCREAMING_SNAKE_CASE	)
            print('''DONE'''	)







def       _UpperCamelCase			(_lowerCamelCase				:  List[str] ,						_lowerCamelCase				:  int						)-> Union[str, Any]:




     '''simple docstring'''





     print(f'''RESTORING CHECKPOINT FROM {save_dir}''' ,						end=''' ... '''						)
     with open(os.path.join(__lowerCAmelCase ,						'''flax_model.msgpack'''						) ,						'''rb'''						) as f:
          __snake_case            =	from_bytes(state.params ,						f.read()						)

     with open(os.path.join(__lowerCAmelCase ,						'''opt_state.msgpack'''						) ,						'''rb'''						) as f:
          __snake_case            =	from_bytes(state.opt_state ,						f.read()						)

     __snake_case            =	joblib.load(os.path.join(__lowerCAmelCase ,						'''args.joblib'''						)						)
     __snake_case            =	joblib.load(os.path.join(__lowerCAmelCase ,						'''data_collator.joblib'''						)						)

     with open(os.path.join(__lowerCAmelCase ,						'''training_state.json'''						) ,						'''r'''						) as f:
          __snake_case            =	json.load(__lowerCAmelCase						)
     __snake_case            =	training_state['''step''']

     print('''DONE'''						)
     return params, opt_state, step, args, data_collator







def       _UpperCamelCase			(_lowerCamelCase				:  List[str] ,						_lowerCamelCase				:  Union[str, Any] ,						_lowerCamelCase				:  Any ,						_lowerCamelCase				:  Union[str, Any]						)-> Tuple:




     '''simple docstring'''





     __snake_case            =	num_train_steps - warmup_steps
     __snake_case            =	optax.linear_schedule(init_value=__lowerCAmelCase ,						end_value=__lowerCAmelCase ,						transition_steps=__lowerCAmelCase						)
     __snake_case            =	optax.linear_schedule(init_value=__lowerCAmelCase ,						end_value=1E-7 ,						transition_steps=__lowerCAmelCase						)
     __snake_case            =	optax.join_schedules(schedules=[warmup_fn, decay_fn] ,						boundaries=[warmup_steps]						)
     return lr







def       _UpperCamelCase			(_lowerCamelCase				:  Optional[Any] ,						_lowerCamelCase				:  Optional[Any] ,						_lowerCamelCase				:  Any ,						_lowerCamelCase				:  List[str] ,						_lowerCamelCase				:  Optional[Any]						)-> int:




     '''simple docstring'''





     def weight_decay_mask(_lowerCamelCase				:  int						):
          __snake_case            =	traverse_util.flatten_dict(__lowerCAmelCase						)
          __snake_case            =	{k: (v[-1] != '''bias''' and v[-2:] != ('''LayerNorm''', '''scale''')) for k, v in params.items()}
          return traverse_util.unflatten_dict(__lowerCAmelCase						)

     __snake_case            =	scheduler_fn(__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase ,						__lowerCAmelCase						)

     __snake_case            =	optax.adamw(learning_rate=__lowerCAmelCase ,						weight_decay=__lowerCAmelCase ,						mask=__lowerCAmelCase						)
     return tx, lr








import inspect
import unittest
from typing import List

import numpy as np

from transformers import EfficientFormerConfig
from transformers.testing_utils import require_tf, require_vision, slow
from transformers.utils import cached_property, is_tf_available, is_vision_available

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


if is_tf_available():
     import tensorflow as tf

     from transformers import (
         TFEfficientFormerForImageClassification,
         TFEfficientFormerForImageClassificationWithTeacher,
         TFEfficientFormerModel,
     )
     from transformers.models.efficientformer.modeling_tf_efficientformer import (
         TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
     )


if is_vision_available():
     from PIL import Image

     from transformers import EfficientFormerImageProcessor




class        _SCREAMING_SNAKE_CASE	:
      def __init__(      self		,						lowerCamelCase		,						lowerCamelCase = 13		,						lowerCamelCase = 64		,						lowerCamelCase = 2		,						lowerCamelCase = 3		,						lowerCamelCase = 3		,						lowerCamelCase = True		,						lowerCamelCase = True		,						lowerCamelCase = 1_28		,						lowerCamelCase=[16, 32, 64, 1_28]		,						lowerCamelCase = 7		,						lowerCamelCase = 4		,						lowerCamelCase = 37		,						lowerCamelCase = "gelu"		,						lowerCamelCase = 0.1		,						lowerCamelCase = 0.1		,						lowerCamelCase = 10		,						lowerCamelCase = 0.0_2		,						lowerCamelCase = 2		,						lowerCamelCase = 1		,						lowerCamelCase = 1_28		,						lowerCamelCase = [2, 2, 2, 2]		,						lowerCamelCase = 2		,						lowerCamelCase = 2		,						):
          snake_case__    =					parent
          snake_case__    =					batch_size
          snake_case__    =					image_size
          snake_case__    =					patch_size
          snake_case__    =					num_channels
          snake_case__    =					is_training
          snake_case__    =					use_labels
          snake_case__    =					hidden_size
          snake_case__    =					num_hidden_layers
          snake_case__    =					num_attention_heads
          snake_case__    =					intermediate_size
          snake_case__    =					hidden_act
          snake_case__    =					hidden_dropout_prob
          snake_case__    =					attention_probs_dropout_prob
          snake_case__    =					type_sequence_label_size
          snake_case__    =					initializer_range
          snake_case__    =					encoder_stride
          snake_case__    =					num_attention_outputs
          snake_case__    =					embed_dim
          snake_case__    =					embed_dim + 1
          snake_case__    =					resolution
          snake_case__    =					depths
          snake_case__    =					hidden_sizes
          snake_case__    =					dim
          snake_case__    =					mlp_expansion_ratio
      def A_      (      self			):
          snake_case__    =					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)

          snake_case__    =					None
          if self.use_labels:
              snake_case__    =					ids_tensor([self.batch_size]		,						self.type_sequence_label_size			)

          snake_case__    =					self.get_config()

          return config, pixel_values, labels
      def A_      (      self			):
          return EfficientFormerConfig(
              image_size=self.image_size		,						patch_size=self.patch_size		,						num_channels=self.num_channels		,						hidden_size=self.hidden_size		,						num_hidden_layers=self.num_hidden_layers		,						num_attention_heads=self.num_attention_heads		,						intermediate_size=self.intermediate_size		,						hidden_act=self.hidden_act		,						hidden_dropout_prob=self.hidden_dropout_prob		,						attention_probs_dropout_prob=self.attention_probs_dropout_prob		,						is_decoder=lowerCamelCase		,						initializer_range=self.initializer_range		,						encoder_stride=self.encoder_stride		,						resolution=self.resolution		,						depths=self.depths		,						hidden_sizes=self.hidden_sizes		,						dim=self.dim		,						mlp_expansion_ratio=self.mlp_expansion_ratio		,						)
      def A_      (      self		,						lowerCamelCase		,						lowerCamelCase		,						lowerCamelCase			):
          snake_case__    =					TFEfficientFormerModel(config=lowerCamelCase			)
          snake_case__    =					model(lowerCamelCase		,						training=lowerCamelCase			)
          self.parent.assertEqual(result.last_hidden_state.shape		,						(self.batch_size, self.seq_length, self.hidden_size)			)
      def A_      (      self		,						lowerCamelCase		,						lowerCamelCase		,						lowerCamelCase			):
          snake_case__    =					self.type_sequence_label_size
          snake_case__    =					TFEfficientFormerForImageClassification(lowerCamelCase			)
          snake_case__    =					model(lowerCamelCase		,						labels=lowerCamelCase		,						training=lowerCamelCase			)
          self.parent.assertEqual(result.logits.shape		,						(self.batch_size, self.type_sequence_label_size)			)

          # test greyscale images
          snake_case__    =					1
          snake_case__    =					TFEfficientFormerForImageClassification(lowerCamelCase			)

          snake_case__    =					floats_tensor([self.batch_size, 1, self.image_size, self.image_size]			)
          snake_case__    =					model(lowerCamelCase		,						labels=lowerCamelCase			)
          self.parent.assertEqual(result.logits.shape		,						(self.batch_size, self.type_sequence_label_size)			)



      def A_      (      self			):
          snake_case__    =					self.prepare_config_and_inputs()
          snake_case__    ,       snake_case__    ,       snake_case__    =					config_and_inputs
          snake_case__    =					{"pixel_values": pixel_values}
          return config, inputs_dict





@require_tf
class        _SCREAMING_SNAKE_CASE	(   __UpperCamelCase							,       __UpperCamelCase							,       unittest.TestCase			):
      _A    :		str     						=  (
          (
              TFEfficientFormerModel,
              TFEfficientFormerForImageClassificationWithTeacher,
              TFEfficientFormerForImageClassification,
          )
          if is_tf_available()
          else ()
      )
      _A    :		List[str]     						=  (
          {
              'feature-extraction': TFEfficientFormerModel,
              'image-classification': (
                  TFEfficientFormerForImageClassification,
                  TFEfficientFormerForImageClassificationWithTeacher,
              ),
          }
          if is_tf_available()
          else {}
      )

      _A    :		Optional[Any]     						=  False

      _A    :		List[Any]     						=  False
      _A    :		Tuple     						=  False
      _A    :		List[Any]     						=  False
      _A    :		Any     						=  False
      def A_      (      self			):
          snake_case__    =					TFEfficientFormerModelTester(self			)
          snake_case__    =					ConfigTester(
              self		,						config_class=lowerCamelCase		,						has_text_modality=lowerCamelCase		,						hidden_size=37			)
      def A_      (      self			):
          self.config_tester.run_common_tests()
      @unittest.skip(reason="EfficientFormer does not use inputs_embeds"			)
      def A_      (      self			):
          pass
      @unittest.skip(reason="EfficientFormer does not support input and output embeddings"			)
      def A_      (      self			):
          pass
      def A_      (      self			):
          snake_case__    ,       snake_case__    =					self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
              snake_case__    =					model_class(lowerCamelCase			)
              snake_case__    =					inspect.signature(model.call			)
              # signature.parameters is an OrderedDict => so arg_names order is deterministic
              snake_case__    =					[*signature.parameters.keys()]

              snake_case__    =					["pixel_values"]
              self.assertListEqual(arg_names[:1]		,						lowerCamelCase			)
      def A_      (      self			):
          def check_hidden_states_output(lowerCamelCase		,						lowerCamelCase		,						lowerCamelCase			):
              snake_case__    =					model_class(lowerCamelCase			)

              snake_case__    =					model(**self._prepare_for_class(lowerCamelCase		,						lowerCamelCase			)		,						training=lowerCamelCase			)
              snake_case__    =					outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

              snake_case__    =					getattr(
                  self.model_tester		,						"expected_num_hidden_layers"		,						self.model_tester.num_hidden_layers + 1			)
              self.assertEqual(len(lowerCamelCase			)		,						lowerCamelCase			)

              if hasattr(self.model_tester		,						"encoder_seq_length"			):
                  snake_case__    =					self.model_tester.encoder_seq_length
                  if hasattr(self.model_tester		,						"chunk_length"			) and self.model_tester.chunk_length > 1:
                      snake_case__    =					seq_length * self.model_tester.chunk_length
              else:
                  snake_case__    =					self.model_tester.seq_length

              self.assertListEqual(
                  list(hidden_states[-1].shape[-2:]			)		,						[seq_length, self.model_tester.hidden_size]		,						)

              if config.is_encoder_decoder:
                  snake_case__    =					outputs.decoder_hidden_states

                  self.asseretIsInstance(lowerCamelCase		,						(list, tuple)			)
                  self.assertEqual(len(lowerCamelCase			)		,						lowerCamelCase			)
                  snake_case__    =					getattr(self.model_tester		,						"seq_length"		,						lowerCamelCase			)
                  snake_case__    =					getattr(self.model_tester		,						"decoder_seq_length"		,						lowerCamelCase			)

                  self.assertListEqual(
                      list(hidden_states[-1].shape[-2:]			)		,						[decoder_seq_length, self.model_tester.hidden_size]		,						)

          snake_case__    ,       snake_case__    =					self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
              snake_case__    =					True
              check_hidden_states_output(lowerCamelCase		,						lowerCamelCase		,						lowerCamelCase			)

              # check that output_hidden_states also work using config
              del inputs_dict["output_hidden_states"]
              snake_case__    =					True

              check_hidden_states_output(lowerCamelCase		,						lowerCamelCase		,						lowerCamelCase			)
      def A_      (      self		,						lowerCamelCase		,						lowerCamelCase		,						lowerCamelCase=False			):
          snake_case__    =					super()._prepare_for_class(lowerCamelCase		,						lowerCamelCase		,						return_labels=lowerCamelCase			)

          if return_labels:
              if model_class.__name__ == "TFEfficientFormerForImageClassificationWithTeacher":
                  del inputs_dict["labels"]

          return inputs_dict
      def A_      (      self			):
          snake_case__    =					self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_model(*lowerCamelCase			)
      @unittest.skip(reason="EfficientFormer does not implement masked image modeling yet"			)
      def A_      (      self			):
          snake_case__    =					self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_masked_image_modeling(*lowerCamelCase			)
      def A_      (      self			):
          snake_case__    =					self.model_tester.prepare_config_and_inputs()
          self.model_tester.create_and_check_for_image_classification(*lowerCamelCase			)
      @slow
      def A_      (      self			):
          for model_name in TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
              snake_case__    =					TFEfficientFormerModel.from_pretrained(lowerCamelCase			)
              self.assertIsNotNone(lowerCamelCase			)
      def A_      (      self			):
          snake_case__    ,       snake_case__    =					self.model_tester.prepare_config_and_inputs_for_common()
          snake_case__    =					True

          snake_case__    =					getattr(self.model_tester		,						"seq_length"		,						lowerCamelCase			)
          snake_case__    =					getattr(self.model_tester		,						"encoder_seq_length"		,						lowerCamelCase			)
          snake_case__    =					getattr(self.model_tester		,						"key_length"		,						lowerCamelCase			)
          snake_case__    =					getattr(self.model_tester		,						"chunk_length"		,						lowerCamelCase			)

          if chunk_length is not None and hasattr(self.model_tester		,						"num_hashes"			):
              snake_case__    =					encoder_seq_length * self.model_tester.num_hashes

          for model_class in self.all_model_classes:
              snake_case__    =					True
              snake_case__    =					False
              snake_case__    =					True
              snake_case__    =					model_class(lowerCamelCase			)

              snake_case__    =					model(**self._prepare_for_class(lowerCamelCase		,						lowerCamelCase			)		,						training=lowerCamelCase			)
              snake_case__    =					outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
              self.assertEqual(len(lowerCamelCase			)		,						self.model_tester.num_attention_outputs			)

              # check that output_attentions also work using config
              del inputs_dict["output_attentions"]
              snake_case__    =					True
              snake_case__    =					model_class(lowerCamelCase			)
              snake_case__    =					model(**self._prepare_for_class(lowerCamelCase		,						lowerCamelCase			)		,						training=lowerCamelCase			)

              snake_case__    =					outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
              self.assertEqual(len(lowerCamelCase			)		,						self.model_tester.num_attention_outputs			)

              if chunk_length is not None:
                  self.assertListEqual(
                      list(attentions[0].shape[-4:]			)		,						[self.model_tester.num_attention_heads, encoder_seq_length, chunk_length, encoder_key_length]		,						)
              else:
                  self.assertListEqual(
                      list(attentions[0].shape[-3:]			)		,						[self.model_tester.num_attention_heads, encoder_seq_length, encoder_key_length]		,						)



      def A_      (      self			):
          # We use a simplified version of this test for EfficientFormer because it requires training=False
          # and Keras refuses to let us force that during functional construction
          snake_case__    ,       snake_case__    =					self.model_tester.prepare_config_and_inputs_for_common()

          for model_class in self.all_model_classes:
              # Prepare our model
              snake_case__    =					model_class(lowerCamelCase			)
              # These are maximally general inputs for the model, with multiple None dimensions
              # Hopefully this will catch any conditionals that fail for flexible shapes
              snake_case__    =					{
                  key: tf.keras.Input(shape=val.shape[1:]		,						dtype=val.dtype		,						name=lowerCamelCase			)
                  for key, val in model.input_signature.items()
                  if key in model.dummy_inputs
              }
              snake_case__    =					model(lowerCamelCase			)
              self.assertTrue(outputs_dict is not None			)





def 						SCREAMING_SNAKE_CASE__     (   ):
    snake_case__    =					Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"				)
    return image





@require_tf
@require_vision
class        _SCREAMING_SNAKE_CASE	(   unittest.TestCase			):
      @cached_property
      def A_      (      self			):
          return (
              EfficientFormerImageProcessor.from_pretrained("snap-research/efficientformer-l1-300"			)
              if is_vision_available()
              else None
          )
      @slow
      def A_      (      self			):
          snake_case__    =					TFEfficientFormerForImageClassification.from_pretrained("snap-research/efficientformer-l1-300"			)
          snake_case__    =					self.default_image_processor
          snake_case__    =					prepare_img()
          snake_case__    =					image_processor(images=lowerCamelCase		,						return_tensors="tf"			)
          # forward pass
          snake_case__    =					model(**lowerCamelCase		,						training=lowerCamelCase			)
          # verify the logits
          snake_case__    =					tf.TensorShape((1, 10_00)			)
          self.assertEqual(outputs.logits.shape		,						lowerCamelCase			)
          snake_case__    =					tf.constant([-0.0_5_5_5, 0.4_8_2_5, -0.0_8_5_2]			)
          self.assertTrue(np.allclose(outputs.logits[0, :3]		,						lowerCamelCase		,						atol=1e-4			)			)



      @slow
      def A_      (      self			):
          snake_case__    =					TFEfficientFormerForImageClassificationWithTeacher.from_pretrained(
              "snap-research/efficientformer-l1-300"			)
          snake_case__    =					self.default_image_processor
          snake_case__    =					prepare_img()
          snake_case__    =					image_processor(images=lowerCamelCase		,						return_tensors="tf"			)
          # forward pass
          snake_case__    =					model(**lowerCamelCase		,						training=lowerCamelCase			)
          # verify the logits
          snake_case__    =					tf.TensorShape((1, 10_00)			)
          self.assertEqual(outputs.logits.shape		,						lowerCamelCase			)
          snake_case__    =					tf.constant([-0.1_3_1_2, 0.4_3_5_3, -1.0_4_9_9]			)
          self.assertTrue(np.allclose(outputs.logits[0, :3]		,						lowerCamelCase		,						atol=1e-4			)			)



