


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        snake_case_   :
					def __init__(      self   :Any			,__snake_case   :Optional[int]			,__snake_case   :Any=2			,__snake_case   :str=True			,__snake_case   :Optional[Any]=False			,__snake_case   :int=10			,__snake_case   :int=3			,__snake_case   :str=32 * 8			,__snake_case   :int=32 * 8			,__snake_case   :int=4			,__snake_case   :Union[str, Any]=64			,)					->    Optional[int]:
												a__             =     parent
												a__             =     batch_size
												a__             =     is_training
												a__             =     use_auxiliary_loss
												a__             =     num_queries
												a__             =     num_channels
												a__             =     min_size
												a__             =     max_size
												a__             =     num_labels
												a__             =     hidden_dim
												a__             =     hidden_dim
					def  lowerCamelCase__(      self   :int      )					->    Union[str, Any]:
												a__             =     floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]      ).to(
												    __snake_case      )

												a__             =     torch.ones([self.batch_size, self.min_size, self.max_size]			,device=__snake_case      )

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

												a__             =     self.get_config()
												return config, pixel_values, pixel_mask, mask_labels, class_labels
					def  lowerCamelCase__(      self   :List[Any]      )					->    int:
												a__             =     MaskaFormerConfig(
												    hidden_size=self.hidden_dim			,)
												a__             =     self.num_queries
												a__             =     self.num_labels

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

												a__             =     64
												a__             =     1_28
												a__             =     self.hidden_dim
												a__             =     self.hidden_dim
												a__             =     self.hidden_dim
												return config
					def  lowerCamelCase__(      self   :List[str]      )					->    Any:
												a__      ,					a__      ,					a__      ,					a__      ,					a__             =     self.prepare_config_and_inputs()
												a__             =     {'pixel_values': pixel_values, 'pixel_mask': pixel_mask}
												return config, inputs_dict
					def  lowerCamelCase__(      self   :Optional[int]			,__snake_case   :List[str]			,__snake_case   :List[Any]      )					->    List[str]:
												a__             =     output.encoder_hidden_states
												a__             =     output.pixel_decoder_hidden_states
												a__             =     output.transformer_decoder_hidden_states

												self.parent.assertTrue(len(__snake_case      )			,len(config.backbone_config.depths      )      )
												self.parent.assertTrue(len(__snake_case      )			,len(config.backbone_config.depths      )      )
												self.parent.assertTrue(len(__snake_case      )			,config.decoder_layers      )
					def  lowerCamelCase__(      self   :List[str]			,__snake_case   :Optional[Any]			,__snake_case   :Optional[int]			,__snake_case   :Optional[int]			,__snake_case   :List[Any]=False      )					->    int:
												with torch.no_grad():
																			a__             =     MaskaFormerModel(config=__snake_case      )
																			model.to(__snake_case      )
																			model.eval()

																			a__             =     model(pixel_values=__snake_case			,pixel_mask=__snake_case      )
																			a__             =     model(__snake_case			,output_hidden_states=__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(__snake_case			,__snake_case      )



					def  lowerCamelCase__(      self   :int			,__snake_case   :Dict			,__snake_case   :Union[str, Any]			,__snake_case   :List[Any]			,__snake_case   :int			,__snake_case   :Any      )					->    str:
												a__             =     MaskaFormerForUniversalSegmentation(config=__snake_case      )
												model.to(__snake_case      )
												model.eval()

												def comm_check_on_output(__snake_case   :str      ):
																			# 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():
																			a__             =     model(pixel_values=__snake_case			,pixel_mask=__snake_case      )
																			a__             =     model(__snake_case      )

																			comm_check_on_output(__snake_case      )

																			a__             =     model(
																			    pixel_values=__snake_case			,pixel_mask=__snake_case			,mask_labels=__snake_case			,class_labels=__snake_case      )

												comm_check_on_output(__snake_case      )

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







@require_torch
class        snake_case_   (lowerCamelCase_	,						lowerCamelCase_	,						unittest.TestCase      ):
					UpperCAmelCase__					:    List[str]          =			(MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
					UpperCAmelCase__					:    Union[str, Any]          =			{'''feature-extraction''': MaskaFormerModel} if is_torch_available() else {}

					UpperCAmelCase__					:    str          =			False
					UpperCAmelCase__					:    Optional[int]          =			False
					UpperCAmelCase__					:    int          =			False
					UpperCAmelCase__					:    List[Any]          =			False
					def  lowerCamelCase__(      self   :str      )					->    str:
												a__             =     MaskaFormerModelTester(self      )
												a__             =     ConfigTester(self			,config_class=__snake_case			,has_text_modality=__snake_case      )
					def  lowerCamelCase__(      self   :Dict      )					->    str:
												self.config_tester.run_common_tests()
					def  lowerCamelCase__(      self   :List[Any]      )					->    Optional[int]:
												a__      ,					a__             =     self.model_tester.prepare_config_and_inputs_for_common()
												self.model_tester.create_and_check_maskaformer_model(__snake_case			,**__snake_case			,output_hidden_states=__snake_case      )
					def  lowerCamelCase__(      self   :Tuple      )					->    str:
												a__             =     self.model_tester.prepare_config_and_inputs()
												self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*__snake_case      )
					@unittest.skip(reason='Mask2Former does not use inputs_embeds'      )
					def  lowerCamelCase__(      self   :str      )					->    int:
												pass
					@unittest.skip(reason='Mask2Former does not have a get_input_embeddings method'      )
					def  lowerCamelCase__(      self   :List[Any]      )					->    Union[str, Any]:
												pass
					@unittest.skip(reason='Mask2Former is not a generative model'      )
					def  lowerCamelCase__(      self   :Optional[Any]      )					->    str:
												pass
					@unittest.skip(reason='Mask2Former does not use token embeddings'      )
					def  lowerCamelCase__(      self   :Any      )					->    List[str]:
												pass
					@require_torch_multi_gpu
					@unittest.skip(
					    reason='Mask2Former has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'      )
					def  lowerCamelCase__(      self   :Optional[Any]      )					->    List[str]:
												pass
					@unittest.skip('Will be fixed soon by reducing the size of the model used for common tests.'      )
					def  lowerCamelCase__(      self   :int      )					->    str:
												pass
					def  lowerCamelCase__(      self   :Optional[Any]      )					->    Tuple:
												a__      ,					a__             =     self.model_tester.prepare_config_and_inputs_for_common()

												for model_class in self.all_model_classes:
																			a__             =     model_class(__snake_case      )
																			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]			,__snake_case      )
					@slow
					def  lowerCamelCase__(      self   :Any      )					->    str:
												for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
																			a__             =     MaskaFormerModel.from_pretrained(__snake_case      )
																			self.assertIsNotNone(__snake_case      )
					def  lowerCamelCase__(      self   :int      )					->    List[str]:
												a__             =     (self.model_tester.min_size,) * 2
												a__             =     {
												    'pixel_values': torch.randn((2, 3, *size)			,device=__snake_case      ),
												    'mask_labels': torch.randn((2, 10, *size)			,device=__snake_case      ),
												    'class_labels': torch.zeros(2			,10			,device=__snake_case      ).long(),
												}
												a__             =     self.model_tester.get_config()

												a__             =     MaskaFormerForUniversalSegmentation(__snake_case      ).to(__snake_case      )
												a__             =     model(**__snake_case      )
												self.assertTrue(outputs.loss is not None      )
					def  lowerCamelCase__(      self   :Tuple      )					->    str:
												a__      ,					a__             =     self.model_tester.prepare_config_and_inputs_for_common()
												self.model_tester.create_and_check_maskaformer_model(__snake_case			,**__snake_case			,output_hidden_states=__snake_case      )
					def  lowerCamelCase__(      self   :Union[str, Any]      )					->    int:
												a__      ,					a__             =     self.model_tester.prepare_config_and_inputs_for_common()

												for model_class in self.all_model_classes:
																			a__             =     model_class(__snake_case      ).to(__snake_case      )
																			a__             =     model(**__snake_case			,output_attentions=__snake_case      )
																			self.assertTrue(outputs.attentions is not None      )
					def  lowerCamelCase__(      self   :List[Any]      )					->    int:
												if not self.model_tester.is_training:
																			return

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

												a__             =     model_class(__snake_case      )
												model.to(__snake_case      )
												model.train()

												a__             =     model(__snake_case			,mask_labels=__snake_case			,class_labels=__snake_case      ).loss
												loss.backward()



					def  lowerCamelCase__(      self   :Tuple      )					->    Tuple:
												a__             =     self.all_model_classes[1]
												a__      ,					a__      ,					a__      ,					a__      ,					a__             =     self.model_tester.prepare_config_and_inputs()
												a__             =     True
												a__             =     True

												a__             =     model_class(__snake_case      ).to(__snake_case      )
												model.train()

												a__             =     model(__snake_case			,mask_labels=__snake_case			,class_labels=__snake_case      )

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

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

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

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

												outputs.loss.backward(retain_graph=__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      )





