

"""simple docstring"""





import argparse
import logging
import os
from datetime import datetime

import numpy as np
import torch
from torch import nn
from torch.utils.data import DataLoader, RandomSampler, TensorDataset
from tqdm import tqdm

from transformers import GPTaLMHeadModel


__UpperCamelCase  :      int							      =		logging.getLogger(__name__)

def   __SCREAMING_SNAKE_CASE  (      A_						,     A_ ):
	# save results
	if os.path.exists(A_ ):
		if os.path.exists(os.path.join(A_						,     '''config.json''' ) ) and os.path.isfile(
		    os.path.join(A_						,     '''config.json''' ) ):
			os.remove(os.path.join(A_						,     '''config.json''' ) )
		if os.path.exists(os.path.join(A_						,     '''pytorch_model.bin''' ) ) and os.path.isfile(
		    os.path.join(A_						,     '''pytorch_model.bin''' ) ):
			os.remove(os.path.join(A_						,     '''pytorch_model.bin''' ) )
	else:
		os.makedirs(A_ )
	model.save_pretrained(A_ )

def   __SCREAMING_SNAKE_CASE  (      A_						,     A_=False ):
	lowerCAmelCase__	:						Optional[Any]  				=	2
	if unlogit:
		lowerCAmelCase__	:						Union[str, Any]  				=	torch.pow(A_						,     A_ )
	lowerCAmelCase__	:						Optional[Any]  				=	p * torch.log(A_ )
	lowerCAmelCase__	:						List[Any]  				=	0
	return -plogp.sum(dim=-1 )

def   __SCREAMING_SNAKE_CASE  (      A_ ):
	logger.info('''lv, h >\t''' + '''\t'''.join(f'{x + 1}' for x in range(len(A_ ) ) ) )
	for row in range(len(A_ ) ):
		if tensor.dtype != torch.long:
			logger.info(f'layer {row + 1}:\t' + '''\t'''.join(f'{x:.5f}' for x in tensor[row].cpu().data ) )
		else:
			logger.info(f'layer {row + 1}:\t' + '''\t'''.join(f'{x:d}' for x in tensor[row].cpu().data ) )

