



import inspect
import unittest

from transformers import DPTConfig
from transformers.file_utils import is_torch_available, is_vision_available
from transformers.models.auto import get_values
from transformers.testing_utils import require_torch, require_vision, slow, torch_device

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 MODEL_MAPPING, DPTForDepthEstimation, DPTForSemanticSegmentation, DPTModel
		from transformers.models.dpt.modeling_dpt import DPT_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
		from PIL import Image

		from transformers import DPTImageProcessor


class 	lowerCamelCase__ :


						def __init__( self						:  List[Any]      ,       _a						:  int      ,       _a						:  int=2      ,       _a						:  List[str]=3_2      ,       _a						:  Optional[int]=1_6      ,       _a						:  Optional[int]=3      ,       _a						:  int=True      ,       _a						:  Dict=True      ,       _a						:  Dict=3_2      ,       _a						:  List[str]=4      ,       _a						:  Any=[0, 1, 2, 3]      ,       _a						:  List[Any]=4      ,       _a						:  Optional[Any]=3_7      ,       _a						:  List[Any]="gelu"      ,       _a						:  Optional[Any]=0.1      ,       _a						:  str=0.1      ,       _a						:  Optional[int]=0.0_2      ,       _a						:  Optional[int]=3      ,       _a						:  str=False      ,       _a						:  Dict=None      ,       ):
													a__:       Any											=parent
													a__:       str											=batch_size
													a__:       int											=image_size
													a__:       str											=patch_size
													a__:       int											=num_channels
													a__:       List[Any]											=is_training
													a__:       Optional[Any]											=use_labels
													a__:       List[Any]											=hidden_size
													a__:       Dict											=num_hidden_layers
													a__:       Union[str, Any]											=backbone_out_indices
													a__:       List[Any]											=num_attention_heads
													a__:       Optional[Any]											=intermediate_size
													a__:       Tuple											=hidden_act
													a__:       int											=hidden_dropout_prob
													a__:       int											=attention_probs_dropout_prob
													a__:       Optional[int]											=initializer_range
													a__:       Any											=num_labels
													a__:       Dict											=scope
													a__:       List[Any]											=is_hybrid
													# sequence length of DPT = num_patches + 1 (we add 1 for the [CLS] token)
													a__:       List[str]											=(image_size // patch_size) ** 2
													a__:       List[Any]											=num_patches + 1


						def   _lowerCamelCase			( self						:  Optional[int]			):
													a__:       Any											=floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]			)

													a__:       Dict											=None
													if self.use_labels:
																				a__:       str											=ids_tensor([self.batch_size, self.image_size, self.image_size]      ,       self.num_labels			)

													a__:       int											=self.get_config()

													return config, pixel_values, labels


						def   _lowerCamelCase			( self						:  Optional[Any]			):
													return DPTConfig(
													    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      ,       backbone_out_indices=self.backbone_out_indices      ,       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=a_      ,       initializer_range=self.initializer_range      ,       is_hybrid=self.is_hybrid      ,       )


						def   _lowerCamelCase			( self						:  List[str]      ,       _a						:  Optional[int]      ,       _a						:  List[str]      ,       _a						:  str			):
													a__:       List[str]											=DPTModel(config=a_			)
													model.to(a_			)
													model.eval()
													a__:       str											=model(a_			)
													self.parent.assertEqual(result.last_hidden_state.shape      ,       (self.batch_size, self.seq_length, self.hidden_size)			)


						def   _lowerCamelCase			( self						:  Dict      ,       _a						:  str      ,       _a						:  Any      ,       _a						:  Union[str, Any]			):
													a__:       Optional[int]											=self.num_labels
													a__:       Dict											=DPTForDepthEstimation(a_			)
													model.to(a_			)
													model.eval()
													a__:       Optional[Any]											=model(a_			)
													self.parent.assertEqual(result.predicted_depth.shape      ,       (self.batch_size, self.image_size, self.image_size)			)


						def   _lowerCamelCase			( self						:  List[Any]      ,       _a						:  List[str]      ,       _a						:  Optional[int]      ,       _a						:  int			):
													a__:       List[str]											=self.num_labels
													a__:       int											=DPTForSemanticSegmentation(a_			)
													model.to(a_			)
													model.eval()
													a__:       int											=model(a_      ,       labels=a_			)
													self.parent.assertEqual(
													    result.logits.shape      ,       (self.batch_size, self.num_labels, self.image_size, self.image_size)			)







						def   _lowerCamelCase			( self						:  str			):
													a__:       List[Any]											=self.prepare_config_and_inputs()
													a__							,			a__							,			a__:       Dict											=config_and_inputs
													a__:       Any											={"pixel_values": pixel_values}
													return config, inputs_dict


