


'''simple docstring'''




import collections
import inspect
import unittest

from transformers import NatConfig
from transformers.testing_utils import require_natten, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
				import torch
				from torch import nn

				from transformers import NatBackbone, NatForImageClassification, NatModel
				from transformers.models.nat.modeling_nat import NAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
				from PIL import Image

				from transformers import AutoImageProcessor





class _snake_case :
							"""simple docstring"""





							def __init__( self      ,					__A      ,					__A=13      ,					__A=64      ,					__A=4      ,					__A=3      ,					__A=16      ,					__A=[1, 2, 1]      ,					__A=[2, 4, 8]      ,					__A=3      ,					__A=2.0      ,					__A=True      ,					__A=0.0      ,					__A=0.0      ,					__A=0.1      ,					__A="gelu"      ,					__A=True      ,					__A=0.02      ,					__A=1E-5      ,					__A=True      ,					__A=None      ,					__A=True      ,					__A=10      ,					__A=["stage1", "stage2"]      ,					__A=[1, 2]      ,					)  ->str:



											"""simple docstring"""







											__lowerCamelCase		:   Optional[Any]           =       parent
											__lowerCamelCase		:   Any           =       batch_size
											__lowerCamelCase		:   Any           =       image_size
											__lowerCamelCase		:   str           =       patch_size
											__lowerCamelCase		:   Optional[int]           =       num_channels
											__lowerCamelCase		:   Optional[Any]           =       embed_dim
											__lowerCamelCase		:   Any           =       depths
											__lowerCamelCase		:   List[Any]           =       num_heads
											__lowerCamelCase		:   List[str]           =       kernel_size
											__lowerCamelCase		:   List[Any]           =       mlp_ratio
											__lowerCamelCase		:   int           =       qkv_bias
											__lowerCamelCase		:   Dict           =       hidden_dropout_prob
											__lowerCamelCase		:   int           =       attention_probs_dropout_prob
											__lowerCamelCase		:   Optional[Any]           =       drop_path_rate
											__lowerCamelCase		:   Optional[int]           =       hidden_act
											__lowerCamelCase		:   Optional[Any]           =       patch_norm
											__lowerCamelCase		:   List[str]           =       layer_norm_eps
											__lowerCamelCase		:   Dict           =       initializer_range
											__lowerCamelCase		:   Optional[Any]           =       is_training
											__lowerCamelCase		:   List[str]           =       scope
											__lowerCamelCase		:   Any           =       use_labels
											__lowerCamelCase		:   List[str]           =       num_labels
											__lowerCamelCase		:   List[str]           =       out_features
											__lowerCamelCase		:   Optional[Any]           =       out_indices







							def 		__SCREAMING_SNAKE_CASE( self						)  ->int:



											"""simple docstring"""







											__lowerCamelCase		:   Union[str, Any]           =       floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]						)

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

											__lowerCamelCase		:   List[Any]           =       self.get_config()

											return config, pixel_values, labels







							def 		__SCREAMING_SNAKE_CASE( self						)  ->List[str]:



											"""simple docstring"""







											return NatConfig(
											    num_labels=self.num_labels      ,					image_size=self.image_size      ,					patch_size=self.patch_size      ,					num_channels=self.num_channels      ,					embed_dim=self.embed_dim      ,					depths=self.depths      ,					num_heads=self.num_heads      ,					kernel_size=self.kernel_size      ,					mlp_ratio=self.mlp_ratio      ,					qkv_bias=self.qkv_bias      ,					hidden_dropout_prob=self.hidden_dropout_prob      ,					attention_probs_dropout_prob=self.attention_probs_dropout_prob      ,					drop_path_rate=self.drop_path_rate      ,					hidden_act=self.hidden_act      ,					patch_norm=self.patch_norm      ,					layer_norm_eps=self.layer_norm_eps      ,					initializer_range=self.initializer_range      ,					out_features=self.out_features      ,					out_indices=self.out_indices      ,					)







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



											"""simple docstring"""







											__lowerCamelCase		:   str           =       NatModel(config=A_						)
											model.to(A_						)
											model.eval()
											__lowerCamelCase		:   Optional[Any]           =       model(A_						)

											__lowerCamelCase		:   Tuple           =       (config.image_size // config.patch_size) // (2 ** (len(config.depths						) - 1))
											__lowerCamelCase		:   Optional[int]           =       int(config.embed_dim * 2 ** (len(config.depths						) - 1)						)

											self.parent.assertEqual(
											    result.last_hidden_state.shape      ,					(self.batch_size, expected_height, expected_width, expected_dim)						)







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



											"""simple docstring"""







											__lowerCamelCase		:   Dict           =       NatForImageClassification(A_						)
											model.to(A_						)
											model.eval()
											__lowerCamelCase		:   Optional[Any]           =       model(A_      ,					labels=A_						)
											self.parent.assertEqual(result.logits.shape      ,					(self.batch_size, self.num_labels)						)

											# test greyscale images
											__lowerCamelCase		:   Tuple           =       1
											__lowerCamelCase		:   Dict           =       NatForImageClassification(A_						)
											model.to(A_						)
											model.eval()

											__lowerCamelCase		:   Any           =       floats_tensor([self.batch_size, 1, self.image_size, self.image_size]						)
											__lowerCamelCase		:   Dict           =       model(A_						)
											self.parent.assertEqual(result.logits.shape      ,					(self.batch_size, self.num_labels)						)







							def 		__SCREAMING_SNAKE_CASE( self      ,					__A      ,					__A      ,					__A						)  ->List[Any]:



											"""simple docstring"""







											__lowerCamelCase		:   Optional[int]           =       NatBackbone(config=A_						)
											model.to(A_						)
											model.eval()
											__lowerCamelCase		:   str           =       model(A_						)

											# verify hidden states
											self.parent.assertEqual(len(result.feature_maps						)      ,					len(config.out_features						)						)
											self.parent.assertListEqual(list(result.feature_maps[0].shape						)      ,					[self.batch_size, model.channels[0], 16, 16]						)

											# verify channels
											self.parent.assertEqual(len(model.channels						)      ,					len(config.out_features						)						)

											# verify backbone works with out_features=None
											__lowerCamelCase		:   List[Any]           =       None
											__lowerCamelCase		:   str           =       NatBackbone(config=A_						)
											model.to(A_						)
											model.eval()
											__lowerCamelCase		:   Any           =       model(A_						)

											# verify feature maps
											self.parent.assertEqual(len(result.feature_maps						)      ,					1						)
											self.parent.assertListEqual(list(result.feature_maps[0].shape						)      ,					[self.batch_size, model.channels[-1], 4, 4]						)

											# verify channels
											self.parent.assertEqual(len(model.channels						)      ,					1						)







							def 		__SCREAMING_SNAKE_CASE( self						)  ->int:



											"""simple docstring"""







											__lowerCamelCase		:   Tuple           =       self.prepare_config_and_inputs()
											__lowerCamelCase    ,       __lowerCamelCase    ,       __lowerCamelCase		:   Dict           =       config_and_inputs
											__lowerCamelCase		:   Dict           =       {"pixel_values": pixel_values}
											return config, inputs_dict





@require_natten
@require_torch
class _snake_case (			A_     , A_     , unittest.TestCase ):
							"""simple docstring"""





							__a				       =  (
							    (
							        NatModel,
							        NatForImageClassification,
							        NatBackbone,
							    )
							    if is_torch_available()
							    else ()
							)
							__a				       =  (
							    {"feature-extraction": NatModel, "image-classification": NatForImageClassification}
							    if is_torch_available()
							    else {}
							)
							__a				       =  False

							__a				       =  False
							__a				       =  False
							__a				       =  False
							__a				       =  False







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



											"""simple docstring"""







											__lowerCamelCase		:   List[Any]           =       NatModelTester(self						)
											__lowerCamelCase		:   Union[str, Any]           =       ConfigTester(self      ,					config_class=A_      ,					embed_dim=37						)







							def 		__SCREAMING_SNAKE_CASE( self						)  ->int:



											"""simple docstring"""







											self.create_and_test_config_common_properties()
											self.config_tester.create_and_test_config_to_json_string()
											self.config_tester.create_and_test_config_to_json_file()
											self.config_tester.create_and_test_config_from_and_save_pretrained()
											self.config_tester.create_and_test_config_with_num_labels()
											self.config_tester.check_config_can_be_init_without_params()
											self.config_tester.check_config_arguments_init()







							def 		__SCREAMING_SNAKE_CASE( self						)  ->Any:



											"""simple docstring"""







											return







							def 		__SCREAMING_SNAKE_CASE( self						)  ->int:



											"""simple docstring"""







											__lowerCamelCase		:   int           =       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_model(*A_						)







							def 		__SCREAMING_SNAKE_CASE( self						)  ->str:



											"""simple docstring"""







											__lowerCamelCase		:   List[str]           =       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_for_image_classification(*A_						)







							def 		__SCREAMING_SNAKE_CASE( self						)  ->List[str]:



											"""simple docstring"""







											__lowerCamelCase		:   str           =       self.model_tester.prepare_config_and_inputs()
											self.model_tester.create_and_check_backbone(*A_						)







							@unittest.skip(reason="Nat does not use inputs_embeds"						)
							def 		__SCREAMING_SNAKE_CASE( self						)  ->Any:



											"""simple docstring"""







											pass







							@unittest.skip(reason="Nat does not use feedforward chunking"						)
							def 		__SCREAMING_SNAKE_CASE( self						)  ->Any:



											"""simple docstring"""







											pass







							def 		__SCREAMING_SNAKE_CASE( self						)  ->int:



											"""simple docstring"""







											__lowerCamelCase    ,       __lowerCamelCase		:   List[Any]           =       self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
															__lowerCamelCase		:   Optional[Any]           =       model_class(A_						)
															self.assertIsInstance(model.get_input_embeddings()      ,					(nn.Module)						)
															__lowerCamelCase		:   Dict           =       model.get_output_embeddings()
															self.assertTrue(x is None or isinstance(A_      ,					nn.Linear						)						)







							def 		__SCREAMING_SNAKE_CASE( self						)  ->Tuple:



											"""simple docstring"""







											__lowerCamelCase    ,       __lowerCamelCase		:   Union[str, Any]           =       self.model_tester.prepare_config_and_inputs_for_common()

											for model_class in self.all_model_classes:
															__lowerCamelCase		:   Optional[Any]           =       model_class(A_						)
															__lowerCamelCase		:   Any           =       inspect.signature(model.forward						)
															# signature.parameters is an OrderedDict => so arg_names order is deterministic
															__lowerCamelCase		:   Optional[Any]           =       [*signature.parameters.keys()]

															__lowerCamelCase		:   List[str]           =       ["pixel_values"]
															self.assertListEqual(arg_names[:1]      ,					A_						)







							def 		__SCREAMING_SNAKE_CASE( self						)  ->Any:



											"""simple docstring"""







											self.skipTest("Nat\'s attention operation is handled entirely by NATTEN."						)







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



											"""simple docstring"""







											__lowerCamelCase		:   Union[str, Any]           =       model_class(A_						)
											model.to(A_						)
											model.eval()

											with torch.no_grad():
															__lowerCamelCase		:   List[Any]           =       model(**self._prepare_for_class(A_      ,					A_						)						)

											__lowerCamelCase		:   List[Any]           =       outputs.hidden_states

											__lowerCamelCase		:   str           =       getattr(
											    self.model_tester      ,					"expected_num_hidden_layers"      ,					len(self.model_tester.depths						) + 1						)
											self.assertEqual(len(A_						)      ,					A_						)

											# Nat has a different seq_length
											__lowerCamelCase		:   Dict           =       (
											    config.patch_size
											    if isinstance(config.patch_size      ,					collections.abc.Iterable						)
											    else (config.patch_size, config.patch_size)
											)

											__lowerCamelCase		:   Tuple           =       image_size[0] // patch_size[0]
											__lowerCamelCase		:   Dict           =       image_size[1] // patch_size[1]

											self.assertListEqual(
											    list(hidden_states[0].shape[-3:]						)      ,					[height, width, self.model_tester.embed_dim]      ,					)

											if model_class.__name__ != "NatBackbone":
															__lowerCamelCase		:   Optional[int]           =       outputs.reshaped_hidden_states
															self.assertEqual(len(A_						)      ,					A_						)

															__lowerCamelCase    ,       __lowerCamelCase    ,       __lowerCamelCase    ,       __lowerCamelCase		:   int           =       reshaped_hidden_states[0].shape
															__lowerCamelCase		:   List[str]           =       (
															    reshaped_hidden_states[0].view(A_      ,					A_      ,					A_      ,					A_						).permute(0      ,					2      ,					3      ,					1						)
															)
															self.assertListEqual(
															    list(reshaped_hidden_states.shape[-3:]						)      ,					[height, width, self.model_tester.embed_dim]      ,					)







							def 		__SCREAMING_SNAKE_CASE( self						)  ->List[str]:



											"""simple docstring"""







											__lowerCamelCase    ,       __lowerCamelCase		:   int           =       self.model_tester.prepare_config_and_inputs_for_common()

											__lowerCamelCase		:   List[str]           =       (
											    self.model_tester.image_size
											    if isinstance(self.model_tester.image_size      ,					collections.abc.Iterable						)
											    else (self.model_tester.image_size, self.model_tester.image_size)
											)

											for model_class in self.all_model_classes:
															__lowerCamelCase		:   Optional[Any]           =       True
															self.check_hidden_states_output(A_      ,					A_      ,					A_      ,					A_						)

															# check that output_hidden_states also work using config
															del inputs_dict["output_hidden_states"]
															__lowerCamelCase		:   int           =       True

															self.check_hidden_states_output(A_      ,					A_      ,					A_      ,					A_						)







							@slow
							def 		__SCREAMING_SNAKE_CASE( self						)  ->int:



											"""simple docstring"""







											for model_name in NAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
															__lowerCamelCase		:   List[str]           =       NatModel.from_pretrained(A_						)
															self.assertIsNotNone(A_						)







							def 		__SCREAMING_SNAKE_CASE( self						)  ->str:



											"""simple docstring"""







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

											__lowerCamelCase		:   List[Any]           =       _config_zero_init(A_						)
											for model_class in self.all_model_classes:
															__lowerCamelCase		:   Dict           =       model_class(config=A_						)
															for name, param in model.named_parameters():
																			if "embeddings" not in name and param.requires_grad:
																							self.assertIn(
																							    ((param.data.mean() * 1E9).round() / 1E9).item()      ,					[0.0, 1.0]      ,					msg=f'Parameter {name} of model {model_class} seems not properly initialized'      ,					)





@require_natten
@require_vision
@require_torch
class _snake_case (			unittest.TestCase ):
							"""simple docstring"""





							@cached_property
							def 		__SCREAMING_SNAKE_CASE( self						)  ->List[Any]:



											"""simple docstring"""







											return AutoImageProcessor.from_pretrained("shi-labs/nat-mini-in1k-224"						) if is_vision_available() else None







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



											"""simple docstring"""







											__lowerCamelCase		:   Optional[Any]           =       NatForImageClassification.from_pretrained("shi-labs/nat-mini-in1k-224"						).to(A_						)
											__lowerCamelCase		:   str           =       self.default_image_processor

											__lowerCamelCase		:   Union[str, Any]           =       Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"						)
											__lowerCamelCase		:   Union[str, Any]           =       image_processor(images=A_      ,					return_tensors="pt"						).to(A_						)

											# forward pass
											with torch.no_grad():
															__lowerCamelCase		:   List[str]           =       model(**A_						)

											# verify the logits
											__lowerCamelCase		:   List[Any]           =       torch.Size((1, 1000)						)
											self.assertEqual(outputs.logits.shape      ,					A_						)
											__lowerCamelCase		:   List[Any]           =       torch.tensor([0.3805, -0.8676, -0.3912]						).to(A_						)
											self.assertTrue(torch.allclose(outputs.logits[0, :3]      ,					A_      ,					atol=1E-4						)						)





@require_torch
@require_natten
class _snake_case (			unittest.TestCase     , A_ ):
							"""simple docstring"""





							__a				       =  (NatBackbone,) if is_torch_available() else ()
							__a				       =  NatConfig







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



											"""simple docstring"""







											__lowerCamelCase		:   Tuple           =       NatModelTester(self						)









import json
import os
import tempfile

import datasets
from utils import generate_example_dataset, get_duration


_A									=		50_0000

_A			,	_A									=		os.path.split(__file__)
_A									=		os.path.join(RESULTS_BASEPATH, 'results', RESULTS_FILENAME.replace('.py', '.json'))

@get_duration
def 		_UpperCAmelCase							(     SCREAMING_SNAKE_CASE__     :   datasets.Dataset							):
     __UpperCamelCase       				=dataset.select(range(0     ,	len(SCREAMING_SNAKE_CASE__							)     ,	2							)							)

@get_duration
def 		_UpperCAmelCase							(     SCREAMING_SNAKE_CASE__     :   datasets.Dataset							):
     __UpperCamelCase       				=dataset.sort('numbers'							)

@get_duration
def 		_UpperCAmelCase							(     SCREAMING_SNAKE_CASE__     :   datasets.Dataset							):
     __UpperCamelCase       				=dataset.shuffle()

@get_duration
def 		_UpperCAmelCase							(     SCREAMING_SNAKE_CASE__     :   datasets.Dataset							):
     __UpperCamelCase       				=dataset.train_test_split(0.1							)

@get_duration
def 		_UpperCAmelCase							(     SCREAMING_SNAKE_CASE__     :   datasets.Dataset     ,	SCREAMING_SNAKE_CASE__     :   int=10							):
     for shard_id in range(SCREAMING_SNAKE_CASE__							):
          __UpperCamelCase       				=dataset.shard(SCREAMING_SNAKE_CASE__     ,	SCREAMING_SNAKE_CASE__							)

def 		_UpperCAmelCase							(     ):
     __UpperCamelCase       				={'num examples': SPEED_TEST_N_EXAMPLES}
     __UpperCamelCase       				=(select, sort, shuffle, train_test_split, shard)
     with tempfile.TemporaryDirectory() as tmp_dir:
          print('generating dataset'							)
          __UpperCamelCase       				=datasets.Features({'text': datasets.Value('string'							), 'numbers': datasets.Value('float32'							)}							)
          __UpperCamelCase       				=generate_example_dataset(
              os.path.join(SCREAMING_SNAKE_CASE__     ,	'dataset.arrow'							)     ,	SCREAMING_SNAKE_CASE__     ,	num_examples=SCREAMING_SNAKE_CASE__							)
          print('Functions'							)
          for func in functions:
               print(func.__name__							)
               __UpperCamelCase       				=func(SCREAMING_SNAKE_CASE__							)

     with open(SCREAMING_SNAKE_CASE__     ,	'wb'							) as f:
          f.write(json.dumps(SCREAMING_SNAKE_CASE__							).encode('utf-8'							)							)


if __name__ == "__main__":  # useful to run the profiler
   benchmark_indices_mapping()





