"""simple docstring"""







import inspect
import unittest

from transformers import EfficientNetConfig
from transformers.testing_utils import is_pipeline_test, require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

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


if is_torch_available():
	import torch

	from transformers import EfficientNetForImageClassification, EfficientNetModel
	from transformers.models.efficientnet.modeling_efficientnet import EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
	from PIL import Image

	from transformers import AutoImageProcessor

class 						UpperCamelCase     :



				def __init__(self							: Union[str, Any]      ,    _A							: List[str]      ,    _A							: Any=13      ,    _A							: Any=32      ,    _A							: Tuple=3      ,    _A							: List[str]=[3, 3, 5]      ,    _A							: Optional[Any]=[32, 16, 24]      ,    _A							: int=[16, 24, 40]      ,    _A							: Tuple=[1, 1, 2]      ,    _A							: Tuple=[1, 1, 2]      ,    _A							: Dict=[1, 6, 6]      ,    _A							: Union[str, Any]=True      ,    _A							: Tuple=True      ,    _A							: Dict=37      ,    _A							: Union[str, Any]="gelu"      ,    _A							: str=10      ,    )       ->      Dict:
					__snake_case     :				Dict 						=					parent
					__snake_case     :				str 						=					batch_size
					__snake_case     :				List[Any] 						=					image_size
					__snake_case     :				List[Any] 						=					num_channels
					__snake_case     :				Union[str, Any] 						=					kernel_sizes
					__snake_case     :				Tuple 						=					in_channels
					__snake_case     :				Any 						=					out_channels
					__snake_case     :				Dict 						=					strides
					__snake_case     :				Dict 						=					num_block_repeats
					__snake_case     :				List[Any] 						=					expand_ratios
					__snake_case     :				Any 						=					is_training
					__snake_case     :				Tuple 						=					hidden_act
					__snake_case     :				List[Any] 						=					num_labels
					__snake_case     :				List[str] 						=					use_labels



				def    _lowercase	(self							: Dict)       ->      Any:
					__snake_case     :				Optional[Any] 						=					floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size])

					__snake_case     :				Union[str, Any] 						=					None
					if self.use_labels:
						__snake_case     :				int 						=					ids_tensor([self.batch_size]      ,    self.num_labels)

					__snake_case     :				Optional[Any] 						=					self.get_config()
					return config, pixel_values, labels



				def    _lowercase	(self							: Tuple)       ->      Any:
					return EfficientNetConfig(
					    num_channels=self.num_channels      ,    kernel_sizes=self.kernel_sizes      ,    in_channels=self.in_channels      ,    out_channels=self.out_channels      ,    strides=self.strides      ,    num_block_repeats=self.num_block_repeats      ,    expand_ratios=self.expand_ratios      ,    hidden_act=self.hidden_act      ,    num_labels=self.num_labels      ,    )



				def    _lowercase	(self							: Tuple      ,    _A							: Optional[int]      ,    _A							: Optional[Any]      ,    _A							: Dict)       ->      Dict:
					__snake_case     :				Union[str, Any] 						=					EfficientNetModel(config=_A)
					model.to(_A)
					model.eval()
					__snake_case     :				List[Any] 						=					model(_A)
					# expected last hidden states: B, C, H // 4, W // 4
					self.parent.assertEqual(
					    result.last_hidden_state.shape      ,    (self.batch_size, config.hidden_dim, self.image_size // 4, self.image_size // 4)      ,    )



				def    _lowercase	(self							: List[str]      ,    _A							: List[Any]      ,    _A							: Any      ,    _A							: List[str])       ->      Dict:
					__snake_case     :				Any 						=					EfficientNetForImageClassification(_A)
					model.to(_A)
					model.eval()
					__snake_case     :				Optional[Any] 						=					model(_A      ,    labels=_A)
					self.parent.assertEqual(result.logits.shape      ,    (self.batch_size, self.num_labels))






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







@require_torch
class 						UpperCamelCase     (       lowercase ,      lowercase ,      unittest.TestCase		):
				UpperCAmelCase							: Optional[Any]       =						(EfficientNetModel, EfficientNetForImageClassification) if is_torch_available() else ()
				UpperCAmelCase							: Optional[Any]       =						(
				    {"""feature-extraction""": EfficientNetModel, """image-classification""": EfficientNetForImageClassification}
				    if is_torch_available()
				    else {}
				)

				UpperCAmelCase							: Tuple       =						False
				UpperCAmelCase							: List[str]       =						False
				UpperCAmelCase							: Dict       =						False
				UpperCAmelCase							: Union[str, Any]       =						False
				UpperCAmelCase							: Union[str, Any]       =						False



				def    _lowercase	(self							: Dict)       ->      Optional[Any]:
					__snake_case     :				List[Any] 						=					EfficientNetModelTester(self)
					__snake_case     :				int 						=					ConfigTester(
					    self      ,    config_class=_A      ,    has_text_modality=_A      ,    hidden_size=37)



				def    _lowercase	(self							: Union[str, Any])       ->      int:
					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    _lowercase	(self							: Any)       ->      Optional[int]:
					return



				@unittest.skip(reason='EfficientNet does not use inputs_embeds')
				def    _lowercase	(self							: List[Any])       ->      Any:
					pass



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



				@unittest.skip(reason='EfficientNet does not use feedforward chunking')
				def    _lowercase	(self							: Tuple)       ->      Dict:
					pass



				def    _lowercase	(self							: Any)       ->      str:
					__snake_case						,							__snake_case     :				int 						=					self.model_tester.prepare_config_and_inputs_for_common()

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

						__snake_case     :				str 						=					['pixel_values']
						self.assertListEqual(arg_names[:1]      ,    _A)



				def    _lowercase	(self							: int)       ->      Union[str, Any]:
					__snake_case     :				List[Any] 						=					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_model(*_A)



				def    _lowercase	(self							: Tuple)       ->      List[str]:
					def check_hidden_states_output(_A							: List[str]      ,    _A							: int      ,    _A							: List[Any]):
						__snake_case     :				Union[str, Any] 						=					model_class(_A)
						model.to(_A)
						model.eval()

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

						__snake_case     :				List[str] 						=					outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states
						__snake_case     :				Dict 						=					sum(config.num_block_repeats) * 4
						self.assertEqual(len(_A)      ,    _A)

						# EfficientNet's feature maps are of shape (batch_size, num_channels, height, width)
						self.assertListEqual(
						    list(hidden_states[0].shape[-2:])      ,    [self.model_tester.image_size // 2, self.model_tester.image_size // 2]      ,    )

					__snake_case						,							__snake_case     :				Optional[Any] 						=					self.model_tester.prepare_config_and_inputs_for_common()

					for model_class in self.all_model_classes:
						__snake_case     :				Any 						=					True
						check_hidden_states_output(_A      ,    _A      ,    _A)

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

						check_hidden_states_output(_A      ,    _A      ,    _A)



				def    _lowercase	(self							: Optional[Any])       ->      List[str]:
					__snake_case     :				int 						=					self.model_tester.prepare_config_and_inputs()
					self.model_tester.create_and_check_for_image_classification(*_A)



				@unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.')
				def    _lowercase	(self							: str)       ->      List[Any]:
					pass



				@slow
				def    _lowercase	(self							: List[Any])       ->      Dict:
					for model_name in EFFICIENTNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
						__snake_case     :				Optional[Any] 						=					EfficientNetModel.from_pretrained(_A)
						self.assertIsNotNone(_A)






				@is_pipeline_test
				@require_vision
				@slow
				def    _lowercase	(self							: Optional[Any])       ->      int:
					super().test_pipeline_image_classification()







def     __UpperCAmelCase							(       )				->     List[Any]:

	'''simple docstring'''





	__snake_case     :				Union[str, Any] 						=					Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'    )
	return image

@require_torch
@require_vision
class 						UpperCamelCase     (       unittest.TestCase		):



				@cached_property
				def    _lowercase	(self							: Tuple)       ->      List[str]:
					return AutoImageProcessor.from_pretrained('google/efficientnet-b7') if is_vision_available() else None






				@slow
				def    _lowercase	(self							: Optional[Any])       ->      List[str]:
					__snake_case     :				Union[str, Any] 						=					EfficientNetForImageClassification.from_pretrained('google/efficientnet-b7').to(_A)

					__snake_case     :				str 						=					self.default_image_processor
					__snake_case     :				List[Any] 						=					prepare_img()
					__snake_case     :				Optional[int] 						=					image_processor(images=_A      ,    return_tensors='pt').to(_A)

					# forward pass
					with torch.no_grad():
						__snake_case     :				Tuple 						=					model(**_A)

					# verify the logits
					__snake_case     :				int 						=					torch.Size((1, 10_00))
					self.assertEqual(outputs.logits.shape      ,    _A)

					__snake_case     :				str 						=					torch.tensor([0.0_001, 0.0_002, 0.0_002]).to(_A)
					self.assertTrue(torch.allclose(outputs.logits[0, :3]      ,    _A      ,    atol=1E-4))






"""simple docstring"""







from typing import TypeVar

from .arrow_dataset import Dataset, _split_by_node_map_style_dataset
from .iterable_dataset import IterableDataset, _split_by_node_iterable_dataset


_a	:	List[Any]=					TypeVar("DatasetType", Dataset, IterableDataset)
def     __UpperCAmelCase							(       UpperCAmelCase_       :							DatasetType     ,				UpperCAmelCase_       :							int     ,				UpperCAmelCase_       :							int    )				->     DatasetType:

	'''simple docstring'''





	if isinstance(UpperCAmelCase_     ,				UpperCAmelCase_    ):
		return _split_by_node_map_style_dataset(UpperCAmelCase_     ,				rank=UpperCAmelCase_     ,				world_size=UpperCAmelCase_    )
	else:
		return _split_by_node_iterable_dataset(UpperCAmelCase_     ,				rank=UpperCAmelCase_     ,				world_size=UpperCAmelCase_    )