def   __SCREAMING_SNAKE_CASE  (      A_						,     A_						,     A_						,     A_=True						,     A_=True						,     A_=None						,     A_=False ):
	lowerCAmelCase__				,lowerCAmelCase__	:						List[Any]  				=	model.config.num_hidden_layers, model.config.num_attention_heads
	lowerCAmelCase__	:						Dict  				=	torch.zeros(A_						,     A_ ).to(args.device )
	lowerCAmelCase__	:						int  				=	torch.zeros(A_						,     A_ ).to(args.device )

	if head_mask is None:
		lowerCAmelCase__	:						Union[str, Any]  				=	torch.ones(A_						,     A_ ).to(args.device )

	head_mask.requires_grad_(requires_grad=A_ )
	# If actually pruned attention multi-head, set head mask to None to avoid shape mismatch
	if actually_pruned:
		lowerCAmelCase__	:						Union[str, Any]  				=	None

	lowerCAmelCase__	:						Optional[int]  				=	0.0
	lowerCAmelCase__	:						Optional[int]  				=	0.0
	for step, inputs in enumerate(tqdm(A_						,     desc='''Iteration'''						,     disable=args.local_rank not in [-1, 0] ) ):
		lowerCAmelCase__	:						Any  				=	tuple(t.to(args.device ) for t in inputs )
		((lowerCAmelCase__)				,)	:						List[Any]  				=	inputs

		# Do a forward pass (not with torch.no_grad() since we need gradients for importance score - see below)
		lowerCAmelCase__	:						Any  				=	model(A_						,     labels=A_						,     head_mask=A_ )
		#  (loss), lm_logits, presents, (all hidden_states), (attentions)
		lowerCAmelCase__				,lowerCAmelCase__				,lowerCAmelCase__	:						Dict  				=	(
		    outputs[0],
		    outputs[1],
		    outputs[-1],
		)  # Loss and logits are the first, attention the last
		loss.backward()  # Backpropagate to populate the gradients in the head mask
		total_loss += loss.detach().cpu().numpy()
		if compute_entropy:
			for layer, attn in enumerate(A_ ):
				lowerCAmelCase__	:						Dict  				=	entropy(attn.detach()						,     A_ )
				attn_entropy[layer] += masked_entropy.sum(-1 ).sum(0 ).sum(0 ).detach()

		if compute_importance:
			head_importance += head_mask.grad.abs().detach()
		tot_tokens += torch.ones_like(A_ ).float().detach().sum().data

	# Normalize
	attn_entropy /= tot_tokens
	head_importance /= tot_tokens
	# Layerwise importance normalization
	if not args.dont_normalize_importance_by_layer:
		lowerCAmelCase__	:						Any  				=	2
		lowerCAmelCase__	:						Dict  				=	torch.pow(torch.pow(A_						,     A_ ).sum(-1 )						,     1 / exponent )
		head_importance /= norm_by_layer.unsqueeze(-1 ) + 1e-2_0

	if not args.dont_normalize_global_importance:
		lowerCAmelCase__	:						List[Any]  				=	(head_importance - head_importance.min()) / (head_importance.max() - head_importance.min())

	# Print matrices
	if compute_entropy:
		logger.info('''Attention entropies''' )
		print_ad_tensor(A_ )
	if compute_importance:
		logger.info('''Head importance scores''' )
		print_ad_tensor(A_ )
	logger.info('''Head ranked by importance scores''' )
	lowerCAmelCase__	:						str  				=	torch.zeros(head_importance.numel()						,     dtype=torch.long						,     device=args.device )
	lowerCAmelCase__	:						Optional[int]  				=	torch.arange(
	    head_importance.numel()						,     device=args.device )
	lowerCAmelCase__	:						int  				=	head_ranks.view_as(A_ )
	print_ad_tensor(A_ )
	return attn_entropy, head_importance, total_loss

def   __SCREAMING_SNAKE_CASE  (      A_						,     A_						,     A_ ):
	lowerCAmelCase__				,lowerCAmelCase__				,lowerCAmelCase__	:						List[str]  				=	compute_heads_importance(A_						,     A_						,     A_						,     compute_entropy=A_ )
	lowerCAmelCase__	:						Union[str, Any]  				=	1 / loss  # instead of downsteam score use the LM loss
	logger.info('''Pruning: original score: %f, threshold: %f'''						,     A_						,     original_score * args.masking_threshold )

	lowerCAmelCase__	:						Union[str, Any]  				=	torch.ones_like(A_ )
	lowerCAmelCase__	:						List[str]  				=	max(1						,     int(new_head_mask.numel() * args.masking_amount ) )

	lowerCAmelCase__	:						int  				=	original_score
	while current_score >= original_score * args.masking_threshold:
		lowerCAmelCase__	:						Union[str, Any]  				=	new_head_mask.clone().detach()  # save current head mask
		# heads from least important to most - keep only not-masked heads
		lowerCAmelCase__	:						str  				=	float('''Inf''' )
		lowerCAmelCase__	:						List[Any]  				=	head_importance.view(-1 ).sort()[1]

		if len(A_ ) <= num_to_mask:
			print('''BREAK BY num_to_mask''' )
			break

		# mask heads
		lowerCAmelCase__	:						List[Any]  				=	current_heads_to_mask[:num_to_mask]
		logger.info('''Heads to mask: %s'''						,     str(current_heads_to_mask.tolist() ) )
		lowerCAmelCase__	:						int  				=	new_head_mask.view(-1 )
		lowerCAmelCase__	:						Optional[int]  				=	0.0
		lowerCAmelCase__	:						Union[str, Any]  				=	new_head_mask.view_as(A_ )
		lowerCAmelCase__	:						Tuple  				=	new_head_mask.clone().detach()
		print_ad_tensor(A_ )

		# Compute metric and head importance again
		lowerCAmelCase__				,lowerCAmelCase__				,lowerCAmelCase__	:						List[Any]  				=	compute_heads_importance(
		    A_						,     A_						,     A_						,     compute_entropy=A_						,     head_mask=A_ )
		lowerCAmelCase__	:						Tuple  				=	1 / loss
		logger.info(
		    '''Masking: current score: %f, remaining heads %d (%.1f percents)'''						,     A_						,     new_head_mask.sum()						,     new_head_mask.sum() / new_head_mask.numel() * 1_00						,     )

	logger.info('''Final head mask''' )
	print_ad_tensor(A_ )
	np.save(os.path.join(args.output_dir						,     '''head_mask.npy''' )						,     head_mask.detach().cpu().numpy() )

	return head_mask