snake_case							:       str     			=      1e-4






def 							__lowercase   ( ):
							a__             =     Image.open('./tests/fixtures/tests_samples/COCO/000000039769.png'   )
							return image







@require_vision
@slow
class        snake_case_   (unittest.TestCase      ):
					@cached_property
					def  lowerCamelCase__(      self   :List[Any]      )					->    List[str]:
												return "facebook/mask2former-swin-small-coco-instance"
					@cached_property
					def  lowerCamelCase__(      self   :Optional[Any]      )					->    Union[str, Any]:
												return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints      ) if is_vision_available() else None
					def  lowerCamelCase__(      self   :Optional[Any]      )					->    Any:
												a__             =     MaskaFormerModel.from_pretrained(self.model_checkpoints      ).to(__snake_case      )
												a__             =     self.default_image_processor
												a__             =     prepare_img()
												a__             =     image_processor(__snake_case			,return_tensors='pt'      ).to(__snake_case      )
												a__             =     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(__snake_case			,(1, 3, 3_84, 3_84)      )

												with torch.no_grad():
																			a__             =     model(**__snake_case      )

												a__             =     torch.tensor(
												    [[-0.27_90, -1.07_17, -1.16_68], [-0.51_28, -0.31_28, -0.49_87], [-0.58_32, 0.19_71, -0.01_97]]      ).to(__snake_case      )
												self.assertTrue(
												    torch.allclose(
												        outputs.encoder_last_hidden_state[0, 0, :3, :3]			,__snake_case			,atol=__snake_case      )      )

												a__             =     torch.tensor(
												    [[0.89_73, 1.18_47, 1.17_76], [1.19_34, 1.50_40, 1.51_28], [1.11_53, 1.44_86, 1.49_51]]      ).to(__snake_case      )
												self.assertTrue(
												    torch.allclose(
												        outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]			,__snake_case			,atol=__snake_case      )      )

												a__             =     torch.tensor(
												    [[2.11_52, 1.70_00, -0.86_03], [1.58_08, 1.80_04, -0.93_53], [1.60_43, 1.74_95, -0.59_99]]      ).to(__snake_case      )
												self.assertTrue(
												    torch.allclose(
												        outputs.transformer_decoder_last_hidden_state[0, :3, :3]			,__snake_case			,atol=__snake_case      )      )
					def  lowerCamelCase__(      self   :Union[str, Any]      )					->    str:
												a__             =     MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints      ).to(__snake_case      ).eval()
												a__             =     self.default_image_processor
												a__             =     prepare_img()
												a__             =     image_processor(__snake_case			,return_tensors='pt'      ).to(__snake_case      )
												a__             =     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(__snake_case			,(1, 3, 3_84, 3_84)      )

												with torch.no_grad():
																			a__             =     model(**__snake_case      )
												# masks_queries_logits
												a__             =     outputs.masks_queries_logits
												self.assertEqual(
												    masks_queries_logits.shape			,(1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)      )
												a__             =     [
												    [-8.78_39, -9.00_56, -8.81_21],
												    [-7.41_04, -7.03_13, -6.54_01],
												    [-6.61_05, -6.34_27, -6.46_75],
												]
												a__             =     torch.tensor(__snake_case      ).to(__snake_case      )
												self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]			,__snake_case			,atol=__snake_case      )      )
												# class_queries_logits
												a__             =     outputs.class_queries_logits
												self.assertEqual(class_queries_logits.shape			,(1, model.config.num_queries, model.config.num_labels + 1)      )
												a__             =     torch.tensor(
												    [
												        [1.83_24, -8.08_35, -4.19_22],
												        [0.84_50, -9.00_50, -3.60_53],
												        [0.30_45, -7.72_93, -3.02_75],
												    ]      ).to(__snake_case      )
												self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]			,__snake_case			,atol=__snake_case      )      )



					def  lowerCamelCase__(      self   :List[str]      )					->    List[Any]:
												a__             =     MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints      ).to(__snake_case      ).eval()
												a__             =     self.default_image_processor

												a__             =     image_processor(
												    [np.zeros((3, 8_00, 13_33)      ), np.zeros((3, 8_00, 13_33)      )]			,segmentation_maps=[np.zeros((3_84, 3_84)      ).astype(np.floataa      ), np.zeros((3_84, 3_84)      ).astype(np.floataa      )]			,return_tensors='pt'			,)

												a__             =     inputs['pixel_values'].to(__snake_case      )
												a__             =     [el.to(__snake_case      ) for el in inputs['mask_labels']]
												a__             =     [el.to(__snake_case      ) for el in inputs['class_labels']]

												with torch.no_grad():
																			a__             =     model(**__snake_case      )

												self.assertTrue(outputs.loss is not None      )








