

"""simple docstring"""





import logging
import pathlib
import re
import sys
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Set, Union

import datasets
import librosa
import numpy as np
import torch
from lang_trans import arabic
from packaging import version
from torch import nn

from transformers import (
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaForCTC,
    WavaVecaProcessor,
    is_apex_available,
    trainer_utils,
)


if is_apex_available():
				from apex import amp

if version.parse(version.parse(torch.__version__).base_version) >= version.parse("""1.6"""):
				_SCREAMING_SNAKE_CASE							:	Tuple										=			True
				from torch.cuda.amp import autocast


_SCREAMING_SNAKE_CASE							:	Union[str, Any]										=			logging.getLogger(__name__)

@dataclass
class      __a :

			"""simple docstring"""

			SCREAMING_SNAKE_CASE_      =    field(
			    metadata={'help': 'Path to pretrained model or model identifier from huggingface.co/models'}     )
			SCREAMING_SNAKE_CASE_      =    field(
			    default=_a,     metadata={'help': 'Where do you want to store the pretrained models downloaded from huggingface.co'},     )
			SCREAMING_SNAKE_CASE_      =    field(
			    default=_a,     metadata={'help': 'Whether to freeze the feature extractor layers of the model.'}     )
			SCREAMING_SNAKE_CASE_      =    field(
			    default=_a,     metadata={'help': 'Whether to log verbose messages or not.'},     )



def 							_lowerCAmelCase   (    UpperCAmelCase  :				Dict   ,				UpperCAmelCase  :				Optional[int]					):


					'''simple docstring'''






					logging.basicConfig(
					    format='''%(asctime)s - %(levelname)s - %(name)s - %(message)s'''   ,				datefmt='''%m/%d/%Y %H:%M:%S'''   ,				handlers=[logging.StreamHandler(sys.stdout					)]   ,				)
					UpperCamelCase__		:     Union[str, Any]								=logging.WARNING
					if model_args.verbose_logging:
										UpperCamelCase__		:     Optional[int]								=logging.DEBUG
					elif trainer_utils.is_main_process(training_args.local_rank					):
										UpperCamelCase__		:     List[str]								=logging.INFO
					logger.setLevel(lowerCamelCase__					)

@dataclass
class      __a :

			"""simple docstring"""

			SCREAMING_SNAKE_CASE_      =    field(
			    default=_a,     metadata={'help': 'The name of the dataset to use (via the datasets library).'}     )
			SCREAMING_SNAKE_CASE_      =    field(
			    default=_a,     metadata={'help': 'The configuration name of the dataset to use (via the datasets library).'}     )
			SCREAMING_SNAKE_CASE_      =    field(
			    default='train',     metadata={
			        'help': 'The name of the training data set split to use (via the datasets library). Defaults to \'train\''
			    },     )
			SCREAMING_SNAKE_CASE_      =    field(
			    default='validation',     metadata={
			        'help': (
			            'The name of the validation data set split to use (via the datasets library). Defaults to \'validation\''
			        )
			    },     )
			SCREAMING_SNAKE_CASE_      =    field(
			    default='text',     metadata={'help': 'Column in the dataset that contains label (target text). Defaults to \'text\''},     )
			SCREAMING_SNAKE_CASE_      =    field(
			    default='file',     metadata={'help': 'Column in the dataset that contains speech file path. Defaults to \'file\''},     )
			SCREAMING_SNAKE_CASE_      =    field(
			    default=_a,     metadata={'help': 'Resample loaded audio to target feature extractor\'s sampling rate or not.'},     )
			SCREAMING_SNAKE_CASE_      =    field(
			    default=_a,     metadata={'help': 'Filters out examples longer than specified. Defaults to no filtering.'},     )
			SCREAMING_SNAKE_CASE_      =    field(
			    default='librispeech',     metadata={
			        'help': (
			            'Orthography used for normalization and tokenization: \'librispeech\' (default), \'timit\', or'
			            ' \'buckwalter\'.'
			        )
			    },     )
			SCREAMING_SNAKE_CASE_      =    field(
			    default=_a,     metadata={'help': 'Overwrite the cached preprocessed datasets or not.'}     )
			SCREAMING_SNAKE_CASE_      =    field(
			    default=_a,     metadata={'help': 'The number of processes to use for the preprocessing.'},     )