def   __SCREAMING_SNAKE_CASE  (      A_						,     A_						,     A_						,     A_ ):
	lowerCAmelCase__	:						Optional[Any]  				=	datetime.now()
	lowerCAmelCase__				,lowerCAmelCase__				,lowerCAmelCase__	:						Optional[Any]  				=	compute_heads_importance(
	    A_						,     A_						,     A_						,     compute_entropy=A_						,     compute_importance=A_						,     head_mask=A_ )
	lowerCAmelCase__	:						Optional[Any]  				=	1 / loss
	lowerCAmelCase__	:						Tuple  				=	datetime.now() - before_time

	lowerCAmelCase__	:						int  				=	sum(p.numel() for p in model.parameters() )
	lowerCAmelCase__	:						List[Any]  				=	{
	    layer: (1 - head_mask[layer].long()).nonzero().squeeze().tolist() for layer in range(len(A_ ) )
	}

	for k, v in heads_to_prune.items():
		if isinstance(A_						,     A_ ):
			lowerCAmelCase__	:						int  				=	[
			    v,
			]

	assert sum(len(A_ ) for h in heads_to_prune.values() ) == (1 - head_mask.long()).sum().item()
	model.prune_heads(A_ )
	lowerCAmelCase__	:						List[Any]  				=	sum(p.numel() for p in model.parameters() )

	lowerCAmelCase__	:						Any  				=	datetime.now()
	lowerCAmelCase__				,lowerCAmelCase__				,lowerCAmelCase__	:						int  				=	compute_heads_importance(
	    A_						,     A_						,     A_						,     compute_entropy=A_						,     compute_importance=A_						,     head_mask=A_						,     actually_pruned=A_						,     )

	lowerCAmelCase__	:						int  				=	1 / loss
	lowerCAmelCase__	:						Dict  				=	datetime.now() - before_time

	logger.info(
	    '''Pruning: original num of params: %.2e, after pruning %.2e (%.1f percents)'''						,     A_						,     A_						,     pruned_num_params / original_num_params * 1_00						,     )
	logger.info('''Pruning: score with masking: %f score with pruning: %f'''						,     A_						,     A_ )
	logger.info('''Pruning: speed ratio (original timing / new timing): %f percents'''						,     original_time / new_time * 1_00 )
	save_model(A_						,     args.output_dir )