import json
import os
import subprocess
import unittest
from ast import literal_eval

import pytest
from parameterized import parameterized_class

from . import is_sagemaker_available


if is_sagemaker_available():
						from sagemaker import Session, TrainingJobAnalytics
						from sagemaker.huggingface import HuggingFace







@pytest.mark.skipif(
    literal_eval(os.getenv('''TEST_SAGEMAKER'''	,						'''False'''      )      ) is not True	,						reason='''Skipping test because should only be run when releasing minor transformers version'''	,						)
@pytest.mark.usefixtures('''sm_env'''      )
@parameterized_class(
    [
        {
            '''framework''': '''pytorch''',
            '''script''': '''run_glue.py''',
            '''model_name_or_path''': '''distilbert-base-cased''',
            '''instance_type''': '''ml.g4dn.xlarge''',
            '''results''': {'''train_runtime''': 6_5_0, '''eval_accuracy''': 0.6, '''eval_loss''': 0.9},
        },
        {
            '''framework''': '''tensorflow''',
            '''script''': '''run_tf.py''',
            '''model_name_or_path''': '''distilbert-base-cased''',
            '''instance_type''': '''ml.g4dn.xlarge''',
            '''results''': {'''train_runtime''': 6_0_0, '''eval_accuracy''': 0.3, '''eval_loss''': 0.9},
        },
    ]      )