@dataclass
class      __a :

			"""simple docstring"""

			SCREAMING_SNAKE_CASE_      =    False
			SCREAMING_SNAKE_CASE_      =    None
			SCREAMING_SNAKE_CASE_      =    "|"
			SCREAMING_SNAKE_CASE_      =    field(default_factory=_a     )
			SCREAMING_SNAKE_CASE_      =    field(default_factory=_a     )
			SCREAMING_SNAKE_CASE_      =    None






			@classmethod
			def  _lowerCAmelCase      (				cls  :       Optional[Any]    ,	lowercase_  :       str    ):
								if name == "librispeech":
													return cls()
								if name == "timit":
													return cls(
													    do_lower_case=lowercase_    ,	translation_table=str.maketrans({'''-''': ''' '''}    )    ,	)
								if name == "buckwalter":
													UpperCamelCase__		:     Tuple								={
													    '''-''': ''' ''',  # sometimes used to represent pauses
													    '''^''': '''v''',  # fixing "tha" in arabic_speech_corpus dataset
													}
													return cls(
													    vocab_file=pathlib.Path(lowercase_    ).parent.joinpath('''vocab/buckwalter.json'''    )    ,	word_delimiter_token='''/'''    ,	translation_table=str.maketrans(lowercase_    )    ,	words_to_remove={'''sil'''}    ,	untransliterator=arabic.buckwalter.untransliterate    ,	)
								raise ValueError(f'''Unsupported orthography: \'{name}\'.'''    )






			def  _lowerCAmelCase      (				self  :       Any    ,	lowercase_  :       str    ):
								# TODO(elgeish) return a pipeline (e.g., from jiwer) instead? Or rely on branch predictor as is
								if len(self.translation_table    ) > 0:
													UpperCamelCase__		:     Dict								=text.translate(self.translation_table    )
								if len(self.words_to_remove    ) == 0:
													UpperCamelCase__		:     List[str]								=''' '''.join(text.split()    )  # clean up whitespaces
								else:
													UpperCamelCase__		:     str								=''' '''.join(w for w in text.split() if w not in self.words_to_remove    )  # and clean up whilespaces
								return text






			def  _lowerCAmelCase      (				self  :       int    ,	lowercase_  :       ModelArguments    ):
								UpperCamelCase__		:     Tuple								=WavaVecaFeatureExtractor.from_pretrained(
								    model_args.model_name_or_path    ,	cache_dir=model_args.cache_dir    )
								if self.vocab_file:
													UpperCamelCase__		:     List[Any]								=WavaVecaCTCTokenizer(
													    self.vocab_file    ,	cache_dir=model_args.cache_dir    ,	do_lower_case=self.do_lower_case    ,	word_delimiter_token=self.word_delimiter_token    ,	)
								else:
													UpperCamelCase__		:     List[Any]								=WavaVecaCTCTokenizer.from_pretrained(
													    model_args.model_name_or_path    ,	cache_dir=model_args.cache_dir    ,	do_lower_case=self.do_lower_case    ,	word_delimiter_token=self.word_delimiter_token    ,	)
								return WavaVecaProcessor(lowercase_    ,	lowercase_    )



@dataclass
class      __a :

			"""simple docstring"""

			SCREAMING_SNAKE_CASE_      =    42
			SCREAMING_SNAKE_CASE_      =    True
			SCREAMING_SNAKE_CASE_      =    None
			SCREAMING_SNAKE_CASE_      =    None
			SCREAMING_SNAKE_CASE_      =    None
			SCREAMING_SNAKE_CASE_      =    None






			def __call__(				self  :       Optional[int]    ,	lowercase_  :       List[Dict[str, Union[List[int], torch.Tensor]]]    ):
								# split inputs and labels since they have to be of different lenghts and need
								# different padding methods
								UpperCamelCase__		:     Optional[Any]								=[{'''input_values''': feature['''input_values''']} for feature in features]
								UpperCamelCase__		:     Optional[Any]								=[{'''input_ids''': feature['''labels''']} for feature in features]

								UpperCamelCase__		:     List[str]								=self.processor.pad(
								    lowercase_    ,	padding=self.padding    ,	max_length=self.max_length    ,	pad_to_multiple_of=self.pad_to_multiple_of    ,	return_tensors='''pt'''    ,	)
								UpperCamelCase__		:     Dict								=self.processor.pad(
								    labels=lowercase_    ,	padding=self.padding    ,	max_length=self.max_length_labels    ,	pad_to_multiple_of=self.pad_to_multiple_of_labels    ,	return_tensors='''pt'''    ,	)

								# replace padding with -100 to ignore loss correctly
								UpperCamelCase__		:     List[str]								=labels_batch['''input_ids'''].masked_fill(labels_batch.attention_mask.ne(1    )    ,	-100    )

								UpperCamelCase__		:     List[str]								=labels

								return batch



