
'''simple docstring'''





import json
import logging
import os
import re
import sys
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Union

import datasets
import numpy as np
import torch
import torchaudio
from packaging import version
from torch import nn

import transformers
from transformers import (
    HfArgumentParser,
    Trainer,
    TrainingArguments,
    WavaVecaCTCTokenizer,
    WavaVecaFeatureExtractor,
    WavaVecaForCTC,
    WavaVecaProcessor,
    is_apex_available,
    set_seed,
)
from transformers.trainer_utils import get_last_checkpoint, is_main_process


if is_apex_available():
				from apex import amp


if version.parse(version.parse(torch.__version__).base_version) >= version.parse("1.6"):
				lowercase__							:							List[Any]									=			True
				from torch.cuda.amp import autocast

lowercase__							:							Tuple									=			logging.getLogger(__name__)




def 	A_ (							snake_case							:   Any=None ,      snake_case							:   List[Any]=None )      ->    Optional[Any]:



						'''simple docstring'''

						return field(default_factory=lambda: default ,      metadata=snake_case )






@dataclass
class        SCREAMING_SNAKE_CASE__    :

							"""simple docstring"""






							_snake_case     		=  field(
							    metadata={'help': 'Path to pretrained model or model identifier from huggingface.co/models'}    )
							_snake_case     		=  field(
							    default=_lowerCAmelCase   ,		metadata={'help': 'Where do you want to store the pretrained models downloaded from huggingface.co'}   ,		)
							_snake_case     		=  field(
							    default=_lowerCAmelCase   ,		metadata={'help': 'Whether to freeze the feature extractor layers of the model.'}    )
							_snake_case     		=  field(
							    default=0.1   ,		metadata={'help': 'The dropout ratio for the attention probabilities.'}    )
							_snake_case     		=  field(
							    default=0.1   ,		metadata={'help': 'The dropout ratio for activations inside the fully connected layer.'}    )
							_snake_case     		=  field(
							    default=0.1   ,		metadata={
							        'help': 'The dropout probabilitiy for all fully connected layers in the embeddings, encoder, and pooler.'
							    }   ,		)
							_snake_case     		=  field(
							    default=0.1   ,		metadata={'help': 'The dropout probabilitiy for all 1D convolutional layers in feature extractor.'}   ,		)
							_snake_case     		=  field(
							    default=0.05   ,		metadata={
							        'help': (
							            'Propability of each feature vector along the time axis to be chosen as the start of the vector'
							            'span to be masked. Approximately ``mask_time_prob * sequence_length // mask_time_length`` feature'
							            'vectors will be masked along the time axis. This is only relevant if ``apply_spec_augment is True``.'
							        )
							    }   ,		)
							_snake_case     		=  field(default=0.0   ,		metadata={'help': 'The LayerDrop probability.'}    )






@dataclass
class        SCREAMING_SNAKE_CASE__    :

							"""simple docstring"""






							_snake_case     		=  field(
							    default=_lowerCAmelCase   ,		metadata={'help': 'The configuration name of the dataset to use (via the datasets library).'}    )
							_snake_case     		=  field(
							    default='train+validation'   ,		metadata={
							        'help': 'The name of the training data set split to use (via the datasets library). Defaults to \'train\''
							    }   ,		)
							_snake_case     		=  field(
							    default=_lowerCAmelCase   ,		metadata={'help': 'Overwrite the cached preprocessed datasets or not.'}    )
							_snake_case     		=  field(
							    default=_lowerCAmelCase   ,		metadata={'help': 'The number of processes to use for the preprocessing.'}   ,		)
							_snake_case     		=  field(
							    default=_lowerCAmelCase   ,		metadata={
							        'help': (
							            'For debugging purposes or quicker training, truncate the number of training examples to this '
							            'value if set.'
							        )
							    }   ,		)
							_snake_case     		=  field(
							    default=_lowerCAmelCase   ,		metadata={
							        'help': (
							            'For debugging purposes or quicker training, truncate the number of validation examples to this '
							            'value if set.'
							        )
							    }   ,		)
							_snake_case     		=  list_field(
							    default=[',', '?', '.', '!', '-', ';', ':', '""', '%', '\'', '"', '�']   ,		metadata={'help': 'A list of characters to remove from the transcripts.'}   ,		)