def   __SCREAMING_SNAKE_CASE  (      ):
	lowerCAmelCase__	:						List[str]  				=	argparse.ArgumentParser()
	# Required parameters
	parser.add_argument(
	    '''--data_dir'''						,     default=A_						,     type=A_						,     required=A_						,     help='''The input data dir. Should contain the .tsv files (or other data files) for the task.'''						,     )
	parser.add_argument(
	    '''--model_name_or_path'''						,     default=A_						,     type=A_						,     required=A_						,     help='''Path to pretrained model or model identifier from huggingface.co/models'''						,     )
	parser.add_argument(
	    '''--output_dir'''						,     default=A_						,     type=A_						,     required=A_						,     help='''The output directory where the model predictions and checkpoints will be written.'''						,     )

	# Other parameters
	parser.add_argument(
	    '''--config_name'''						,     default=''''''						,     type=A_						,     help='''Pretrained config name or path if not the same as model_name_or_path'''						,     )
	parser.add_argument(
	    '''--tokenizer_name'''						,     default=''''''						,     type=A_						,     help='''Pretrained tokenizer name or path if not the same as model_name_or_path'''						,     )
	parser.add_argument(
	    '''--cache_dir'''						,     default=A_						,     type=A_						,     help='''Where do you want to store the pre-trained models downloaded from s3'''						,     )
	parser.add_argument(
	    '''--data_subset'''						,     type=A_						,     default=-1						,     help='''If > 0: limit the data to a subset of data_subset instances.''' )
	parser.add_argument(
	    '''--overwrite_output_dir'''						,     action='''store_true'''						,     help='''Whether to overwrite data in output directory''' )
	parser.add_argument(
	    '''--overwrite_cache'''						,     action='''store_true'''						,     help='''Overwrite the cached training and evaluation sets''' )

	parser.add_argument(
	    '''--dont_normalize_importance_by_layer'''						,     action='''store_true'''						,     help='''Don\'t normalize importance score by layers''' )
	parser.add_argument(
	    '''--dont_normalize_global_importance'''						,     action='''store_true'''						,     help='''Don\'t normalize all importance scores between 0 and 1'''						,     )

	parser.add_argument(
	    '''--try_masking'''						,     action='''store_true'''						,     help='''Whether to try to mask head until a threshold of accuracy.''' )
	parser.add_argument(
	    '''--masking_threshold'''						,     default=0.9						,     type=A_						,     help='''masking threshold in term of metrics (stop masking when metric < threshold * original metric value).'''						,     )
	parser.add_argument(
	    '''--masking_amount'''						,     default=0.1						,     type=A_						,     help='''Amount to heads to masking at each masking step.''' )
	parser.add_argument('''--metric_name'''						,     default='''acc'''						,     type=A_						,     help='''Metric to use for head masking.''' )

	parser.add_argument(
	    '''--max_seq_length'''						,     default=1_28						,     type=A_						,     help=(
	        '''The maximum total input sequence length after WordPiece tokenization. \n'''
	        '''Sequences longer than this will be truncated, sequences shorter padded.'''
	    )						,     )
	parser.add_argument('''--batch_size'''						,     default=1						,     type=A_						,     help='''Batch size.''' )

	parser.add_argument('''--seed'''						,     type=A_						,     default=42 )
	parser.add_argument('''--local_rank'''						,     type=A_						,     default=-1						,     help='''local_rank for distributed training on gpus''' )
	parser.add_argument('''--no_cuda'''						,     action='''store_true'''						,     help='''Whether not to use CUDA when available''' )
	parser.add_argument('''--server_ip'''						,     type=A_						,     default=''''''						,     help='''Can be used for distant debugging.''' )
	parser.add_argument('''--server_port'''						,     type=A_						,     default=''''''						,     help='''Can be used for distant debugging.''' )
	lowerCAmelCase__	:						Optional[Any]  				=	parser.parse_args()

	if args.server_ip and args.server_port:
		# Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
		import ptvsd

		print('''Waiting for debugger attach''' )
		ptvsd.enable_attach(address=(args.server_ip, args.server_port)						,     redirect_output=A_ )
		ptvsd.wait_for_attach()

	# Setup devices and distributed training
	if args.local_rank == -1 or args.no_cuda:
		lowerCAmelCase__	:						Union[str, Any]  				=	torch.device('''cuda''' if torch.cuda.is_available() and not args.no_cuda else '''cpu''' )
		lowerCAmelCase__	:						str  				=	0 if args.no_cuda else torch.cuda.device_count()
	else:
		torch.cuda.set_device(args.local_rank )
		lowerCAmelCase__	:						Dict  				=	torch.device('''cuda'''						,     args.local_rank )
		lowerCAmelCase__	:						Union[str, Any]  				=	1
		torch.distributed.init_process_group(backend='''nccl''' )  # Initializes the distributed backend

	# Setup logging
	logging.basicConfig(level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN )
	logger.info('''device: {} n_gpu: {}, distributed: {}'''.format(args.device						,     args.n_gpu						,     bool(args.local_rank != -1 ) ) )

	lowerCAmelCase__	:						List[str]  				=	GPTaLMHeadModel.from_pretrained(args.model_name_or_path )

	# Distributed and parallel training
	model.to(args.device )
	if args.local_rank != -1:
		lowerCAmelCase__	:						Dict  				=	nn.parallel.DistributedDataParallel(
		    A_						,     device_ids=[args.local_rank]						,     output_device=args.local_rank						,     find_unused_parameters=A_ )
	elif args.n_gpu > 1:
		lowerCAmelCase__	:						List[Any]  				=	nn.DataParallel(A_ )

	# Print/save training arguments
	os.makedirs(args.output_dir						,     exist_ok=A_ )
	torch.save(A_						,     os.path.join(args.output_dir						,     '''run_args.bin''' ) )
	logger.info('''Training/evaluation parameters %s'''						,     A_ )

	# Prepare dataset
	lowerCAmelCase__	:						str  				=	np.concatenate(
	    [
	        np.loadtxt(args.data_dir						,     dtype=np.intaa ),
	    ] )
	lowerCAmelCase__	:						Union[str, Any]  				=	(torch.from_numpy(A_ ),)
	lowerCAmelCase__	:						Tuple  				=	TensorDataset(*A_ )
	lowerCAmelCase__	:						Optional[int]  				=	RandomSampler(A_ )
	lowerCAmelCase__	:						Dict  				=	DataLoader(A_						,     sampler=A_						,     batch_size=args.batch_size )

	# Compute head entropy and importance score
	compute_heads_importance(A_						,     A_						,     A_ )

	# Try head masking (set heads to zero until the score goes under a threshole)
	# and head pruning (remove masked heads and see the effect on the network)
	if args.try_masking and args.masking_threshold > 0.0 and args.masking_threshold < 1.0:
		lowerCAmelCase__	:						Tuple  				=	mask_heads(A_						,     A_						,     A_ )
		prune_heads(A_						,     A_						,     A_						,     A_ )