@require_torch
class 	lowerCamelCase__ (		_a				,					_a				,					unittest.TestCase   ):
						_lowerCAmelCase								= (DPTModel, DPTForDepthEstimation, DPTForSemanticSegmentation) if is_torch_available() else ()
						_lowerCAmelCase								= (
						    {
						        '''depth-estimation''': DPTForDepthEstimation,
						        '''feature-extraction''': DPTModel,
						        '''image-segmentation''': DPTForSemanticSegmentation,
						    }
						    if is_torch_available()
						    else {}
						)

						_lowerCAmelCase								= False
						_lowerCAmelCase								= False
						_lowerCAmelCase								= False


						def   _lowerCamelCase			( self						:  Optional[int]			):
													a__:       Any											=DPTModelTester(self			)
													a__:       Dict											=ConfigTester(self      ,       config_class=a_      ,       has_text_modality=a_      ,       hidden_size=3_7			)


						def   _lowerCamelCase			( self						:  List[str]			):
													self.config_tester.run_common_tests()


						@unittest.skip(reason="DPT does not use inputs_embeds"			)
						def   _lowerCamelCase			( self						:  Dict			):
													pass


						def   _lowerCamelCase			( self						:  str			):
													a__							,			a__:       str											=self.model_tester.prepare_config_and_inputs_for_common()

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


						def   _lowerCamelCase			( self						:  Optional[Any]			):
													a__							,			a__:       str											=self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																				a__:       Dict											=model_class(a_			)
																				a__:       List[str]											=inspect.signature(model.forward			)
																				# signature.parameters is an OrderedDict => so arg_names order is deterministic
																				a__:       Tuple											=[*signature.parameters.keys()]

																				a__:       int											=["pixel_values"]
																				self.assertListEqual(arg_names[:1]      ,       a_			)


						def   _lowerCamelCase			( self						:  List[Any]			):
													a__:       List[Any]											=self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_model(*a_			)


						def   _lowerCamelCase			( self						:  Optional[int]			):
													a__:       Any											=self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_depth_estimation(*a_			)


						def   _lowerCamelCase			( self						:  List[Any]			):
													a__:       Dict											=self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_for_semantic_segmentation(*a_			)


						def   _lowerCamelCase			( self						:  Tuple			):
													for model_class in self.all_model_classes:
																				if model_class.__name__ == "DPTForDepthEstimation":
																											continue

																				a__							,			a__:       Tuple											=self.model_tester.prepare_config_and_inputs_for_common()
																				a__:       Dict											=True

																				if model_class in get_values(a_			):
																											continue

																				a__:       Tuple											=model_class(a_			)
																				model.to(a_			)
																				model.train()
																				a__:       List[str]											=self._prepare_for_class(a_      ,       a_      ,       return_labels=a_			)
																				a__:       Any											=model(**a_			).loss
																				loss.backward()


						def   _lowerCamelCase			( self						:  Optional[Any]			):
													for model_class in self.all_model_classes:
																				if model_class.__name__ == "DPTForDepthEstimation":
																											continue

																				a__							,			a__:       str											=self.model_tester.prepare_config_and_inputs_for_common()
																				a__:       Union[str, Any]											=False
																				a__:       List[str]											=True

																				if model_class in get_values(a_			) or not model_class.supports_gradient_checkpointing:
																											continue
																				a__:       List[str]											=model_class(a_			)
																				model.to(a_			)
																				model.gradient_checkpointing_enable()
																				model.train()
																				a__:       Union[str, Any]											=self._prepare_for_class(a_      ,       a_      ,       return_labels=a_			)
																				a__:       str											=model(**a_			).loss
																				loss.backward()


						def   _lowerCamelCase			( self						:  str			):
													a__							,			a__:       Any											=self.model_tester.prepare_config_and_inputs_for_common()

													a__:       Optional[Any]											=_config_zero_init(a_			)
													for model_class in self.all_model_classes:
																				a__:       Optional[Any]											=model_class(config=a_			)
																				# Skip the check for the backbone
																				a__:       Union[str, Any]											=[]
																				for name, module in model.named_modules():
																											if module.__class__.__name__ == "DPTViTHybridEmbeddings":
																																		a__:       int											=[F"{name}.{key}" for key in module.state_dict().keys()]
																																		break

																				for name, param in model.named_parameters():
																											if param.requires_grad:
																																		if name in backbone_params:
																																									continue
																																		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"      ,       )


						@unittest.skip("Will be fixed soon by reducing the size of the model used for common tests."			)
						def   _lowerCamelCase			( self						:  Any			):
													pass







						@slow
						def   _lowerCamelCase			( self						:  List[Any]			):
													for model_name in DPT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
																				a__:       List[Any]											=DPTModel.from_pretrained(a_			)
																				self.assertIsNotNone(a_			)