@dataclass
class        SCREAMING_SNAKE_CASE__    :

							"""simple docstring"""






							_snake_case     		=  42
							_snake_case     		=  True
							_snake_case     		=  None
							_snake_case     		=  None
							_snake_case     		=  None
							_snake_case     		=  None



							def __call__(  self     ,	SCREAMING_SNAKE_CASE_			)->   Dict[str, torch.Tensor]:





													'''simple docstring'''






													__UpperCamelCase       =			[{'''input_values''': feature['''input_values''']} for feature in features]
													__UpperCamelCase       =			[{'''input_ids''': feature['''labels''']} for feature in features]

													__UpperCamelCase       =			self.processor.pad(
													    _lowerCAmelCase     ,	padding=self.padding     ,	max_length=self.max_length     ,	pad_to_multiple_of=self.pad_to_multiple_of     ,	return_tensors='''pt'''     ,	)
													__UpperCamelCase       =			self.processor.pad(
													    labels=_lowerCAmelCase     ,	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       =			labels_batch['''input_ids'''].masked_fill(labels_batch.attention_mask.ne(1			)     ,	-100			)

													__UpperCamelCase       =			labels

													return batch






class        SCREAMING_SNAKE_CASE__    (							_lowerCAmelCase    ):

							"""simple docstring"""






							def     A__		(  self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_			)->   torch.Tensor:





													'''simple docstring'''






													model.train()
													__UpperCamelCase       =			self._prepare_inputs(_lowerCAmelCase			)

													if self.use_amp:
																			with autocast():
																									__UpperCamelCase       =			self.compute_loss(_lowerCAmelCase     ,	_lowerCAmelCase			)
													else:
																			__UpperCamelCase       =			self.compute_loss(_lowerCAmelCase     ,	_lowerCAmelCase			)

													if self.args.n_gpu > 1:
																			if model.module.config.ctc_loss_reduction == "mean":
																									__UpperCamelCase       =			loss.mean()
																			elif model.module.config.ctc_loss_reduction == "sum":
																									__UpperCamelCase       =			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       =			loss / self.args.gradient_accumulation_steps

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

													return loss.detach()





def 	A_ (							)      ->    str:



						'''simple docstring'''

						__UpperCamelCase       =			HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments) )
						if len(sys.argv ) == 2 and sys.argv[1].endswith('''.json''' ):
												# If we pass only one argument to the script and it's the path to a json file,
												# let's parse it to get our arguments.
												__UpperCamelCase  ,							__UpperCamelCase  ,							__UpperCamelCase       =			parser.parse_json_file(json_file=os.path.abspath(sys.argv[1] ) )
						else:
												__UpperCamelCase  ,							__UpperCamelCase  ,							__UpperCamelCase       =			parser.parse_args_into_dataclasses()

						# Detecting last checkpoint.
						__UpperCamelCase       =			None
						if os.path.isdir(training_args.output_dir ) and training_args.do_train and not training_args.overwrite_output_dir:
												__UpperCamelCase       =			get_last_checkpoint(training_args.output_dir )
												if last_checkpoint is None and len(os.listdir(training_args.output_dir ) ) > 0:
																		raise ValueError(
																		    f"Output directory ({training_args.output_dir}) already exists and is not empty. "
																		    '''Use --overwrite_output_dir to overcome.''' )
												elif last_checkpoint is not None:
																		logger.info(
																		    f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
																		    '''the `--output_dir` or add `--overwrite_output_dir` to train from scratch.''' )

    # Setup logging
						logging.basicConfig(
						    format='''%(asctime)s - %(levelname)s - %(name)s - %(message)s''' ,      datefmt='''%m/%d/%Y %H:%M:%S''' ,      handlers=[logging.StreamHandler(sys.stdout )] ,      )
						logger.setLevel(logging.INFO if is_main_process(training_args.local_rank ) else logging.WARN )

						# Log on each process the small summary:
						logger.warning(
						    f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
						    + f"distributed training: {bool(training_args.local_rank != -1 )}, 16-bits training: {training_args.fpaa}" )
						# Set the verbosity to info of the Transformers logger (on main process only):
						if is_main_process(training_args.local_rank ):
												transformers.utils.logging.set_verbosity_info()
						logger.info('''Training/evaluation parameters %s''' ,      snake_case )

						# Set seed before initializing model.
						set_seed(training_args.seed )

						# Get the datasets:
						__UpperCamelCase       =			datasets.load_dataset(
						    '''common_voice''' ,      data_args.dataset_config_name ,      split=data_args.train_split_name )
						__UpperCamelCase       =			datasets.load_dataset('''common_voice''' ,      data_args.dataset_config_name ,      split='''test''' )

						# Create and save tokenizer
						__UpperCamelCase       =			f"[{''.join(data_args.chars_to_ignore )}]"

						def remove_special_characters(snake_case							:   List[str] ):
												__UpperCamelCase       =			re.sub(snake_case ,      '''''' ,      batch['''sentence'''] ).lower() + ''' '''
												return batch

						__UpperCamelCase       =			train_dataset.map(snake_case ,      remove_columns=['''sentence'''] )
						__UpperCamelCase       =			eval_dataset.map(snake_case ,      remove_columns=['''sentence'''] )

						def extract_all_chars(snake_case							:   Union[str, Any] ):
												__UpperCamelCase       =			''' '''.join(batch['''text'''] )
												__UpperCamelCase       =			list(set(snake_case ) )
												return {"vocab": [vocab], "all_text": [all_text]}

						__UpperCamelCase       =			train_dataset.map(
						    snake_case ,      batched=snake_case ,      batch_size=-1 ,      keep_in_memory=snake_case ,      remove_columns=train_dataset.column_names ,      )
						__UpperCamelCase       =			train_dataset.map(
						    snake_case ,      batched=snake_case ,      batch_size=-1 ,      keep_in_memory=snake_case ,      remove_columns=eval_dataset.column_names ,      )

						__UpperCamelCase       =			list(set(vocab_train['''vocab'''][0] ) | set(vocab_test['''vocab'''][0] ) )
						__UpperCamelCase       =			{v: k for k, v in enumerate(snake_case )}
						__UpperCamelCase       =			vocab_dict[''' ''']
						del vocab_dict[" "]
						__UpperCamelCase       =			len(snake_case )
						__UpperCamelCase       =			len(snake_case )

						with open('''vocab.json''' ,      '''w''' ) as vocab_file:
												json.dump(snake_case ,      snake_case )

						# Load pretrained model and tokenizer
						#
						# Distributed training:
						# The .from_pretrained methods guarantee that only one local process can concurrently
						# download model & vocab.
						__UpperCamelCase       =			WavaVecaCTCTokenizer(
						    '''vocab.json''' ,      unk_token='''[UNK]''' ,      pad_token='''[PAD]''' ,      word_delimiter_token='''|''' ,      )
						__UpperCamelCase       =			WavaVecaFeatureExtractor(
						    feature_size=1 ,      sampling_rate=16000 ,      padding_value=0.0 ,      do_normalize=snake_case ,      return_attention_mask=snake_case )
						__UpperCamelCase       =			WavaVecaProcessor(feature_extractor=snake_case ,      tokenizer=snake_case )
						__UpperCamelCase       =			WavaVecaForCTC.from_pretrained(
						    model_args.model_name_or_path ,      cache_dir=model_args.cache_dir ,      activation_dropout=model_args.activation_dropout ,      attention_dropout=model_args.attention_dropout ,      hidden_dropout=model_args.hidden_dropout ,      feat_proj_dropout=model_args.feat_proj_dropout ,      mask_time_prob=model_args.mask_time_prob ,      gradient_checkpointing=training_args.gradient_checkpointing ,      layerdrop=model_args.layerdrop ,      ctc_loss_reduction='''mean''' ,      pad_token_id=processor.tokenizer.pad_token_id ,      vocab_size=len(processor.tokenizer ) ,      )

						if data_args.max_train_samples is not None:
												__UpperCamelCase       =			min(len(snake_case ) ,      data_args.max_train_samples )
												__UpperCamelCase       =			train_dataset.select(range(snake_case ) )

						if data_args.max_val_samples is not None:
												__UpperCamelCase       =			eval_dataset.select(range(data_args.max_val_samples ) )

						__UpperCamelCase       =			torchaudio.transforms.Resample(48000 ,      16000 )

						# Preprocessing the datasets.
						# We need to read the aduio files as arrays and tokenize the targets.
						def speech_file_to_array_fn(snake_case							:   Dict ):
												__UpperCamelCase  ,							__UpperCamelCase       =			torchaudio.load(batch['''path'''] )
												__UpperCamelCase       =			resampler(snake_case ).squeeze().numpy()
												__UpperCamelCase       =			16000
												__UpperCamelCase       =			batch['''text''']
												return batch

						__UpperCamelCase       =			train_dataset.map(
						    snake_case ,      remove_columns=train_dataset.column_names ,      num_proc=data_args.preprocessing_num_workers ,      )
						__UpperCamelCase       =			eval_dataset.map(
						    snake_case ,      remove_columns=eval_dataset.column_names ,      num_proc=data_args.preprocessing_num_workers ,      )

						def prepare_dataset(snake_case							:   Dict ):
												# 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       =			processor(
												    audio=batch['''speech'''] ,      text=batch['''target_text'''] ,      sampling_rate=batch['''sampling_rate'''][0] )
												batch.update(snake_case )
												return batch

						__UpperCamelCase       =			train_dataset.map(
						    snake_case ,      remove_columns=train_dataset.column_names ,      batch_size=training_args.per_device_train_batch_size ,      batched=snake_case ,      num_proc=data_args.preprocessing_num_workers ,      )
						__UpperCamelCase       =			eval_dataset.map(
						    snake_case ,      remove_columns=eval_dataset.column_names ,      batch_size=training_args.per_device_train_batch_size ,      batched=snake_case ,      num_proc=data_args.preprocessing_num_workers ,      )

						# Metric
						__UpperCamelCase       =			datasets.load_metric('''wer''' )

						def compute_metrics(snake_case							:   Optional[Any] ):
												__UpperCamelCase       =			pred.predictions
												__UpperCamelCase       =			np.argmax(snake_case ,      axis=-1 )

												__UpperCamelCase       =			processor.tokenizer.pad_token_id

												__UpperCamelCase       =			processor.batch_decode(snake_case )
												# we do not want to group tokens when computing the metrics
												__UpperCamelCase       =			processor.batch_decode(pred.label_ids ,      group_tokens=snake_case )

												__UpperCamelCase       =			wer_metric.compute(predictions=snake_case ,      references=snake_case )

												return {"wer": wer}

						if model_args.freeze_feature_extractor:
												model.freeze_feature_extractor()

						# Data collator
						__UpperCamelCase       =			DataCollatorCTCWithPadding(processor=snake_case ,      padding=snake_case )

						# Initialize our Trainer
						__UpperCamelCase       =			CTCTrainer(
						    model=snake_case ,      data_collator=snake_case ,      args=snake_case ,      compute_metrics=snake_case ,      train_dataset=train_dataset if training_args.do_train else None ,      eval_dataset=eval_dataset if training_args.do_eval else None ,      tokenizer=processor.feature_extractor ,      )

						# Training
						if training_args.do_train:
												if last_checkpoint is not None:
																		__UpperCamelCase       =			last_checkpoint
												elif os.path.isdir(model_args.model_name_or_path ):
																		__UpperCamelCase       =			model_args.model_name_or_path
												else:
																		__UpperCamelCase       =			None

												# Save the feature_extractor and the tokenizer
												if is_main_process(training_args.local_rank ):
																		processor.save_pretrained(training_args.output_dir )

												__UpperCamelCase       =			trainer.train(resume_from_checkpoint=snake_case )
												trainer.save_model()

												__UpperCamelCase       =			train_result.metrics
												__UpperCamelCase       =			(
												    data_args.max_train_samples if data_args.max_train_samples is not None else len(snake_case )
												)
												__UpperCamelCase       =			min(snake_case ,      len(snake_case ) )

												trainer.log_metrics('''train''' ,      snake_case )
												trainer.save_metrics('''train''' ,      snake_case )
												trainer.save_state()

						# Evaluation
						__UpperCamelCase       =			{}
						if training_args.do_eval:
												logger.info('''*** Evaluate ***''' )
												__UpperCamelCase       =			trainer.evaluate()
												__UpperCamelCase       =			data_args.max_val_samples if data_args.max_val_samples is not None else len(snake_case )
												__UpperCamelCase       =			min(snake_case ,      len(snake_case ) )

												trainer.log_metrics('''eval''' ,      snake_case )
												trainer.save_metrics('''eval''' ,      snake_case )

						return results