if __name__ == "__main__":
							main()










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 				UpperCamelCase_    :

    '''simple docstring'''






    def __init__(			self      :	Optional[int]					, UpperCAmelCase__      :	List[Any]					, UpperCAmelCase__      :	int = 13					, UpperCAmelCase__      :	int = 64					, UpperCAmelCase__      :	int = 2					, UpperCAmelCase__      :	int = 3					, UpperCAmelCase__      :	int = 3					, UpperCAmelCase__      :	bool = True					, UpperCAmelCase__      :	bool = True					, UpperCAmelCase__      :	int = 128					, UpperCAmelCase__      :	Optional[Any]=[16, 32, 64, 128]					, UpperCAmelCase__      :	int = 7					, UpperCAmelCase__      :	int = 4					, UpperCAmelCase__      :	int = 37					, UpperCAmelCase__      :	str = "gelu"					, UpperCAmelCase__      :	float = 0.1					, UpperCAmelCase__      :	float = 0.1					, UpperCAmelCase__      :	int = 10					, UpperCAmelCase__      :	float = 0.02					, UpperCAmelCase__      :	int = 2					, UpperCAmelCase__      :	int = 1					, UpperCAmelCase__      :	int = 128					, UpperCAmelCase__      :	List[int] = [2, 2, 2, 2]					, UpperCAmelCase__      :	int = 2					, UpperCAmelCase__      :	int = 2					, )						->List[Any]:





           '''simple docstring'''

           A__       =    parent
           A__       =    batch_size
           A__       =    image_size
           A__       =    patch_size
           A__       =    num_channels
           A__       =    is_training
           A__       =    use_labels
           A__       =    hidden_size
           A__       =    num_hidden_layers
           A__       =    num_attention_heads
           A__       =    intermediate_size
           A__       =    hidden_act
           A__       =    hidden_dropout_prob
           A__       =    attention_probs_dropout_prob
           A__       =    type_sequence_label_size
           A__       =    initializer_range
           A__       =    encoder_stride
           A__       =    num_attention_outputs
           A__       =    embed_dim
           A__       =    embed_dim + 1
           A__       =    resolution
           A__       =    depths
           A__       =    hidden_sizes
           A__       =    dim
           A__       =    mlp_expansion_ratio


    def        SCREAMING_SNAKE_CASE							(			self      :	List[Any])						->str:





           '''simple docstring'''

           A__       =    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

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

           A__       =    self.get_config()

           return config, pixel_values, labels


    def        SCREAMING_SNAKE_CASE							(			self      :	int)						->str:





           '''simple docstring'''

           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=UpperCAmelCase__					, 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        SCREAMING_SNAKE_CASE							(			self      :	Optional[int]					, UpperCAmelCase__      :	Dict					, UpperCAmelCase__      :	List[str]					, UpperCAmelCase__      :	Dict)						->Dict:





           '''simple docstring'''

           A__       =    TFEfficientFormerModel(config=UpperCAmelCase__)
           A__       =    model(UpperCAmelCase__					, training=UpperCAmelCase__)
           self.parent.assertEqual(result.last_hidden_state.shape					, (self.batch_size, self.seq_length, self.hidden_size))


    def        SCREAMING_SNAKE_CASE							(			self      :	int					, UpperCAmelCase__      :	Dict					, UpperCAmelCase__      :	Dict					, UpperCAmelCase__      :	str)						->Union[str, Any]:





           '''simple docstring'''

           A__       =    self.type_sequence_label_size
           A__       =    TFEfficientFormerForImageClassification(UpperCAmelCase__)
           A__       =    model(UpperCAmelCase__					, labels=UpperCAmelCase__					, training=UpperCAmelCase__)
           self.parent.assertEqual(result.logits.shape					, (self.batch_size, self.type_sequence_label_size))

           # test greyscale images
           A__       =    1
           A__       =    TFEfficientFormerForImageClassification(UpperCAmelCase__)

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





    def        SCREAMING_SNAKE_CASE							(			self      :	int)						->List[str]:





           '''simple docstring'''

           A__       =    self.prepare_config_and_inputs()
           A__ ,							A__ ,							A__       =    config_and_inputs
           A__       =    {'''pixel_values''': pixel_values}
           return config, inputs_dict