class      __a (   _a     ):

			"""simple docstring"""

			def  _lowerCAmelCase      (				self  :       Tuple    ,	lowercase_  :       nn.Module    ,	lowercase_  :       Dict[str, Union[torch.Tensor, Any]]    ):
								model.train()
								UpperCamelCase__		:     Optional[Any]								=self._prepare_inputs(lowercase_    )

								if self.use_amp:
													with autocast():
																		UpperCamelCase__		:     Optional[Any]								=self.compute_loss(lowercase_    ,	lowercase_    )
								else:
													UpperCamelCase__		:     str								=self.compute_loss(lowercase_    ,	lowercase_    )

								if self.args.n_gpu > 1:
													if model.module.config.ctc_loss_reduction == "mean":
																		UpperCamelCase__		:     Dict								=loss.mean()
													elif model.module.config.ctc_loss_reduction == "sum":
																		UpperCamelCase__		:     Union[str, Any]								=loss.sum() / (inputs['''labels'''] >= 0).sum()
													else:
																		raise ValueError(f'''{model.config.ctc_loss_reduction} is not valid. Choose one of [\'mean\', \'sum\']'''    )

								if self.args.gradient_accumulation_steps > 1:
													UpperCamelCase__		:     List[Any]								=loss / self.args.gradient_accumulation_steps

								if self.use_amp:
													self.scaler.scale(lowercase_    ).backward()
								elif self.use_apex:
													with amp.scale_loss(lowercase_    ,	self.optimizer    ) as scaled_loss:
																		scaled_loss.backward()
								elif self.deepspeed:
													self.deepspeed.backward(lowercase_    )
								else:
													loss.backward()

								return loss.detach()