if __name__ == "__main__":
				main()






import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import MaskaFormerConfig, is_torch_available, is_vision_available
from transformers.testing_utils import require_torch, require_torch_multi_gpu, require_vision, slow, torch_device
from transformers.utils import cached_property

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
				import torch

				from transformers import MaskaFormerForUniversalSegmentation, MaskaFormerModel

				if is_vision_available():
								from transformers import MaskaFormerImageProcessor

if is_vision_available():
				from PIL import Image






class        SCREAMING_SNAKE_CASE__    :

							"""simple docstring"""






							def __init__(  self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_=2     ,	SCREAMING_SNAKE_CASE_=True     ,	SCREAMING_SNAKE_CASE_=False     ,	SCREAMING_SNAKE_CASE_=10     ,	SCREAMING_SNAKE_CASE_=3     ,	SCREAMING_SNAKE_CASE_=32 * 8     ,	SCREAMING_SNAKE_CASE_=32 * 8     ,	SCREAMING_SNAKE_CASE_=4     ,	SCREAMING_SNAKE_CASE_=64     ,	)->   Dict:





													'''simple docstring'''






													__UpperCamelCase       =			parent
													__UpperCamelCase       =			batch_size
													__UpperCamelCase       =			is_training
													__UpperCamelCase       =			use_auxiliary_loss
													__UpperCamelCase       =			num_queries
													__UpperCamelCase       =			num_channels
													__UpperCamelCase       =			min_size
													__UpperCamelCase       =			max_size
													__UpperCamelCase       =			num_labels
													__UpperCamelCase       =			hidden_dim
													__UpperCamelCase       =			hidden_dim



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





													'''simple docstring'''






													__UpperCamelCase       =			floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]			).to(
													    SCREAMING_SNAKE_CASE_			)

													__UpperCamelCase       =			torch.ones([self.batch_size, self.min_size, self.max_size]     ,	device=SCREAMING_SNAKE_CASE_			)

													__UpperCamelCase       =			(
													    torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]     ,	device=SCREAMING_SNAKE_CASE_			) > 0.5
													).float()
													__UpperCamelCase       =			(torch.rand((self.batch_size, self.num_labels)     ,	device=SCREAMING_SNAKE_CASE_			) > 0.5).long()

													__UpperCamelCase       =			self.get_config()
													return config, pixel_values, pixel_mask, mask_labels, class_labels



							def     A__		(  self			)->   Optional[Any]:





													'''simple docstring'''






													__UpperCamelCase       =			MaskaFormerConfig(
													    hidden_size=self.hidden_dim     ,	)
													__UpperCamelCase       =			self.num_queries
													__UpperCamelCase       =			self.num_labels

													__UpperCamelCase       =			[1, 1, 1, 1]
													__UpperCamelCase       =			self.num_channels

													__UpperCamelCase       =			64
													__UpperCamelCase       =			128
													__UpperCamelCase       =			self.hidden_dim
													__UpperCamelCase       =			self.hidden_dim
													__UpperCamelCase       =			self.hidden_dim
													return config



							def     A__		(  self			)->   Any:





													'''simple docstring'''






													__UpperCamelCase  ,							__UpperCamelCase  ,							__UpperCamelCase  ,							__UpperCamelCase  ,							__UpperCamelCase       =			self.prepare_config_and_inputs()
													__UpperCamelCase       =			{'''pixel_values''': pixel_values, '''pixel_mask''': pixel_mask}
													return config, inputs_dict



							def     A__		(  self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_			)->   Any:





													'''simple docstring'''






													__UpperCamelCase       =			output.encoder_hidden_states
													__UpperCamelCase       =			output.pixel_decoder_hidden_states
													__UpperCamelCase       =			output.transformer_decoder_hidden_states

													self.parent.assertTrue(len(SCREAMING_SNAKE_CASE_			)     ,	len(config.backbone_config.depths			)			)
													self.parent.assertTrue(len(SCREAMING_SNAKE_CASE_			)     ,	len(config.backbone_config.depths			)			)
													self.parent.assertTrue(len(SCREAMING_SNAKE_CASE_			)     ,	config.decoder_layers			)



							def     A__		(  self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_=False			)->   Tuple:





													'''simple docstring'''






													with torch.no_grad():
																			__UpperCamelCase       =			MaskaFormerModel(config=SCREAMING_SNAKE_CASE_			)
																			model.to(SCREAMING_SNAKE_CASE_			)
																			model.eval()

																			__UpperCamelCase       =			model(pixel_values=SCREAMING_SNAKE_CASE_     ,	pixel_mask=SCREAMING_SNAKE_CASE_			)
																			__UpperCamelCase       =			model(SCREAMING_SNAKE_CASE_     ,	output_hidden_states=SCREAMING_SNAKE_CASE_			)

													self.parent.assertEqual(
													    output.transformer_decoder_last_hidden_state.shape     ,	(self.batch_size, self.num_queries, self.hidden_dim)     ,	)
													# let's ensure the other two hidden state exists
													self.parent.assertTrue(output.pixel_decoder_last_hidden_state is not None			)
													self.parent.assertTrue(output.encoder_last_hidden_state is not None			)

													if output_hidden_states:
																			self.check_output_hidden_state(SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_			)



							def     A__		(  self     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_     ,	SCREAMING_SNAKE_CASE_			)->   List[str]:





													'''simple docstring'''






													__UpperCamelCase       =			MaskaFormerForUniversalSegmentation(config=SCREAMING_SNAKE_CASE_			)
													model.to(SCREAMING_SNAKE_CASE_			)
													model.eval()

													def comm_check_on_output(SCREAMING_SNAKE_CASE_			):
																			# let's still check that all the required stuff is there
																			self.parent.assertTrue(result.transformer_decoder_last_hidden_state is not None			)
																			self.parent.assertTrue(result.pixel_decoder_last_hidden_state is not None			)
																			self.parent.assertTrue(result.encoder_last_hidden_state is not None			)
																			# okay, now we need to check the logits shape
																			# due to the encoder compression, masks have a //4 spatial size
																			self.parent.assertEqual(
																			    result.masks_queries_logits.shape     ,	(self.batch_size, self.num_queries, self.min_size // 4, self.max_size // 4)     ,	)
																			# + 1 for null class
																			self.parent.assertEqual(
																			    result.class_queries_logits.shape     ,	(self.batch_size, self.num_queries, self.num_labels + 1)			)

													with torch.no_grad():
																			__UpperCamelCase       =			model(pixel_values=SCREAMING_SNAKE_CASE_     ,	pixel_mask=SCREAMING_SNAKE_CASE_			)
																			__UpperCamelCase       =			model(SCREAMING_SNAKE_CASE_			)

																			comm_check_on_output(SCREAMING_SNAKE_CASE_			)

																			__UpperCamelCase       =			model(
																			    pixel_values=SCREAMING_SNAKE_CASE_     ,	pixel_mask=SCREAMING_SNAKE_CASE_     ,	mask_labels=SCREAMING_SNAKE_CASE_     ,	class_labels=SCREAMING_SNAKE_CASE_			)

													comm_check_on_output(SCREAMING_SNAKE_CASE_			)

													self.parent.assertTrue(result.loss is not None			)
													self.parent.assertEqual(result.loss.shape     ,	torch.Size([1]			)			)






@require_torch
class        SCREAMING_SNAKE_CASE__    (							SCREAMING_SNAKE_CASE_   ,		SCREAMING_SNAKE_CASE_   ,		unittest.TestCase    ):

							"""simple docstring"""






							_snake_case     		=  (MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
							_snake_case     		=  {'feature-extraction': MaskaFormerModel} if is_torch_available() else {}

							_snake_case     		=  False
							_snake_case     		=  False
							_snake_case     		=  False
							_snake_case     		=  False



							def     A__		(  self			)->   Optional[Any]:





													'''simple docstring'''






													__UpperCamelCase       =			MaskaFormerModelTester(self			)
													__UpperCamelCase       =			ConfigTester(self     ,	config_class=SCREAMING_SNAKE_CASE_     ,	has_text_modality=SCREAMING_SNAKE_CASE_			)



							def     A__		(  self			)->   List[str]:





													'''simple docstring'''






													self.config_tester.run_common_tests()



							def     A__		(  self			)->   int:





													'''simple docstring'''






													__UpperCamelCase  ,							__UpperCamelCase       =			self.model_tester.prepare_config_and_inputs_for_common()
													self.model_tester.create_and_check_maskaformer_model(SCREAMING_SNAKE_CASE_     ,	**SCREAMING_SNAKE_CASE_     ,	output_hidden_states=SCREAMING_SNAKE_CASE_			)



							def     A__		(  self			)->   Tuple:





													'''simple docstring'''






													__UpperCamelCase       =			self.model_tester.prepare_config_and_inputs()
													self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*SCREAMING_SNAKE_CASE_			)



							@unittest.skip(reason='''Mask2Former does not use inputs_embeds'''			)
							def     A__		(  self			)->   Any:





													'''simple docstring'''






													pass



							@unittest.skip(reason='''Mask2Former does not have a get_input_embeddings method'''			)
							def     A__		(  self			)->   Union[str, Any]:





													'''simple docstring'''






													pass



							@unittest.skip(reason='''Mask2Former is not a generative model'''			)
							def     A__		(  self			)->   Optional[Any]:





													'''simple docstring'''






													pass



							@unittest.skip(reason='''Mask2Former does not use token embeddings'''			)
							def     A__		(  self			)->   List[str]:





													'''simple docstring'''






													pass



							@require_torch_multi_gpu
							@unittest.skip(
							    reason='''Mask2Former has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'''			)
							def     A__		(  self			)->   Dict:





													'''simple docstring'''






													pass



							@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''			)
							def     A__		(  self			)->   str:





													'''simple docstring'''






													pass



							def     A__		(  self			)->   str:





													'''simple docstring'''






													__UpperCamelCase  ,							__UpperCamelCase       =			self.model_tester.prepare_config_and_inputs_for_common()

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

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



							@slow
							def     A__		(  self			)->   Any:





													'''simple docstring'''






													for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
																			__UpperCamelCase       =			MaskaFormerModel.from_pretrained(SCREAMING_SNAKE_CASE_			)
																			self.assertIsNotNone(SCREAMING_SNAKE_CASE_			)



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





													'''simple docstring'''






													__UpperCamelCase       =			(self.model_tester.min_size,) * 2
													__UpperCamelCase       =			{
													    '''pixel_values''': torch.randn((2, 3, *size)     ,	device=SCREAMING_SNAKE_CASE_			),
													    '''mask_labels''': torch.randn((2, 10, *size)     ,	device=SCREAMING_SNAKE_CASE_			),
													    '''class_labels''': torch.zeros(2     ,	10     ,	device=SCREAMING_SNAKE_CASE_			).long(),
													}
													__UpperCamelCase       =			self.model_tester.get_config()

													__UpperCamelCase       =			MaskaFormerForUniversalSegmentation(SCREAMING_SNAKE_CASE_			).to(SCREAMING_SNAKE_CASE_			)
													__UpperCamelCase       =			model(**SCREAMING_SNAKE_CASE_			)
													self.assertTrue(outputs.loss is not None			)



							def     A__		(  self			)->   Tuple:





													'''simple docstring'''






													__UpperCamelCase  ,							__UpperCamelCase       =			self.model_tester.prepare_config_and_inputs_for_common()
													self.model_tester.create_and_check_maskaformer_model(SCREAMING_SNAKE_CASE_     ,	**SCREAMING_SNAKE_CASE_     ,	output_hidden_states=SCREAMING_SNAKE_CASE_			)



							def     A__		(  self			)->   Any:





													'''simple docstring'''






													__UpperCamelCase  ,							__UpperCamelCase       =			self.model_tester.prepare_config_and_inputs_for_common()

													for model_class in self.all_model_classes:
																			__UpperCamelCase       =			model_class(SCREAMING_SNAKE_CASE_			).to(SCREAMING_SNAKE_CASE_			)
																			__UpperCamelCase       =			model(**SCREAMING_SNAKE_CASE_     ,	output_attentions=SCREAMING_SNAKE_CASE_			)
																			self.assertTrue(outputs.attentions is not None			)



							def     A__		(  self			)->   Any:





													'''simple docstring'''






													if not self.model_tester.is_training:
																			return

													__UpperCamelCase       =			self.all_model_classes[1]
													__UpperCamelCase  ,							__UpperCamelCase  ,							__UpperCamelCase  ,							__UpperCamelCase  ,							__UpperCamelCase       =			self.model_tester.prepare_config_and_inputs()

													__UpperCamelCase       =			model_class(SCREAMING_SNAKE_CASE_			)
													model.to(SCREAMING_SNAKE_CASE_			)
													model.train()

													__UpperCamelCase       =			model(SCREAMING_SNAKE_CASE_     ,	mask_labels=SCREAMING_SNAKE_CASE_     ,	class_labels=SCREAMING_SNAKE_CASE_			).loss
													loss.backward()



							def     A__		(  self			)->   Tuple:





													'''simple docstring'''






													__UpperCamelCase       =			self.all_model_classes[1]
													__UpperCamelCase  ,							__UpperCamelCase  ,							__UpperCamelCase  ,							__UpperCamelCase  ,							__UpperCamelCase       =			self.model_tester.prepare_config_and_inputs()
													__UpperCamelCase       =			True
													__UpperCamelCase       =			True

													__UpperCamelCase       =			model_class(SCREAMING_SNAKE_CASE_			).to(SCREAMING_SNAKE_CASE_			)
													model.train()

													__UpperCamelCase       =			model(SCREAMING_SNAKE_CASE_     ,	mask_labels=SCREAMING_SNAKE_CASE_     ,	class_labels=SCREAMING_SNAKE_CASE_			)

													__UpperCamelCase       =			outputs.encoder_hidden_states[0]
													encoder_hidden_states.retain_grad()

													__UpperCamelCase       =			outputs.pixel_decoder_hidden_states[0]
													pixel_decoder_hidden_states.retain_grad()

													__UpperCamelCase       =			outputs.transformer_decoder_hidden_states[0]
													transformer_decoder_hidden_states.retain_grad()

													__UpperCamelCase       =			outputs.attentions[0]
													attentions.retain_grad()

													outputs.loss.backward(retain_graph=SCREAMING_SNAKE_CASE_			)

													self.assertIsNotNone(encoder_hidden_states.grad			)
													self.assertIsNotNone(pixel_decoder_hidden_states.grad			)
													self.assertIsNotNone(transformer_decoder_hidden_states.grad			)
													self.assertIsNotNone(attentions.grad			)





lowercase__							:							Any									=			1e-4




def 	A_ (							)      ->    List[Any]:



						'''simple docstring'''

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






@require_vision
@slow
class        SCREAMING_SNAKE_CASE__    (							unittest.TestCase    ):

							"""simple docstring"""






							@cached_property
							def     A__		(  self			)->   List[Any]:





													'''simple docstring'''






													return "facebook/mask2former-swin-small-coco-instance"



							@cached_property
							def     A__		(  self			)->   Any:





													'''simple docstring'''






													return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints			) if is_vision_available() else None



							def     A__		(  self			)->   List[str]:





													'''simple docstring'''






													__UpperCamelCase       =			MaskaFormerModel.from_pretrained(self.model_checkpoints			).to(SCREAMING_SNAKE_CASE_			)
													__UpperCamelCase       =			self.default_image_processor
													__UpperCamelCase       =			prepare_img()
													__UpperCamelCase       =			image_processor(SCREAMING_SNAKE_CASE_     ,	return_tensors='''pt'''			).to(SCREAMING_SNAKE_CASE_			)
													__UpperCamelCase       =			inputs['''pixel_values'''].shape
													# check size is divisible by 32
													self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0			)
													# check size
													self.assertEqual(SCREAMING_SNAKE_CASE_     ,	(1, 3, 384, 384)			)

													with torch.no_grad():
																			__UpperCamelCase       =			model(**SCREAMING_SNAKE_CASE_			)

													__UpperCamelCase       =			torch.tensor(
													    [[-0.2_7_9_0, -1.0_7_1_7, -1.1_6_6_8], [-0.5_1_2_8, -0.3_1_2_8, -0.4_9_8_7], [-0.5_8_3_2, 0.1_9_7_1, -0.0_1_9_7]]			).to(SCREAMING_SNAKE_CASE_			)
													self.assertTrue(
													    torch.allclose(
													        outputs.encoder_last_hidden_state[0, 0, :3, :3]     ,	SCREAMING_SNAKE_CASE_     ,	atol=SCREAMING_SNAKE_CASE_			)			)

													__UpperCamelCase       =			torch.tensor(
													    [[0.8_9_7_3, 1.1_8_4_7, 1.1_7_7_6], [1.1_9_3_4, 1.5_0_4_0, 1.5_1_2_8], [1.1_1_5_3, 1.4_4_8_6, 1.4_9_5_1]]			).to(SCREAMING_SNAKE_CASE_			)
													self.assertTrue(
													    torch.allclose(
													        outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]     ,	SCREAMING_SNAKE_CASE_     ,	atol=SCREAMING_SNAKE_CASE_			)			)

													__UpperCamelCase       =			torch.tensor(
													    [[2.1_1_5_2, 1.7_0_0_0, -0.8_6_0_3], [1.5_8_0_8, 1.8_0_0_4, -0.9_3_5_3], [1.6_0_4_3, 1.7_4_9_5, -0.5_9_9_9]]			).to(SCREAMING_SNAKE_CASE_			)
													self.assertTrue(
													    torch.allclose(
													        outputs.transformer_decoder_last_hidden_state[0, :3, :3]     ,	SCREAMING_SNAKE_CASE_     ,	atol=SCREAMING_SNAKE_CASE_			)			)



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





													'''simple docstring'''






													__UpperCamelCase       =			MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints			).to(SCREAMING_SNAKE_CASE_			).eval()
													__UpperCamelCase       =			self.default_image_processor
													__UpperCamelCase       =			prepare_img()
													__UpperCamelCase       =			image_processor(SCREAMING_SNAKE_CASE_     ,	return_tensors='''pt'''			).to(SCREAMING_SNAKE_CASE_			)
													__UpperCamelCase       =			inputs['''pixel_values'''].shape
													# check size is divisible by 32
													self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0			)
													# check size
													self.assertEqual(SCREAMING_SNAKE_CASE_     ,	(1, 3, 384, 384)			)

													with torch.no_grad():
																			__UpperCamelCase       =			model(**SCREAMING_SNAKE_CASE_			)
													# masks_queries_logits
													__UpperCamelCase       =			outputs.masks_queries_logits
													self.assertEqual(
													    masks_queries_logits.shape     ,	(1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)			)
													__UpperCamelCase       =			[
													    [-8.7_8_3_9, -9.0_0_5_6, -8.8_1_2_1],
													    [-7.4_1_0_4, -7.0_3_1_3, -6.5_4_0_1],
													    [-6.6_1_0_5, -6.3_4_2_7, -6.4_6_7_5],
													]
													__UpperCamelCase       =			torch.tensor(SCREAMING_SNAKE_CASE_			).to(SCREAMING_SNAKE_CASE_			)
													self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]     ,	SCREAMING_SNAKE_CASE_     ,	atol=SCREAMING_SNAKE_CASE_			)			)
													# class_queries_logits
													__UpperCamelCase       =			outputs.class_queries_logits
													self.assertEqual(class_queries_logits.shape     ,	(1, model.config.num_queries, model.config.num_labels + 1)			)
													__UpperCamelCase       =			torch.tensor(
													    [
													        [1.8_3_2_4, -8.0_8_3_5, -4.1_9_2_2],
													        [0.8_4_5_0, -9.0_0_5_0, -3.6_0_5_3],
													        [0.3_0_4_5, -7.7_2_9_3, -3.0_2_7_5],
													    ]			).to(SCREAMING_SNAKE_CASE_			)
													self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]     ,	SCREAMING_SNAKE_CASE_     ,	atol=SCREAMING_SNAKE_CASE_			)			)



							def     A__		(  self			)->   str:





													'''simple docstring'''






													__UpperCamelCase       =			MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints			).to(SCREAMING_SNAKE_CASE_			).eval()
													__UpperCamelCase       =			self.default_image_processor

													__UpperCamelCase       =			image_processor(
													    [np.zeros((3, 800, 1333)			), np.zeros((3, 800, 1333)			)]     ,	segmentation_maps=[np.zeros((384, 384)			).astype(np.floataa			), np.zeros((384, 384)			).astype(np.floataa			)]     ,	return_tensors='''pt'''     ,	)

													__UpperCamelCase       =			inputs['''pixel_values'''].to(SCREAMING_SNAKE_CASE_			)
													__UpperCamelCase       =			[el.to(SCREAMING_SNAKE_CASE_			) for el in inputs['''mask_labels''']]
													__UpperCamelCase       =			[el.to(SCREAMING_SNAKE_CASE_			) for el in inputs['''class_labels''']]

													with torch.no_grad():
																			__UpperCamelCase       =			model(**SCREAMING_SNAKE_CASE_			)

													self.assertTrue(outputs.loss is not None			)