@require_tf
class 				UpperCamelCase_    ( UpperCAmelCase__		,					UpperCAmelCase__		,					unittest.TestCase     ):

    '''simple docstring'''






    UpperCAmelCase__							  =      (
        (
            TFEfficientFormerModel,
            TFEfficientFormerForImageClassificationWithTeacher,
            TFEfficientFormerForImageClassification,
        )
        if is_tf_available()
        else ()
    )
    UpperCAmelCase__							  =      (
        {
            '''feature-extraction''': TFEfficientFormerModel,
            '''image-classification''': (
                TFEfficientFormerForImageClassification,
                TFEfficientFormerForImageClassificationWithTeacher,
            ),
        }
        if is_tf_available()
        else {}
    )

    UpperCAmelCase__							  =      False

    UpperCAmelCase__							  =      False
    UpperCAmelCase__							  =      False
    UpperCAmelCase__							  =      False
    UpperCAmelCase__							  =      False


    def        SCREAMING_SNAKE_CASE							(			self      :	Optional[Any])						->List[str]:





           '''simple docstring'''

           A__       =    TFEfficientFormerModelTester(self)
           A__       =    ConfigTester(
               self					, config_class=UpperCAmelCase__					, has_text_modality=UpperCAmelCase__					, hidden_size=37)


    def        SCREAMING_SNAKE_CASE							(			self      :	int)						->Any:





           '''simple docstring'''

           self.config_tester.run_common_tests()


    @unittest.skip(reason='''EfficientFormer does not use inputs_embeds''')
    def        SCREAMING_SNAKE_CASE							(			self      :	List[str])						->Dict:





           '''simple docstring'''

           pass


    @unittest.skip(reason='''EfficientFormer does not support input and output embeddings''')
    def        SCREAMING_SNAKE_CASE							(			self      :	List[Any])						->Optional[Any]:





           '''simple docstring'''

           pass


    def        SCREAMING_SNAKE_CASE							(			self      :	Any)						->Optional[Any]:





           '''simple docstring'''

           A__ ,							A__       =    self.model_tester.prepare_config_and_inputs_for_common()

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

                  A__       =    ['''pixel_values''']
                  self.assertListEqual(arg_names[:1]					, UpperCAmelCase__)


    def        SCREAMING_SNAKE_CASE							(			self      :	str)						->Any:





           '''simple docstring'''

           def check_hidden_states_output(UpperCAmelCase__      :	Optional[Any]					, UpperCAmelCase__      :	Any					, UpperCAmelCase__      :	Dict):
                  A__       =    model_class(UpperCAmelCase__)

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

                  A__       =    getattr(
                      self.model_tester					, '''expected_num_hidden_layers'''					, self.model_tester.num_hidden_layers + 1)
                  self.assertEqual(len(UpperCAmelCase__)					, UpperCAmelCase__)

                  if hasattr(self.model_tester					, '''encoder_seq_length'''):
                         A__       =    self.model_tester.encoder_seq_length
                         if hasattr(self.model_tester					, '''chunk_length''') and self.model_tester.chunk_length > 1:
                                A__       =    seq_length * self.model_tester.chunk_length
                  else:
                         A__       =    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:
                         A__       =    outputs.decoder_hidden_states

                         self.asseretIsInstance(UpperCAmelCase__					, (list, tuple))
                         self.assertEqual(len(UpperCAmelCase__)					, UpperCAmelCase__)
                         A__       =    getattr(self.model_tester					, '''seq_length'''					, UpperCAmelCase__)
                         A__       =    getattr(self.model_tester					, '''decoder_seq_length'''					, UpperCAmelCase__)

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

           A__ ,							A__       =    self.model_tester.prepare_config_and_inputs_for_common()

           for model_class in self.all_model_classes:
                  A__       =    True
                  check_hidden_states_output(UpperCAmelCase__					, UpperCAmelCase__					, UpperCAmelCase__)

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

                  check_hidden_states_output(UpperCAmelCase__					, UpperCAmelCase__					, UpperCAmelCase__)


    def        SCREAMING_SNAKE_CASE							(			self      :	Optional[Any]					, UpperCAmelCase__      :	Dict					, UpperCAmelCase__      :	Dict					, UpperCAmelCase__      :	Dict=False)						->int:





           '''simple docstring'''

           A__       =    super()._prepare_for_class(UpperCAmelCase__					, UpperCAmelCase__					, return_labels=UpperCAmelCase__)

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

           return inputs_dict


    def        SCREAMING_SNAKE_CASE							(			self      :	Optional[Any])						->Union[str, Any]:





           '''simple docstring'''

           A__       =    self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_model(*UpperCAmelCase__)


    @unittest.skip(reason='''EfficientFormer does not implement masked image modeling yet''')
    def        SCREAMING_SNAKE_CASE							(			self      :	str)						->str:





           '''simple docstring'''

           A__       =    self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_masked_image_modeling(*UpperCAmelCase__)


    def        SCREAMING_SNAKE_CASE							(			self      :	Any)						->Tuple:





           '''simple docstring'''

           A__       =    self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_image_classification(*UpperCAmelCase__)


    @slow
    def        SCREAMING_SNAKE_CASE							(			self      :	Tuple)						->Optional[int]:





           '''simple docstring'''

           for model_name in TF_EFFICIENTFORMER_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
                  A__       =    TFEfficientFormerModel.from_pretrained(UpperCAmelCase__)
                  self.assertIsNotNone(UpperCAmelCase__)


    def        SCREAMING_SNAKE_CASE							(			self      :	Any)						->str:





           '''simple docstring'''

           A__ ,							A__       =    self.model_tester.prepare_config_and_inputs_for_common()
           A__       =    True

           A__       =    getattr(self.model_tester					, '''seq_length'''					, UpperCAmelCase__)
           A__       =    getattr(self.model_tester					, '''encoder_seq_length'''					, UpperCAmelCase__)
           A__       =    getattr(self.model_tester					, '''key_length'''					, UpperCAmelCase__)
           A__       =    getattr(self.model_tester					, '''chunk_length'''					, UpperCAmelCase__)

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

           for model_class in self.all_model_classes:
                  A__       =    True
                  A__       =    False
                  A__       =    True
                  A__       =    model_class(UpperCAmelCase__)

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

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

                  A__       =    outputs.encoder_attentions if config.is_encoder_decoder else outputs.attentions
                  self.assertEqual(len(UpperCAmelCase__)					, 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        SCREAMING_SNAKE_CASE							(			self      :	List[str])						->Optional[Any]:





           '''simple docstring'''

           A__ ,							A__       =    self.model_tester.prepare_config_and_inputs_for_common()

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







def SCREAMING_SNAKE_CASE				(       )   ->       Any:



       """simple docstring"""

       A__       =    Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''						)
       return image