def 							_lowerCAmelCase   (    ):


					'''simple docstring'''






					UpperCamelCase__		:     Union[str, Any]								=HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)					)

					UpperCamelCase__	,      UpperCamelCase__	,      UpperCamelCase__		:     Tuple								=parser.parse_args_into_dataclasses()
					configure_logger(lowerCamelCase__   ,				lowerCamelCase__					)

					UpperCamelCase__		:     int								=Orthography.from_name(data_args.orthography.lower()					)
					UpperCamelCase__		:     Union[str, Any]								=orthography.create_processor(lowerCamelCase__					)
					UpperCamelCase__		:     str								=WavaVecaForCTC.from_pretrained(
					    model_args.model_name_or_path   ,				cache_dir=model_args.cache_dir   ,				gradient_checkpointing=training_args.gradient_checkpointing   ,				vocab_size=len(processor.tokenizer					)   ,				)

					UpperCamelCase__		:     Optional[Any]								=datasets.load_dataset(
					    data_args.dataset_name   ,				data_args.dataset_config_name   ,				split=data_args.train_split_name					)
					UpperCamelCase__		:     str								=datasets.load_dataset(
					    data_args.dataset_name   ,				data_args.dataset_config_name   ,				split=data_args.validation_split_name					)

					UpperCamelCase__		:     Any								=datasets.load_metric('''wer'''					)
					UpperCamelCase__		:     Union[str, Any]								=processor.feature_extractor.sampling_rate if data_args.target_feature_extractor_sampling_rate else None
					UpperCamelCase__		:     List[Any]								=''''''.join(t for t in processor.tokenizer.get_vocab().keys() if len(lowerCamelCase__					) == 1					)
					UpperCamelCase__		:     str								=re.compile(  # remove characters not in vocabulary
					    rF'''[^\s{re.escape(lowerCamelCase__					)}]'''   ,				flags=re.IGNORECASE if processor.tokenizer.do_lower_case else 0   ,				)
					UpperCamelCase__		:     Optional[Any]								=[]

					def prepare_example(UpperCAmelCase  :				Tuple					):  # TODO(elgeish) make use of multiprocessing?
										UpperCamelCase__	,      UpperCamelCase__		:     int								=librosa.load(example[data_args.speech_file_column]   ,				sr=lowerCamelCase__					)
										if data_args.max_duration_in_seconds is not None:
															UpperCamelCase__		:     Any								=len(example['''speech''']					) / example['''sampling_rate''']
										# Normalize and clean up text; order matters!
										UpperCamelCase__		:     Optional[int]								=orthography.preprocess_for_training(example[data_args.target_text_column]					)
										UpperCamelCase__		:     List[str]								=vocabulary_text_cleaner.sub(''''''   ,				lowerCamelCase__					)
										if updated_text != example[data_args.target_text_column]:
															text_updates.append((example[data_args.target_text_column], updated_text)					)
															UpperCamelCase__		:     Union[str, Any]								=updated_text
										return example

					UpperCamelCase__		:     Union[str, Any]								=train_dataset.map(lowerCamelCase__   ,				remove_columns=[data_args.speech_file_column]					)
					UpperCamelCase__		:     Union[str, Any]								=val_dataset.map(lowerCamelCase__   ,				remove_columns=[data_args.speech_file_column]					)

					if data_args.max_duration_in_seconds is not None:

										def filter_by_max_duration(UpperCAmelCase  :				str					):
															return example["duration_in_seconds"] <= data_args.max_duration_in_seconds

										UpperCamelCase__		:     Any								=len(lowerCamelCase__					)
										UpperCamelCase__		:     Any								=len(lowerCamelCase__					)
										UpperCamelCase__		:     List[str]								=train_dataset.filter(lowerCamelCase__   ,				remove_columns=['''duration_in_seconds''']					)
										UpperCamelCase__		:     Tuple								=val_dataset.filter(lowerCamelCase__   ,				remove_columns=['''duration_in_seconds''']					)
										if len(lowerCamelCase__					) > old_train_size:
															logger.warning(
															    F'''Filtered out {len(lowerCamelCase__					) - old_train_size} train example(s) longer than'''
															    F''' {data_args.max_duration_in_seconds} second(s).'''					)
										if len(lowerCamelCase__					) > old_val_size:
															logger.warning(
															    F'''Filtered out {len(lowerCamelCase__					) - old_val_size} validation example(s) longer than'''
															    F''' {data_args.max_duration_in_seconds} second(s).'''					)
					logger.info(F'''Split sizes: {len(lowerCamelCase__					)} train and {len(lowerCamelCase__					)} validation.'''					)

					logger.warning(F'''Updated {len(lowerCamelCase__					)} transcript(s) using \'{data_args.orthography}\' orthography rules.'''					)
					if logger.isEnabledFor(logging.DEBUG					):
										for original_text, updated_text in text_updates:
															logger.debug(F'''Updated text: \"{original_text}\" -> \"{updated_text}\"'''					)
					UpperCamelCase__		:     List[Any]								=None

					def prepare_dataset(UpperCAmelCase  :				Tuple					):
										# check that all files have the correct sampling rate
										assert (
										    len(set(batch['''sampling_rate''']					)					) == 1
										), F'''Make sure all inputs have the same sampling rate of {processor.feature_extractor.sampling_rate}.'''

										UpperCamelCase__		:     Dict								=processor(
										    audio=batch['''speech''']   ,				text=batch[data_args.target_text_column]   ,				sampling_rate=batch['''sampling_rate'''][0]					)
										batch.update(lowerCamelCase__					)
										return batch

					UpperCamelCase__		:     Dict								=train_dataset.map(
					    lowerCamelCase__   ,				batch_size=training_args.per_device_train_batch_size   ,				batched=lowerCamelCase__   ,				num_proc=data_args.preprocessing_num_workers   ,				)
					UpperCamelCase__		:     Optional[Any]								=val_dataset.map(
					    lowerCamelCase__   ,				batch_size=training_args.per_device_train_batch_size   ,				batched=lowerCamelCase__   ,				num_proc=data_args.preprocessing_num_workers   ,				)

					UpperCamelCase__		:     Dict								=DataCollatorCTCWithPadding(processor=lowerCamelCase__   ,				padding=lowerCamelCase__					)

					def compute_metrics(UpperCAmelCase  :				Tuple					):
										UpperCamelCase__		:     List[Any]								=pred.predictions
										UpperCamelCase__		:     Tuple								=np.argmax(lowerCamelCase__   ,				axis=-1					)

										UpperCamelCase__		:     Any								=processor.tokenizer.pad_token_id

										UpperCamelCase__		:     List[Any]								=processor.batch_decode(lowerCamelCase__					)
										# we do not want to group tokens when computing the metrics
										UpperCamelCase__		:     str								=processor.batch_decode(pred.label_ids   ,				group_tokens=lowerCamelCase__					)
										if logger.isEnabledFor(logging.DEBUG					):
															for reference, predicted in zip(lowerCamelCase__   ,				lowerCamelCase__					):
																				logger.debug(F'''reference: \"{reference}\"'''					)
																				logger.debug(F'''predicted: \"{predicted}\"'''					)
																				if orthography.untransliterator is not None:
																									logger.debug(F'''reference (untransliterated): \"{orthography.untransliterator(lowerCamelCase__					)}\"'''					)
																									logger.debug(F'''predicted (untransliterated): \"{orthography.untransliterator(lowerCamelCase__					)}\"'''					)

										UpperCamelCase__		:     Optional[Any]								=wer_metric.compute(predictions=lowerCamelCase__   ,				references=lowerCamelCase__					)

										return {"wer": wer}

					if model_args.freeze_feature_extractor:
										model.freeze_feature_extractor()

					UpperCamelCase__		:     Union[str, Any]								=CTCTrainer(
					    model=lowerCamelCase__   ,				data_collator=lowerCamelCase__   ,				args=lowerCamelCase__   ,				compute_metrics=lowerCamelCase__   ,				train_dataset=lowerCamelCase__   ,				eval_dataset=lowerCamelCase__   ,				tokenizer=processor.feature_extractor   ,				)

					trainer.train()


