
"""simple docstring"""






import inspect
import unittest

from transformers import BitConfig
from transformers.testing_utils import 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, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
					import torch
					from torch import nn

					from transformers import BitBackbone, BitForImageClassification, BitImageProcessor, BitModel
					from transformers.models.bit.modeling_bit import BIT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
					from PIL import Image



class 			lowercase					:






			def __init__(self  :	int      ,SCREAMING_SNAKE_CASE_  :	Tuple      ,SCREAMING_SNAKE_CASE_  :	int=3      ,SCREAMING_SNAKE_CASE_  :	Dict=32      ,SCREAMING_SNAKE_CASE_  :	Dict=3      ,SCREAMING_SNAKE_CASE_  :	List[Any]=10      ,SCREAMING_SNAKE_CASE_  :	Any=[8, 16, 32, 64]      ,SCREAMING_SNAKE_CASE_  :	List[Any]=[1, 1, 2, 1]      ,SCREAMING_SNAKE_CASE_  :	List[Any]=True      ,SCREAMING_SNAKE_CASE_  :	str=True      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]="relu"      ,SCREAMING_SNAKE_CASE_  :	str=3      ,SCREAMING_SNAKE_CASE_  :	Dict=None      ,SCREAMING_SNAKE_CASE_  :	str=["stage2", "stage3", "stage4"]      ,SCREAMING_SNAKE_CASE_  :	List[str]=[2, 3, 4]      ,SCREAMING_SNAKE_CASE_  :	Tuple=1      ,)       ->	Optional[Any]:





				"""simple docstring"""


				lowerCAmelCase												=  parent
				lowerCAmelCase												=  batch_size
				lowerCAmelCase												=  image_size
				lowerCAmelCase												=  num_channels
				lowerCAmelCase												=  embeddings_size
				lowerCAmelCase												=  hidden_sizes
				lowerCAmelCase												=  depths
				lowerCAmelCase												=  is_training
				lowerCAmelCase												=  use_labels
				lowerCAmelCase												=  hidden_act
				lowerCAmelCase												=  num_labels
				lowerCAmelCase												=  scope
				lowerCAmelCase												=  len(SCREAMING_SNAKE_CASE_			)
				lowerCAmelCase												=  out_features
				lowerCAmelCase												=  out_indices
				lowerCAmelCase												=  num_groups






			def 					UpperCAmelCase							(self  :	Any			)       ->	Dict:





				"""simple docstring"""


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

				lowerCAmelCase												=  None
				if self.use_labels:
					lowerCAmelCase												=  ids_tensor([self.batch_size]      ,self.num_labels			)

				lowerCAmelCase												=  self.get_config()

				return config, pixel_values, labels






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





				"""simple docstring"""


				return BitConfig(
				    num_channels=self.num_channels      ,embeddings_size=self.embeddings_size      ,hidden_sizes=self.hidden_sizes      ,depths=self.depths      ,hidden_act=self.hidden_act      ,num_labels=self.num_labels      ,out_features=self.out_features      ,out_indices=self.out_indices      ,num_groups=self.num_groups      ,)






			def 					UpperCAmelCase							(self  :	Optional[int]      ,SCREAMING_SNAKE_CASE_  :	Any      ,SCREAMING_SNAKE_CASE_  :	Any      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]			)       ->	Any:





				"""simple docstring"""


				lowerCAmelCase												=  BitModel(config=SCREAMING_SNAKE_CASE_			)
				model.to(SCREAMING_SNAKE_CASE_			)
				model.eval()
				lowerCAmelCase												=  model(SCREAMING_SNAKE_CASE_			)
				self.parent.assertEqual(
				    result.last_hidden_state.shape      ,(self.batch_size, self.hidden_sizes[-1], self.image_size // 32, self.image_size // 32)      ,)






			def 					UpperCAmelCase							(self  :	List[Any]      ,SCREAMING_SNAKE_CASE_  :	List[str]      ,SCREAMING_SNAKE_CASE_  :	List[Any]      ,SCREAMING_SNAKE_CASE_  :	int			)       ->	int:





				"""simple docstring"""


				lowerCAmelCase												=  self.num_labels
				lowerCAmelCase												=  BitForImageClassification(SCREAMING_SNAKE_CASE_			)
				model.to(SCREAMING_SNAKE_CASE_			)
				model.eval()
				lowerCAmelCase												=  model(SCREAMING_SNAKE_CASE_      ,labels=SCREAMING_SNAKE_CASE_			)
				self.parent.assertEqual(result.logits.shape      ,(self.batch_size, self.num_labels)			)






			def 					UpperCAmelCase							(self  :	str      ,SCREAMING_SNAKE_CASE_  :	Optional[int]      ,SCREAMING_SNAKE_CASE_  :	List[str]      ,SCREAMING_SNAKE_CASE_  :	Dict			)       ->	str:





				"""simple docstring"""


				lowerCAmelCase												=  BitBackbone(config=SCREAMING_SNAKE_CASE_			)
				model.to(SCREAMING_SNAKE_CASE_			)
				model.eval()
				lowerCAmelCase												=  model(SCREAMING_SNAKE_CASE_			)

				# verify feature maps
				self.parent.assertEqual(len(result.feature_maps			)      ,len(config.out_features			)			)
				self.parent.assertListEqual(list(result.feature_maps[0].shape			)      ,[self.batch_size, self.hidden_sizes[1], 4, 4]			)

				# verify channels
				self.parent.assertEqual(len(model.channels			)      ,len(config.out_features			)			)
				self.parent.assertListEqual(model.channels      ,config.hidden_sizes[1:]			)

				# verify backbone works with out_features=None
				lowerCAmelCase												=  None
				lowerCAmelCase												=  BitBackbone(config=SCREAMING_SNAKE_CASE_			)
				model.to(SCREAMING_SNAKE_CASE_			)
				model.eval()
				lowerCAmelCase												=  model(SCREAMING_SNAKE_CASE_			)

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

				# verify channels
				self.parent.assertEqual(len(model.channels			)      ,1			)
				self.parent.assertListEqual(model.channels      ,[config.hidden_sizes[-1]]			)






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





				"""simple docstring"""


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





@require_torch
class 			lowercase					( lowercase__   ,lowercase__   ,unittest.TestCase		):
			lowercase        =     (BitModel, BitForImageClassification, BitBackbone) if is_torch_available() else ()
			lowercase        =     (
			    {'''feature-extraction''': BitModel, '''image-classification''': BitForImageClassification}
			    if is_torch_available()
			    else {}
			)

			lowercase        =     False
			lowercase        =     False
			lowercase        =     False
			lowercase        =     False
			lowercase        =     False






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





				"""simple docstring"""


				lowerCAmelCase												=  BitModelTester(self			)
				lowerCAmelCase												=  ConfigTester(self      ,config_class=SCREAMING_SNAKE_CASE_      ,has_text_modality=SCREAMING_SNAKE_CASE_			)






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





				"""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 					UpperCAmelCase							(self  :	Tuple			)       ->	Tuple:





				"""simple docstring"""


				return






			@unittest.skip(reason='''Bit does not output attentions'''			)
			def 					UpperCAmelCase							(self  :	Union[str, Any]			)       ->	Union[str, Any]:





				"""simple docstring"""


				pass






			@unittest.skip(reason='''Bit does not use inputs_embeds'''			)
			def 					UpperCAmelCase							(self  :	str			)       ->	Optional[int]:





				"""simple docstring"""


				pass






			@unittest.skip(reason='''Bit does not support input and output embeddings'''			)
			def 					UpperCAmelCase							(self  :	Any			)       ->	int:





				"""simple docstring"""


				pass






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





				"""simple docstring"""


				lowerCAmelCase							,       lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs_for_common()

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

					lowerCAmelCase												=  ['''pixel_values''']
					self.assertListEqual(arg_names[:1]      ,SCREAMING_SNAKE_CASE_			)






			def 					UpperCAmelCase							(self  :	Dict			)       ->	Dict:





				"""simple docstring"""


				lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_model(*SCREAMING_SNAKE_CASE_			)






			def 					UpperCAmelCase							(self  :	Tuple			)       ->	List[Any]:





				"""simple docstring"""


				lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_backbone(*SCREAMING_SNAKE_CASE_			)






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





				"""simple docstring"""


				lowerCAmelCase							,       lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs_for_common()

				for model_class in self.all_model_classes:
					lowerCAmelCase												=  model_class(config=SCREAMING_SNAKE_CASE_			)
					for name, module in model.named_modules():
						if isinstance(SCREAMING_SNAKE_CASE_      ,(nn.BatchNormad, nn.GroupNorm)			):
							self.assertTrue(
							    torch.all(module.weight == 1			)      ,msg=F"""Parameter {name} of model {model_class} seems not properly initialized"""      ,)
							self.assertTrue(
							    torch.all(module.bias == 0			)      ,msg=F"""Parameter {name} of model {model_class} seems not properly initialized"""      ,)






			def 					UpperCAmelCase							(self  :	str			)       ->	Dict:





				"""simple docstring"""


				def check_hidden_states_output(SCREAMING_SNAKE_CASE_  :	Optional[Any]      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]      ,SCREAMING_SNAKE_CASE_  :	Optional[Any]			):
					lowerCAmelCase												=  model_class(SCREAMING_SNAKE_CASE_			)
					model.to(SCREAMING_SNAKE_CASE_			)
					model.eval()

					with torch.no_grad():
						lowerCAmelCase												=  model(**self._prepare_for_class(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)			)

					lowerCAmelCase												=  outputs.encoder_hidden_states if config.is_encoder_decoder else outputs.hidden_states

					lowerCAmelCase												=  self.model_tester.num_stages
					self.assertEqual(len(SCREAMING_SNAKE_CASE_			)      ,expected_num_stages + 1			)

					# Bit'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 // 4, self.model_tester.image_size // 4]      ,)

				lowerCAmelCase							,       lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs_for_common()
				lowerCAmelCase												=  ['''preactivation''', '''bottleneck''']
				for model_class in self.all_model_classes:
					for layer_type in layers_type:
						lowerCAmelCase												=  layer_type
						lowerCAmelCase												=  True
						check_hidden_states_output(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)

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

						check_hidden_states_output(SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_      ,SCREAMING_SNAKE_CASE_			)






			@unittest.skip(reason='''Bit does not use feedforward chunking'''			)
			def 					UpperCAmelCase							(self  :	List[Any]			)       ->	str:





				"""simple docstring"""


				pass






			def 					UpperCAmelCase							(self  :	Tuple			)       ->	Dict:





				"""simple docstring"""


				lowerCAmelCase												=  self.model_tester.prepare_config_and_inputs()
				self.model_tester.create_and_check_for_image_classification(*SCREAMING_SNAKE_CASE_			)






			@slow
			def 					UpperCAmelCase							(self  :	Any			)       ->	Union[str, Any]:





				"""simple docstring"""


				for model_name in BIT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
					lowerCAmelCase												=  BitModel.from_pretrained(SCREAMING_SNAKE_CASE_			)
					self.assertIsNotNone(SCREAMING_SNAKE_CASE_			)





def __magic_name__	(				)					->						int:

	'''simple docstring'''




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



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






			@cached_property
			def 					UpperCAmelCase							(self  :	Union[str, Any]			)       ->	List[str]:





				"""simple docstring"""


				return (
				    BitImageProcessor.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]			) if is_vision_available() else None
				)






			@slow
			def 					UpperCAmelCase							(self  :	List[Any]			)       ->	int:





				"""simple docstring"""


				lowerCAmelCase												=  BitForImageClassification.from_pretrained(BIT_PRETRAINED_MODEL_ARCHIVE_LIST[0]			).to(SCREAMING_SNAKE_CASE_			)

				lowerCAmelCase												=  self.default_image_processor
				lowerCAmelCase												=  prepare_img()
				lowerCAmelCase												=  image_processor(images=SCREAMING_SNAKE_CASE_      ,return_tensors='''pt'''			).to(SCREAMING_SNAKE_CASE_			)

				# forward pass
				with torch.no_grad():
					lowerCAmelCase												=  model(**SCREAMING_SNAKE_CASE_			)

				# verify the logits
				lowerCAmelCase												=  torch.Size((1, 1_000)			)
				self.assertEqual(outputs.logits.shape      ,SCREAMING_SNAKE_CASE_			)

				lowerCAmelCase												=  torch.tensor([[-0.65_26, -0.52_63, -1.43_98]]			).to(SCREAMING_SNAKE_CASE_			)

				self.assertTrue(torch.allclose(outputs.logits[0, :3]      ,SCREAMING_SNAKE_CASE_      ,atol=1e-4			)			)





@require_torch
class 			lowercase					( lowercase__   ,unittest.TestCase		):
			lowercase        =     (BitBackbone,) if is_torch_available() else ()
			lowercase        =     BitConfig

			lowercase        =     False






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





				"""simple docstring"""


				lowerCAmelCase												=  BitModelTester(self			)




"""simple docstring"""






from typing import Dict

from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
from transformers.testing_utils import (
    TestCasePlus,
    execute_subprocess_async,
    get_torch_dist_unique_port,
    require_torch_multi_gpu,
    require_torch_neuroncore,
)
from transformers.training_args import ParallelMode
from transformers.utils import logging


UpperCAmelCase        =				logging.get_logger(__name__)


if is_torch_available():
					import torch
					from torch import nn
					from torch.utils.data import Dataset

					from transformers import Trainer



					class 			lowercase					( lowercase__		):






								def __init__(self  :	int      ,SCREAMING_SNAKE_CASE_  :	int = 101			)       ->	List[str]:





									"""simple docstring"""


									lowerCAmelCase												=  length






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





									"""simple docstring"""


									return self.length






								def __getitem__(self  :	str      ,SCREAMING_SNAKE_CASE_  :	Union[str, Any]			)       ->	int:





									"""simple docstring"""


									return i





					class 			lowercase					:






								def __call__(self  :	List[Any]      ,SCREAMING_SNAKE_CASE_  :	List[Any]			)       ->	Optional[Any]:





									"""simple docstring"""


									return {"input_ids": torch.tensor(SCREAMING_SNAKE_CASE_			), "labels": torch.tensor(SCREAMING_SNAKE_CASE_			)}





					class 			lowercase					( nn.Module		):






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





									"""simple docstring"""


									super().__init__()
									# Add some (unused) params otherwise DDP will complain.
									lowerCAmelCase												=  nn.Linear(120      ,80			)






								def 					UpperCAmelCase							(self  :	Optional[int]      ,SCREAMING_SNAKE_CASE_  :	str      ,SCREAMING_SNAKE_CASE_  :	Optional[int]=None			)       ->	Union[str, Any]:





									"""simple docstring"""


									if labels is not None:
										return torch.tensor(0.0      ,device=input_ids.device			), input_ids
									else:
										return input_ids



class 			lowercase					( lowercase__		):






			@require_torch_neuroncore
			def 					UpperCAmelCase							(self  :	Union[str, Any]			)       ->	Union[str, Any]:





				"""simple docstring"""


				lowerCAmelCase												=  F"""--nproc_per_node=2
            --master_port={get_torch_dist_unique_port()}
            {self.test_file_dir}/test_trainer_distributed.py
        """.split()
				lowerCAmelCase												=  self.get_auto_remove_tmp_dir()
				lowerCAmelCase												=  F"""--output_dir {output_dir}""".split()
				lowerCAmelCase												=  ['''torchrun'''] + distributed_args + args
				execute_subprocess_async(SCREAMING_SNAKE_CASE_      ,env=self.get_env()			)
				# successful return here == success - any errors would have caused an error in the sub-call





class 			lowercase					( lowercase__		):






			@require_torch_multi_gpu
			def 					UpperCAmelCase							(self  :	List[str]			)       ->	List[str]:





				"""simple docstring"""


				lowerCAmelCase												=  F"""--nproc_per_node={torch.cuda.device_count()}
            --master_port={get_torch_dist_unique_port()}
            {self.test_file_dir}/test_trainer_distributed.py
        """.split()
				lowerCAmelCase												=  self.get_auto_remove_tmp_dir()
				lowerCAmelCase												=  F"""--output_dir {output_dir}""".split()
				lowerCAmelCase												=  ['''torchrun'''] + distributed_args + args
				execute_subprocess_async(SCREAMING_SNAKE_CASE_      ,env=self.get_env()			)
				# successful return here == success - any errors would have caused an error in the sub-call


if __name__ == "__main__":
					# The script below is meant to be run under torch.distributed, on a machine with multiple GPUs:
					#
					# PYTHONPATH="src" python -m torch.distributed.run --nproc_per_node 2 --output_dir output_dir ./tests/test_trainer_distributed.py

					UpperCAmelCase        =				HfArgumentParser((TrainingArguments,))
					UpperCAmelCase        =				parser.parse_args_into_dataclasses()[0]

					logger.warning(
					    f"""Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}, """
					    f"""distributed training: {training_args.parallel_mode != ParallelMode.NOT_DISTRIBUTED}"""
					)

					# Essentially, what we want to verify in the distributed case is that we get all samples back,
					# in the right order. (this is crucial for prediction for instance)
					for dataset_length in [1_0_1, 4_0, 7]:
										UpperCAmelCase        =				DummyDataset(dataset_length)




										def __magic_name__	(				_lowerCamelCase:		EvalPrediction			)					->						Dict:

											'''simple docstring'''




											lowerCAmelCase												=  list(range(len(_lowerCamelCase			)			)			)
											lowerCAmelCase												=  p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
											if not success and training_args.local_rank == 0:
												logger.warning(
												    '''Predictions and/or labels do not match expected results:\n  - predictions: '''
												    F"""{p.predictions.tolist()}\n  - labels: {p.label_ids.tolist()}\n  - expected: {sequential}"""			)
											return {"success": success}

										UpperCAmelCase        =				Trainer(
										    model=DummyModel(),
										    args=training_args,
										    data_collator=DummyDataCollator(),
										    eval_dataset=dataset,
										    compute_metrics=compute_metrics,
										)
										UpperCAmelCase        =				trainer.evaluate()
										logger.info(metrics)
										if metrics["eval_success"] is not True:
															logger.error(metrics)
															exit(1)

										UpperCAmelCase        =				trainer.predict(dataset)
										logger.info(p.metrics)
										if p.metrics["test_success"] is not True:
															logger.error(p.metrics)
															exit(1)

										UpperCAmelCase        =				2

										UpperCAmelCase        =				trainer.evaluate()
										logger.info(metrics)
										if metrics["eval_success"] is not True:
															logger.error(metrics)
															exit(1)

										UpperCAmelCase        =				trainer.predict(dataset)
										logger.info(p.metrics)
										if p.metrics["test_success"] is not True:
															logger.error(p.metrics)
															exit(1)



										UpperCAmelCase        =				None