class        snake_case_   (unittest.TestCase      ):
					def  lowerCamelCase__(      self   :Dict      )					->    Union[str, Any]:
												if self.framework == "pytorch":
																			subprocess.run(
																			    F'cp ./examples/pytorch/text-classification/run_glue.py {self.env.test_path}/run_glue.py'.split()			,encoding='utf-8'			,check=__snake_case			,)
												assert hasattr(self			,'env'      )
					def  lowerCamelCase__(      self   :List[Any]			,__snake_case   :Optional[Any]=1      )					->    int:
												# creates estimator
												return HuggingFace(
												    entry_point=self.script			,source_dir=self.env.test_path			,role=self.env.role			,image_uri=self.env.image_uri			,base_job_name=F'{self.env.base_job_name}-single'			,instance_count=__snake_case			,instance_type=self.instance_type			,debugger_hook_config=__snake_case			,hyperparameters={**self.env.hyperparameters, 'model_name_or_path': self.model_name_or_path}			,metric_definitions=self.env.metric_definitions			,py_version='py36'			,)
					def  lowerCamelCase__(      self   :Dict			,__snake_case   :str      )					->    Optional[int]:
												TrainingJobAnalytics(__snake_case      ).export_csv(F'{self.env.test_path}/{job_name}_metrics.csv'      )



					def  lowerCamelCase__(      self   :Tuple      )					->    Union[str, Any]:
												# create estimator
												a__             =     self.create_estimator()

												# run training
												estimator.fit()

												# result dataframe
												a__             =     TrainingJobAnalytics(estimator.latest_training_job.name      ).dataframe()

												# extract kpis
												a__             =     list(result_metrics_df[result_metrics_df.metric_name == 'eval_accuracy']['value']      )
												a__             =     list(result_metrics_df[result_metrics_df.metric_name == 'eval_loss']['value']      )
												# get train time from SageMaker job, this includes starting, preprocessing, stopping
												a__             =     (
												    Session().describe_training_job(estimator.latest_training_job.name      ).get('TrainingTimeInSeconds'			,99_99_99      )
												)

												# assert kpis
												assert train_runtime <= self.results["train_runtime"]
												assert all(t >= self.results['eval_accuracy'] for t in eval_accuracy      )
												assert all(t <= self.results['eval_loss'] for t in eval_loss      )

												# dump tests result into json file to share in PR
												with open(F'{estimator.latest_training_job.name}.json'			,'w'      ) as outfile:
																			json.dump({'train_time': train_runtime, 'eval_accuracy': eval_accuracy, 'eval_loss': eval_loss}			,__snake_case      )