@require_tf
@require_vision
class 				UpperCamelCase_    ( unittest.TestCase     ):

    '''simple docstring'''






    @cached_property
    def        SCREAMING_SNAKE_CASE							(			self      :	List[str])						->List[str]:





           '''simple docstring'''

           return (
               EfficientFormerImageProcessor.from_pretrained('''snap-research/efficientformer-l1-300''')
               if is_vision_available()
               else None
           )


    @slow
    def        SCREAMING_SNAKE_CASE							(			self      :	List[str])						->Any:





           '''simple docstring'''

           A__       =    TFEfficientFormerForImageClassification.from_pretrained('''snap-research/efficientformer-l1-300''')
           A__       =    self.default_image_processor
           A__       =    prepare_img()
           A__       =    image_processor(images=UpperCAmelCase__					, return_tensors='''tf''')
           # forward pass
           A__       =    model(**UpperCAmelCase__					, training=UpperCAmelCase__)
           # verify the logits
           A__       =    tf.TensorShape((1, 1_000))
           self.assertEqual(outputs.logits.shape					, UpperCAmelCase__)
           A__       =    tf.constant([-0.0555, 0.4825, -0.0852])
           self.assertTrue(np.allclose(outputs.logits[0, :3]					, UpperCAmelCase__					, atol=1e-4))





    @slow
    def        SCREAMING_SNAKE_CASE							(			self      :	Dict)						->int:





           '''simple docstring'''

           A__       =    TFEfficientFormerForImageClassificationWithTeacher.from_pretrained(
               '''snap-research/efficientformer-l1-300''')
           A__       =    self.default_image_processor
           A__       =    prepare_img()
           A__       =    image_processor(images=UpperCAmelCase__					, return_tensors='''tf''')
           # forward pass
           A__       =    model(**UpperCAmelCase__					, training=UpperCAmelCase__)
           # verify the logits
           A__       =    tf.TensorShape((1, 1_000))
           self.assertEqual(outputs.logits.shape					, UpperCAmelCase__)
           A__       =    tf.constant([-0.1312, 0.4353, -1.0499])
           self.assertTrue(np.allclose(outputs.logits[0, :3]					, UpperCAmelCase__					, atol=1e-4))