def 						__lowerCamelCase       (      ):
							a__:       Dict											=Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"						)
							return image







@require_torch
@require_vision
@slow
class 	lowerCamelCase__ (		unittest.TestCase   ):


						def   _lowerCamelCase			( self						:  Union[str, Any]			):
													a__:       Dict											=DPTImageProcessor.from_pretrained("Intel/dpt-large"			)
													a__:       Optional[int]											=DPTForDepthEstimation.from_pretrained("Intel/dpt-large"			).to(a_			)

													a__:       Optional[Any]											=prepare_img()
													a__:       List[str]											=image_processor(images=a_      ,       return_tensors="pt"			).to(a_			)

													# forward pass
													with torch.no_grad():
																				a__:       List[str]											=model(**a_			)
																				a__:       List[str]											=outputs.predicted_depth

													# verify the predicted depth
													a__:       Dict											=torch.Size((1, 3_8_4, 3_8_4)			)
													self.assertEqual(predicted_depth.shape      ,       a_			)

													a__:       List[str]											=torch.tensor(
													    [[6.3_1_9_9, 6.3_6_2_9, 6.4_1_4_8], [6.3_8_5_0, 6.3_6_1_5, 6.4_1_6_6], [6.3_5_1_9, 6.3_1_7_6, 6.3_5_7_5]]			).to(a_			)

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


						def   _lowerCamelCase			( self						:  Dict			):
													a__:       Optional[int]											=DPTImageProcessor.from_pretrained("Intel/dpt-large-ade"			)
													a__:       int											=DPTForSemanticSegmentation.from_pretrained("Intel/dpt-large-ade"			).to(a_			)

													a__:       Tuple											=prepare_img()
													a__:       List[str]											=image_processor(images=a_      ,       return_tensors="pt"			).to(a_			)

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

													# verify the logits
													a__:       Any											=torch.Size((1, 1_5_0, 4_8_0, 4_8_0)			)
													self.assertEqual(outputs.logits.shape      ,       a_			)

													a__:       Any											=torch.tensor(
													    [[4.0_4_8_0, 4.2_4_2_0, 4.4_3_6_0], [4.3_1_2_4, 4.5_6_9_3, 4.8_2_6_1], [4.5_7_6_8, 4.8_9_6_5, 5.2_1_6_3]]			).to(a_			)

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







						def   _lowerCamelCase			( self						:  Tuple			):
													a__:       Optional[Any]											=DPTImageProcessor.from_pretrained("Intel/dpt-large-ade"			)
													a__:       Tuple											=DPTForSemanticSegmentation.from_pretrained("Intel/dpt-large-ade"			).to(a_			)

													a__:       Union[str, Any]											=prepare_img()
													a__:       Optional[Any]											=image_processor(images=a_      ,       return_tensors="pt"			).to(a_			)

													# forward pass
													with torch.no_grad():
																				a__:       int											=model(**a_			)

													a__:       str											=outputs.logits.detach().cpu()

													a__:       Dict											=image_processor.post_process_semantic_segmentation(outputs=a_      ,       target_sizes=[(5_0_0, 3_0_0)]			)
													a__:       List[str]											=torch.Size((5_0_0, 3_0_0)			)
													self.assertEqual(segmentation[0].shape      ,       a_			)

													a__:       Tuple											=image_processor.post_process_semantic_segmentation(outputs=a_			)
													a__:       Dict											=torch.Size((4_8_0, 4_8_0)			)
													self.assertEqual(segmentation[0].shape      ,       a_			)