if __name__ == "__main__":
				main()



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




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




    def     __lowerCamelCase			(       self	: List[Any] )						->			str:
      a__  		=  floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size] )

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

      a__  		=  self.get_config()

      return config, pixel_values, labels




    def     __lowerCamelCase			(       self	: List[Any] )						->			Union[str, 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=lowerCAmelCase   ,   initializer_range=self.initializer_range   ,   is_hybrid=self.is_hybrid   ,   )




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




    def     __lowerCamelCase			(       self	: Union[str, Any]   ,   lowerCAmelCase	: Any   ,   lowerCAmelCase	: Dict   ,   lowerCAmelCase	: str )						->			Optional[Any]:
      a__  		=  self.num_labels
      a__  		=  DPTForDepthEstimation(lowerCAmelCase )
      model.to(lowerCAmelCase )
      model.eval()
      a__  		=  model(lowerCAmelCase )
      self.parent.assertEqual(result.predicted_depth.shape   ,   (self.batch_size, self.image_size, self.image_size) )




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




    def     __lowerCamelCase			(       self	: int )						->			Any:
      a__  		=  self.prepare_config_and_inputs()
      a__		,a__		,a__  		=  config_and_inputs
      a__  		=  {"pixel_values": pixel_values}
      return config, inputs_dict







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

    UpperCAmelCase_      :							List[str]    			=				False
    UpperCAmelCase_      :							int    			=				False
    UpperCAmelCase_      :							Optional[Any]    			=				False




    def     __lowerCamelCase			(       self	: List[Any] )						->			Any:
      a__  		=  DPTModelTester(self )
      a__  		=  ConfigTester(self   ,   config_class=lowerCAmelCase   ,   has_text_modality=lowerCAmelCase   ,   hidden_size=37 )




    def     __lowerCamelCase			(       self	: int )						->			List[Any]:
      self.config_tester.run_common_tests()




    @unittest.skip(reason="DPT does not use inputs_embeds" )
    def     __lowerCamelCase			(       self	: str )						->			Dict:
      pass




    def     __lowerCamelCase			(       self	: List[str] )						->			int:
      a__		,a__  		=  self.model_tester.prepare_config_and_inputs_for_common()

      for model_class in self.all_model_classes:
        a__  		=  model_class(lowerCAmelCase )
        self.assertIsInstance(model.get_input_embeddings()   ,   (nn.Module) )
        a__  		=  model.get_output_embeddings()
        self.assertTrue(x is None or isinstance(lowerCAmelCase   ,   nn.Linear ) )




    def     __lowerCamelCase			(       self	: Union[str, Any] )						->			Any:
      a__		,a__  		=  self.model_tester.prepare_config_and_inputs_for_common()

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

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




    def     __lowerCamelCase			(       self	: Optional[int] )						->			Optional[Any]:
      a__  		=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_model(*lowerCAmelCase )




    def     __lowerCamelCase			(       self	: Tuple )						->			List[str]:
      a__  		=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_depth_estimation(*lowerCAmelCase )




    def     __lowerCamelCase			(       self	: List[Any] )						->			Any:
      a__  		=  self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_for_semantic_segmentation(*lowerCAmelCase )




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

        a__		,a__  		=  self.model_tester.prepare_config_and_inputs_for_common()
        a__  		=  True

        if model_class in get_values(lowerCAmelCase ):
          continue

        a__  		=  model_class(lowerCAmelCase )
        model.to(lowerCAmelCase )
        model.train()
        a__  		=  self._prepare_for_class(lowerCAmelCase   ,   lowerCAmelCase   ,   return_labels=lowerCAmelCase )
        a__  		=  model(**lowerCAmelCase ).loss
        loss.backward()




    def     __lowerCamelCase			(       self	: Any )						->			Dict:
      for model_class in self.all_model_classes:
        if model_class.__name__ == "DPTForDepthEstimation":
          continue

        a__		,a__  		=  self.model_tester.prepare_config_and_inputs_for_common()
        a__  		=  False
        a__  		=  True

        if model_class in get_values(lowerCAmelCase ) or not model_class.supports_gradient_checkpointing:
          continue
        a__  		=  model_class(lowerCAmelCase )
        model.to(lowerCAmelCase )
        model.gradient_checkpointing_enable()
        model.train()
        a__  		=  self._prepare_for_class(lowerCAmelCase   ,   lowerCAmelCase   ,   return_labels=lowerCAmelCase )
        a__  		=  model(**lowerCAmelCase ).loss
        loss.backward()




    def     __lowerCamelCase			(       self	: List[str] )						->			int:
      a__		,a__  		=  self.model_tester.prepare_config_and_inputs_for_common()

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




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

