




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

     def __init__(  self		,     _a		,     _a=2		,     _a=32		,     _a=16		,     _a=3		,     _a=True		,     _a=True		,     _a=32		,     _a=4		,     _a=[0, 1, 2, 3]		,     _a=4		,     _a=37		,     _a="gelu"		,     _a=0.1		,     _a=0.1		,     _a=0.02		,     _a=3		,     _a=[1, 384, 24, 24]		,     _a=True		,     _a=None		,     )						->							Optional[int]:
           _A							:     Optional[int]			=						parent
           _A							:     int			=						batch_size
           _A							:     Optional[int]			=						image_size
           _A							:     List[str]			=						patch_size
           _A							:     List[Any]			=						num_channels
           _A							:     str			=						is_training
           _A							:     List[str]			=						use_labels
           _A							:     Union[str, Any]			=						hidden_size
           _A							:     List[Any]			=						num_hidden_layers
           _A							:     List[Any]			=						backbone_out_indices
           _A							:     Any			=						num_attention_heads
           _A							:     str			=						intermediate_size
           _A							:     Optional[int]			=						hidden_act
           _A							:     Dict			=						hidden_dropout_prob
           _A							:     List[str]			=						attention_probs_dropout_prob
           _A							:     List[str]			=						initializer_range
           _A							:     Union[str, Any]			=						num_labels
           _A							:     int			=						backbone_featmap_shape
           _A							:     str			=						scope
           _A							:     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							:     str			=						num_patches + 1

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

           _A							:     Tuple			=						None
           if self.use_labels:
                 _A							:     List[Any]			=						ids_tensor([self.batch_size, self.image_size, self.image_size]		,     self.num_labels				)

           _A							:     Optional[Any]			=						self.get_config()

           return config, pixel_values, labels

     def 					a__	(  self				)						->							Optional[Any]:
           _A							:     Optional[Any]			=						{
               """global_padding""": """same""",
               """layer_type""": """bottleneck""",
               """depths""": [3, 4, 9],
               """out_features""": ["""stage1""", """stage2""", """stage3"""],
               """embedding_dynamic_padding""": True,
               """hidden_sizes""": [96, 192, 384, 768],
               """num_groups""": 2,
           }

           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=lowerCAmelCase_		,     initializer_range=self.initializer_range		,     is_hybrid=self.is_hybrid		,     backbone_config=lowerCAmelCase_		,     backbone_featmap_shape=self.backbone_featmap_shape		,     )

     def 					a__	(  self		,     _a		,     _a		,     _a				)						->							int:
           _A							:     str			=						DPTModel(config=lowerCAmelCase_				)
           model.to(lowerCAmelCase_				)
           model.eval()
           _A							:     Optional[Any]			=						model(lowerCAmelCase_				)
           self.parent.assertEqual(result.last_hidden_state.shape		,     (self.batch_size, self.seq_length, self.hidden_size)				)

     def 					a__	(  self		,     _a		,     _a		,     _a				)						->							List[Any]:
           _A							:     Union[str, Any]			=						self.num_labels
           _A							:     Optional[Any]			=						DPTForDepthEstimation(lowerCAmelCase_				)
           model.to(lowerCAmelCase_				)
           model.eval()
           _A							:     int			=						model(lowerCAmelCase_				)
           self.parent.assertEqual(result.predicted_depth.shape		,     (self.batch_size, self.image_size, self.image_size)				)

     def 					a__	(  self		,     _a		,     _a		,     _a				)						->							Any:
           _A							:     List[str]			=						self.num_labels
           _A							:     Union[str, Any]			=						DPTForSemanticSegmentation(lowerCAmelCase_				)
           model.to(lowerCAmelCase_				)
           model.eval()
           _A							:     Optional[Any]			=						model(lowerCAmelCase_		,     labels=lowerCAmelCase_				)
           self.parent.assertEqual(
               result.logits.shape		,     (self.batch_size, self.num_labels, self.image_size, self.image_size)				)



     def 					a__	(  self				)						->							int:
           _A							:     List[Any]			=						self.prepare_config_and_inputs()
           _A			,   _A			,   _A							:     int			=						config_and_inputs
           _A							:     Union[str, Any]			=						{"""pixel_values""": pixel_values}
           return config, inputs_dict







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

     _a							 =						False
     _a							 =						False
     _a							 =						False

     def 					a__	(  self				)						->							Optional[int]:
           _A							:     List[Any]			=						DPTModelTester(self				)
           _A							:     int			=						ConfigTester(self		,     config_class=lowerCAmelCase_		,     has_text_modality=lowerCAmelCase_		,     hidden_size=37				)

     def 					a__	(  self				)						->							Optional[Any]:
           self.config_tester.run_common_tests()

     @unittest.skip(reason="""DPT does not use inputs_embeds"""				)
     def 					a__	(  self				)						->							Optional[Any]:
           pass

     def 					a__	(  self				)						->							Tuple:
           _A			,   _A							:     Optional[int]			=						self.model_tester.prepare_config_and_inputs_for_common()

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

     def 					a__	(  self				)						->							Optional[Any]:
           _A			,   _A							:     List[Any]			=						self.model_tester.prepare_config_and_inputs_for_common()

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

                 _A							:     Dict			=						["""pixel_values"""]
                 self.assertListEqual(arg_names[:1]		,     lowerCAmelCase_				)

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

     def 					a__	(  self				)						->							Tuple:
           _A							:     Tuple			=						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_depth_estimation(*lowerCAmelCase_				)

     def 					a__	(  self				)						->							int:
           _A							:     Tuple			=						self.model_tester.prepare_config_and_inputs()
           self.model_tester.create_and_check_for_semantic_segmentation(*lowerCAmelCase_				)

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

                 _A			,   _A							:     Optional[Any]			=						self.model_tester.prepare_config_and_inputs_for_common()
                 _A							:     Optional[int]			=						True

                 if model_class in get_values(lowerCAmelCase_				):
                       continue

                 _A							:     Dict			=						model_class(lowerCAmelCase_				)
                 model.to(lowerCAmelCase_				)
                 model.train()
                 _A							:     Any			=						self._prepare_for_class(lowerCAmelCase_		,     lowerCAmelCase_		,     return_labels=lowerCAmelCase_				)
                 _A							:     Optional[Any]			=						model(**lowerCAmelCase_				).loss
                 loss.backward()

     def 					a__	(  self				)						->							List[str]:
           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							:     List[Any]			=						False
                 _A							:     Any			=						True

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

     def 					a__	(  self				)						->							Optional[int]:
           _A			,   _A							:     Any			=						self.model_tester.prepare_config_and_inputs_for_common()

           _A							:     List[str]			=						_config_zero_init(lowerCAmelCase_				)
           for model_class in self.all_model_classes:
                 _A							:     Union[str, Any]			=						model_class(config=lowerCAmelCase_				)
                 # Skip the check for the backbone
                 _A							:     List[Any]			=						[]
                 for name, module in model.named_modules():
                       if module.__class__.__name__ == "DPTViTHybridEmbeddings":
                             _A							:     Tuple			=						[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 					a__	(  self				)						->							Any:
           pass

     @slow
     def 					a__	(  self				)						->							List[str]:
           for model_name in DPT_PRETRAINED_MODEL_ARCHIVE_LIST[1:]:
                 _A							:     List[Any]			=						DPTModel.from_pretrained(lowerCAmelCase_				)
                 self.assertIsNotNone(lowerCAmelCase_				)



     def 					a__	(  self				)						->							Any:
           # We do this test only for DPTForDepthEstimation since it is the only model that uses readout_type
           _A			,   _A							:     List[Any]			=						self.model_tester.prepare_config_and_inputs_for_common()
           _A							:     Union[str, Any]			=						"""add"""
           with self.assertRaises(lowerCAmelCase_				):
                 _A							:     List[Any]			=						DPTForDepthEstimation(lowerCAmelCase_				)







def 				lowerCAmelCase_	(     ):
      _A							:     Any			=						Image.open("""./tests/fixtures/tests_samples/COCO/000000039769.png"""     )
      return image







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

     def 					a__	(  self				)						->							List[Any]:
           _A							:     List[Any]			=						DPTImageProcessor.from_pretrained("""Intel/dpt-hybrid-midas"""				)
           _A							:     str			=						DPTForDepthEstimation.from_pretrained("""Intel/dpt-hybrid-midas"""				).to(lowerCAmelCase_				)

           _A							:     str			=						prepare_img()
           _A							:     Tuple			=						image_processor(images=lowerCAmelCase_		,     return_tensors="""pt"""				).to(lowerCAmelCase_				)

           # forward pass
           with torch.no_grad():
                 _A							:     Any			=						model(**lowerCAmelCase_				)
                 _A							:     Optional[int]			=						outputs.predicted_depth

           # verify the predicted depth
           _A							:     Optional[Any]			=						torch.Size((1, 384, 384)				)
           self.assertEqual(predicted_depth.shape		,     lowerCAmelCase_				)

           _A							:     Optional[int]			=						torch.tensor(
               [[[5.6437, 5.6146, 5.6511], [5.4371, 5.5649, 5.5958], [5.5215, 5.5184, 5.5293]]]				).to(lowerCAmelCase_				)

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


from unittest.mock import patch

import pyspark

from datasets.packaged_modules.spark.spark import (
    Spark,
    SparkExamplesIterable,
    _generate_iterable_examples,
)

from ..utils import (
    require_dill_gt_0_3_2,
    require_not_windows,
)






def  snake_case  (       snake_case__			:Union[str, Any]   ,	snake_case__			:Dict)						->   Any:
       _A        =  []
       for part_id in partition_order:
              _A        =  df.where(F'''SPARK_PARTITION_ID() = {part_id}''').collect()
              for row_idx, row in enumerate(snake_case__):
                     expected_row_ids_and_row_dicts.append((F'''{part_id}_{row_idx}''', row.asDict()))
       return expected_row_ids_and_row_dicts






@require_not_windows
@require_dill_gt_0_3_2
def  snake_case  (       )						->   Optional[Any]:
       _A        =  pyspark.sql.SparkSession.builder.master("""local[*]""").appName("""pyspark""").getOrCreate()
       _A        =  spark.range(100).repartition(1)
       _A        =  Spark(snake_case__)
       # The id ints will be converted to Pyarrow int64s, so each row will be 8 bytes. Setting a max_shard_size of 16 means
       # that each partition can hold 2 rows.
       spark_builder._repartition_df_if_needed(max_shard_size=16)
       # Given that the dataframe has 100 rows and each partition has 2 rows, we expect 50 partitions.
       assert spark_builder.df.rdd.getNumPartitions() == 50






@require_not_windows
@require_dill_gt_0_3_2
def  snake_case  (       )						->   Union[str, Any]:
       _A        =  pyspark.sql.SparkSession.builder.master("""local[*]""").appName("""pyspark""").getOrCreate()
       _A        =  spark.range(10).repartition(2)
       _A        =  [1, 0]
       _A        =  _generate_iterable_examples(snake_case__   ,	snake_case__)  # Reverse the partitions.
       _A        =  _get_expected_row_ids_and_row_dicts_for_partition_order(snake_case__   ,	snake_case__)

       for i, (row_id, row_dict) in enumerate(generate_fn()):
              _A   ,  _A        =  expected_row_ids_and_row_dicts[i]
              assert row_id == expected_row_id
              assert row_dict == expected_row_dict






@require_not_windows
@require_dill_gt_0_3_2
def  snake_case  (       )						->   int:
       _A        =  pyspark.sql.SparkSession.builder.master("""local[*]""").appName("""pyspark""").getOrCreate()
       _A        =  spark.range(10).repartition(1)
       _A        =  SparkExamplesIterable(snake_case__)
       assert it.n_shards == 1
       for i, (row_id, row_dict) in enumerate(snake_case__):
              assert row_id == F'''0_{i}'''
              assert row_dict == {"id": i}






@require_not_windows
@require_dill_gt_0_3_2
def  snake_case  (       )						->   Union[str, Any]:
       _A        =  pyspark.sql.SparkSession.builder.master("""local[*]""").appName("""pyspark""").getOrCreate()
       _A        =  spark.range(30).repartition(3)
       # Mock the generator so that shuffle reverses the partition indices.
       with patch("""numpy.random.Generator""") as generator_mock:
              _A        =  lambda snake_case__: x.reverse()
              _A        =  _get_expected_row_ids_and_row_dicts_for_partition_order(snake_case__   ,	[2, 1, 0])

              _A        =  SparkExamplesIterable(snake_case__).shuffle_data_sources(snake_case__)
              assert shuffled_it.n_shards == 3
              for i, (row_id, row_dict) in enumerate(snake_case__):
                     _A   ,  _A        =  expected_row_ids_and_row_dicts[i]
                     assert row_id == expected_row_id
                     assert row_dict == expected_row_dict






@require_not_windows
@require_dill_gt_0_3_2
def  snake_case  (       )						->   List[str]:
       _A        =  pyspark.sql.SparkSession.builder.master("""local[*]""").appName("""pyspark""").getOrCreate()
       _A        =  spark.range(20).repartition(4)

       # Partitions 0 and 2
       _A        =  SparkExamplesIterable(snake_case__).shard_data_sources(worker_id=0   ,	num_workers=2)
       assert shard_it_a.n_shards == 2
       _A        =  _get_expected_row_ids_and_row_dicts_for_partition_order(snake_case__   ,	[0, 2])
       for i, (row_id, row_dict) in enumerate(snake_case__):
              _A   ,  _A        =  expected_row_ids_and_row_dicts_a[i]
              assert row_id == expected_row_id
              assert row_dict == expected_row_dict

       # Partitions 1 and 3
       _A        =  SparkExamplesIterable(snake_case__).shard_data_sources(worker_id=1   ,	num_workers=2)
       assert shard_it_a.n_shards == 2
       _A        =  _get_expected_row_ids_and_row_dicts_for_partition_order(snake_case__   ,	[1, 3])
       for i, (row_id, row_dict) in enumerate(snake_case__):
              _A   ,  _A        =  expected_row_ids_and_row_dicts_a[i]
              assert row_id == expected_row_id
              assert row_dict == expected_row_dict







@require_not_windows
@require_dill_gt_0_3_2
def  snake_case  (       )						->   Tuple:
       _A        =  pyspark.sql.SparkSession.builder.master("""local[*]""").appName("""pyspark""").getOrCreate()
       _A        =  spark.range(100).repartition(1)
       _A        =  Spark(snake_case__)
       # Choose a small max_shard_size for maximum partitioning.
       spark_builder._repartition_df_if_needed(max_shard_size=1)
       # The new number of partitions should not be greater than the number of rows.
       assert spark_builder.df.rdd.getNumPartitions() == 100