# This test is meant to be run in on an instance with TPUs like this:
#
#   python examples/pytorch/xla_spawn.py --num_cores=8 tests/test_trainer_tpu.py
#
# Replace 8 with the number of TPU cores you have.
#

import sys
from typing import Dict

from transformers import EvalPrediction, HfArgumentParser, TrainingArguments, is_torch_available
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 	lowerCamelCase__ (		_a   ):


								def __init__( self						:  Any      ,       _a						:  int = 1_0_1			):
															a__:       Dict											=length


								def __len__( self						:  Optional[Any]			):
															return self.length







								def __getitem__( self						:  Dict      ,       _a						:  int			):
															return i


		class 	lowerCamelCase__ :


								def __call__( self						:  Any      ,       _a						:  List[str]			):
															return {"input_ids": torch.tensor(_a			), "labels": torch.tensor(_a			)}


		class 	lowerCamelCase__ (		nn.Module   ):


								def __init__( self						:  Tuple			):
															super().__init__()
															# Add some (unused) params otherwise DDP will complain.
															a__:       int											=nn.Linear(1_2_0      ,       8_0			)







								def   _lowerCamelCase			( self						:  str      ,       _a						:  List[Any]      ,       _a						:  List[str]=None			):
															if labels is not None:
																						return torch.tensor(0.0      ,       device=input_ids.device			), input_ids
															else:
																						return input_ids





def 						__lowerCamelCase       (      ):
							a__:       List[str]											=HfArgumentParser((TrainingArguments,)						)
							sys.argv += ["--output_dir", "./examples"]
							a__:       int											=parser.parse_args_into_dataclasses()[0]

							logger.warning(
							    F"Process rank: {training_args.local_rank}, device: {training_args.device}, "
							    F"tpu_num_cores: {training_args.tpu_num_cores}"     , )

							# 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_001, 256, 15]:
														a__:       Optional[Any]											=DummyDataset(__magic_name__						)

														def compute_metrics(__magic_name__	:  EvalPrediction						) -> Dict:
																					a__:       List[str]											=list(range(len(__magic_name__						)						)						)
																					a__:       Union[str, Any]											=p.predictions.tolist() == sequential and p.label_ids.tolist() == sequential
																					return {"success": success}

														a__:       Union[str, Any]											=Trainer(
														    model=DummyModel()     , args=__magic_name__     , data_collator=DummyDataCollator()     , eval_dataset=__magic_name__     , compute_metrics=__magic_name__     , )
														a__:       str											=trainer.evaluate()
														logger.info(__magic_name__						)
														if metrics["eval_success"] is not True:
																					logger.error(__magic_name__						)
																					exit(1						)

														a__:       List[str]											=trainer.predict(__magic_name__						)
														logger.info(p.metrics						)
														if p.metrics["test_success"] is not True:
																					logger.error(p.metrics						)
																					exit(1						)

														a__:       Tuple											=2

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

														a__:       List[str]											=trainer.predict(__magic_name__						)
														logger.info(p.metrics						)
														if p.metrics["test_success"] is not True:
																					logger.error(p.metrics						)
																					exit(1						)

														a__:       List[str]											=None

							logger.info("🔥 All distributed tests successful"						)





def 						__lowerCamelCase       (      __magic_name__	:  Optional[Any]						):
							# For xla_spawn (TPUs)
							main()


if __name__ == "__main__":
		main()