def 					__lowercase						(    )					->int:

  '''simple docstring'''




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







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




    def     __lowerCamelCase			(       self	: Any )						->			Union[str, Any]:
      a__  		=  DPTImageProcessor.from_pretrained("Intel/dpt-large" )
      a__  		=  DPTForDepthEstimation.from_pretrained("Intel/dpt-large" ).to(lowerCAmelCase )

      a__  		=  prepare_img()
      a__  		=  image_processor(images=lowerCAmelCase   ,   return_tensors="pt" ).to(lowerCAmelCase )

      # forward pass
      with torch.no_grad():
        a__  		=  model(**lowerCAmelCase )
        a__  		=  outputs.predicted_depth

      # verify the predicted depth
      a__  		=  torch.Size((1, 384, 384) )
      self.assertEqual(predicted_depth.shape   ,   lowerCAmelCase )

      a__  		=  torch.tensor(
          [[6.3_199, 6.3_629, 6.4_148], [6.3_850, 6.3_615, 6.4_166], [6.3_519, 6.3_176, 6.3_575]] ).to(lowerCAmelCase )

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




    def     __lowerCamelCase			(       self	: Union[str, Any] )						->			str:
      a__  		=  DPTImageProcessor.from_pretrained("Intel/dpt-large-ade" )
      a__  		=  DPTForSemanticSegmentation.from_pretrained("Intel/dpt-large-ade" ).to(lowerCAmelCase )

      a__  		=  prepare_img()
      a__  		=  image_processor(images=lowerCAmelCase   ,   return_tensors="pt" ).to(lowerCAmelCase )

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

      # verify the logits
      a__  		=  torch.Size((1, 150, 480, 480) )
      self.assertEqual(outputs.logits.shape   ,   lowerCAmelCase )

      a__  		=  torch.tensor(
          [[4.0_480, 4.2_420, 4.4_360], [4.3_124, 4.5_693, 4.8_261], [4.5_768, 4.8_965, 5.2_163]] ).to(lowerCAmelCase )

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




    def     __lowerCamelCase			(       self	: str )						->			Optional[int]:
      a__  		=  DPTImageProcessor.from_pretrained("Intel/dpt-large-ade" )
      a__  		=  DPTForSemanticSegmentation.from_pretrained("Intel/dpt-large-ade" ).to(lowerCAmelCase )

      a__  		=  prepare_img()
      a__  		=  image_processor(images=lowerCAmelCase   ,   return_tensors="pt" ).to(lowerCAmelCase )

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

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

      a__  		=  image_processor.post_process_semantic_segmentation(outputs=lowerCAmelCase   ,   target_sizes=[(500, 300)] )
      a__  		=  torch.Size((500, 300) )
      self.assertEqual(segmentation[0].shape   ,   lowerCAmelCase )

      a__  		=  image_processor.post_process_semantic_segmentation(outputs=lowerCAmelCase )
      a__  		=  torch.Size((480, 480) )
      self.assertEqual(segmentation[0].shape   ,   lowerCAmelCase )

