
import inspect
import unittest

import numpy as np

from tests.test_modeling_common import floats_tensor
from transformers import DetrConfig, MaskFormerConfig, SwinConfig, 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 MaskFormerForInstanceSegmentation, MaskFormerModel

					if is_vision_available():
										from transformers import MaskFormerImageProcessor

if is_vision_available():
					from PIL import Image

class    lowerCAmelCase__     :






		def __init__(		self		:							Union[str, Any]							,			SCREAMING_SNAKE_CASE__		:							str							,			SCREAMING_SNAKE_CASE__		:							Optional[int]=2							,			SCREAMING_SNAKE_CASE__		:							Tuple=True							,			SCREAMING_SNAKE_CASE__		:							str=False							,			SCREAMING_SNAKE_CASE__		:							Optional[Any]=10							,			SCREAMING_SNAKE_CASE__		:							Union[str, Any]=3							,			SCREAMING_SNAKE_CASE__		:							Tuple=32 * 4							,			SCREAMING_SNAKE_CASE__		:							Optional[Any]=32 * 6							,			SCREAMING_SNAKE_CASE__		:							List[str]=4							,			SCREAMING_SNAKE_CASE__		:							List[Any]=32							,			)     ->							str:
						__lowerCamelCase    					=	parent
						__lowerCamelCase    					=	batch_size
						__lowerCamelCase    					=	is_training
						__lowerCamelCase    					=	use_auxiliary_loss
						__lowerCamelCase    					=	num_queries
						__lowerCamelCase    					=	num_channels
						__lowerCamelCase    					=	min_size
						__lowerCamelCase    					=	max_size
						__lowerCamelCase    					=	num_labels
						__lowerCamelCase    					=	mask_feature_size






		def  __A       (		self		:							List[Any]  )     ->							Optional[Any]:
						__lowerCamelCase    					=	floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]  ).to(
						    SCREAMING_SNAKE_CASE__  )

						__lowerCamelCase    					=	torch.ones([self.batch_size, self.min_size, self.max_size]							,			device=SCREAMING_SNAKE_CASE__  )

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

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






		def  __A       (		self		:							List[str]  )     ->							str:
						return MaskFormerConfig.from_backbone_and_decoder_configs(
						    backbone_config=SwinConfig(
						        depths=[1, 1, 1, 1]							,			)							,			decoder_config=DetrConfig(
						        decoder_ffn_dim=1_28							,			num_queries=self.num_queries							,			decoder_attention_heads=2							,			d_model=self.mask_feature_size							,			)							,			mask_feature_size=self.mask_feature_size							,			fpn_feature_size=self.mask_feature_size							,			num_channels=self.num_channels							,			num_labels=self.num_labels							,			)






		def  __A       (		self		:							Any  )     ->							str:
						__lowerCamelCase					,    __lowerCamelCase					,    __lowerCamelCase					,    __lowerCamelCase					,    __lowerCamelCase    					=	self.prepare_config_and_inputs()
						__lowerCamelCase    					=	{'''pixel_values''': pixel_values, '''pixel_mask''': pixel_mask}
						return config, inputs_dict






		def  __A       (		self		:							Optional[int]							,			SCREAMING_SNAKE_CASE__		:							List[str]							,			SCREAMING_SNAKE_CASE__		:							Optional[int]  )     ->							Union[str, Any]:
						__lowerCamelCase    					=	output.encoder_hidden_states
						__lowerCamelCase    					=	output.pixel_decoder_hidden_states
						__lowerCamelCase    					=	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_config.decoder_layers  )






		def  __A       (		self		:							Tuple							,			SCREAMING_SNAKE_CASE__		:							List[Any]							,			SCREAMING_SNAKE_CASE__		:							str							,			SCREAMING_SNAKE_CASE__		:							int							,			SCREAMING_SNAKE_CASE__		:							Any=False  )     ->							str:
						with torch.no_grad():
										__lowerCamelCase    					=	MaskFormerModel(config=SCREAMING_SNAKE_CASE__  )
										model.to(SCREAMING_SNAKE_CASE__  )
										model.eval()

										__lowerCamelCase    					=	model(pixel_values=SCREAMING_SNAKE_CASE__							,			pixel_mask=SCREAMING_SNAKE_CASE__  )
										__lowerCamelCase    					=	model(SCREAMING_SNAKE_CASE__							,			output_hidden_states=SCREAMING_SNAKE_CASE__  )
						# the correct shape of output.transformer_decoder_hidden_states ensure the correcteness of the
						# encoder and pixel decoder
						self.parent.assertEqual(
						    output.transformer_decoder_last_hidden_state.shape							,			(self.batch_size, self.num_queries, self.mask_feature_size)							,			)
						# 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		:							Union[str, Any]							,			SCREAMING_SNAKE_CASE__		:							Tuple							,			SCREAMING_SNAKE_CASE__		:							Tuple							,			SCREAMING_SNAKE_CASE__		:							Optional[int]							,			SCREAMING_SNAKE_CASE__		:							int							,			SCREAMING_SNAKE_CASE__		:							Dict  )     ->							Optional[int]:
						__lowerCamelCase    					=	MaskFormerForInstanceSegmentation(config=SCREAMING_SNAKE_CASE__  )
						model.to(SCREAMING_SNAKE_CASE__  )
						model.eval()

						def comm_check_on_output(SCREAMING_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():
										__lowerCamelCase    					=	model(pixel_values=SCREAMING_SNAKE_CASE__							,			pixel_mask=SCREAMING_SNAKE_CASE__  )
										__lowerCamelCase    					=	model(SCREAMING_SNAKE_CASE__  )

										comm_check_on_output(SCREAMING_SNAKE_CASE__  )

										__lowerCamelCase    					=	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    lowerCAmelCase__     (	__lowercase  ,	__lowercase  ,	unittest.TestCase			):
		a__       :						Optional[int]        =		(MaskFormerModel, MaskFormerForInstanceSegmentation) if is_torch_available() else ()
		a__       :						Dict        =		(
		    {"""feature-extraction""": MaskFormerModel, """image-segmentation""": MaskFormerForInstanceSegmentation}
		    if is_torch_available()
		    else {}
		)

		a__       :						Tuple        =		False
		a__       :						Union[str, Any]        =		False
		a__       :						List[Any]        =		False
		a__       :						Tuple        =		False






		def  __A       (		self		:							Any  )     ->							Dict:
						__lowerCamelCase    					=	MaskFormerModelTester(self  )
						__lowerCamelCase    					=	ConfigTester(self							,			config_class=SCREAMING_SNAKE_CASE__							,			has_text_modality=SCREAMING_SNAKE_CASE__  )






		def  __A       (		self		:							str  )     ->							Union[str, Any]:
						self.config_tester.run_common_tests()






		def  __A       (		self		:							Any  )     ->							str:
						__lowerCamelCase					,    __lowerCamelCase    					=	self.model_tester.prepare_config_and_inputs_for_common()
						self.model_tester.create_and_check_maskformer_model(SCREAMING_SNAKE_CASE__							,			**SCREAMING_SNAKE_CASE__							,			output_hidden_states=SCREAMING_SNAKE_CASE__  )






		def  __A       (		self		:							List[str]  )     ->							Optional[int]:
						__lowerCamelCase    					=	self.model_tester.prepare_config_and_inputs()
						self.model_tester.create_and_check_maskformer_instance_segmentation_head_model(*SCREAMING_SNAKE_CASE__  )






		@unittest.skip(reason='''MaskFormer does not use inputs_embeds'''  )
		def  __A       (		self		:							Optional[int]  )     ->							str:
						pass






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






		@unittest.skip(reason='''MaskFormer is not a generative model'''  )
		def  __A       (		self		:							Any  )     ->							List[str]:
						pass






		@unittest.skip(reason='''MaskFormer does not use token embeddings'''  )
		def  __A       (		self		:							Tuple  )     ->							Optional[int]:
						pass






		@require_torch_multi_gpu
		@unittest.skip(
		    reason='''MaskFormer has some layers using `add_module` which doesn\'t work well with `nn.DataParallel`'''  )
		def  __A       (		self		:							Optional[int]  )     ->							str:
						pass






		@unittest.skip('''Will be fixed soon by reducing the size of the model used for common tests.'''  )
		def  __A       (		self		:							Optional[Any]  )     ->							Any:
						pass






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

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

										__lowerCamelCase    					=	['''pixel_values''']
										self.assertListEqual(arg_names[:1]							,			SCREAMING_SNAKE_CASE__  )






		@slow
		def  __A       (		self		:							Tuple  )     ->							List[str]:
						for model_name in ["facebook/maskformer-swin-small-coco"]:
										__lowerCamelCase    					=	MaskFormerModel.from_pretrained(SCREAMING_SNAKE_CASE__  )
										self.assertIsNotNone(SCREAMING_SNAKE_CASE__  )






		def  __A       (		self		:							Dict  )     ->							Any:
						__lowerCamelCase    					=	(self.model_tester.min_size,) * 2
						__lowerCamelCase    					=	{
						    '''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(),
						}

						__lowerCamelCase    					=	MaskFormerForInstanceSegmentation(MaskFormerConfig()  ).to(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	model(**SCREAMING_SNAKE_CASE__  )
						self.assertTrue(outputs.loss is not None  )






		def  __A       (		self		:							Optional[Any]  )     ->							Union[str, Any]:
						__lowerCamelCase					,    __lowerCamelCase    					=	self.model_tester.prepare_config_and_inputs_for_common()
						self.model_tester.create_and_check_maskformer_model(SCREAMING_SNAKE_CASE__							,			**SCREAMING_SNAKE_CASE__							,			output_hidden_states=SCREAMING_SNAKE_CASE__  )






		def  __A       (		self		:							int  )     ->							Tuple:
						__lowerCamelCase					,    __lowerCamelCase    					=	self.model_tester.prepare_config_and_inputs_for_common()

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






		def  __A       (		self		:							Dict  )     ->							Any:
						if not self.model_tester.is_training:
										return
						# only MaskFormerForInstanceSegmentation has the loss
						__lowerCamelCase    					=	self.all_model_classes[1]
						__lowerCamelCase					,    __lowerCamelCase					,    __lowerCamelCase					,    __lowerCamelCase					,    __lowerCamelCase    					=	self.model_tester.prepare_config_and_inputs()

						__lowerCamelCase    					=	model_class(SCREAMING_SNAKE_CASE__  )
						model.to(SCREAMING_SNAKE_CASE__  )
						model.train()

						__lowerCamelCase    					=	model(SCREAMING_SNAKE_CASE__							,			mask_labels=SCREAMING_SNAKE_CASE__							,			class_labels=SCREAMING_SNAKE_CASE__  ).loss
						loss.backward()







		def  __A       (		self		:							Union[str, Any]  )     ->							List[str]:
						# only MaskFormerForInstanceSegmentation has the loss
						__lowerCamelCase    					=	self.all_model_classes[1]
						__lowerCamelCase					,    __lowerCamelCase					,    __lowerCamelCase					,    __lowerCamelCase					,    __lowerCamelCase    					=	self.model_tester.prepare_config_and_inputs()
						__lowerCamelCase    					=	True
						__lowerCamelCase    					=	True

						__lowerCamelCase    					=	model_class(SCREAMING_SNAKE_CASE__  )
						model.to(SCREAMING_SNAKE_CASE__  )
						model.train()

						__lowerCamelCase    					=	model(SCREAMING_SNAKE_CASE__							,			mask_labels=SCREAMING_SNAKE_CASE__							,			class_labels=SCREAMING_SNAKE_CASE__  )

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

						__lowerCamelCase    					=	outputs.pixel_decoder_hidden_states[0]
						pixel_decoder_hidden_states.retain_grad()
						# we requires_grad=True in inputs_embeds (line 2152), the original implementation don't
						__lowerCamelCase    					=	outputs.transformer_decoder_hidden_states[0]
						transformer_decoder_hidden_states.retain_grad()

						__lowerCamelCase    					=	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  )





SCREAMING_SNAKE_CASE__   :			Dict												=    1E-4




def  __magic_name__ (   )			->   str:
				__lowerCamelCase    					=	Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''							)
				return image







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






		@cached_property
		def  __A       (		self		:							str  )     ->							Optional[int]:
						return (
						    MaskFormerImageProcessor.from_pretrained('''facebook/maskformer-swin-small-coco'''  )
						    if is_vision_available()
						    else None
						)






		def  __A       (		self		:							Union[str, Any]  )     ->							Optional[Any]:
						__lowerCamelCase    					=	MaskFormerModel.from_pretrained('''facebook/maskformer-swin-small-coco'''  ).to(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	self.default_image_processor
						__lowerCamelCase    					=	prepare_img()
						__lowerCamelCase    					=	image_processor(SCREAMING_SNAKE_CASE__							,			return_tensors='''pt'''  ).to(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	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, 8_00, 10_88)  )

						with torch.no_grad():
										__lowerCamelCase    					=	model(**SCREAMING_SNAKE_CASE__  )

						__lowerCamelCase    					=	torch.tensor(
						    [[-0.0482, 0.9228, 0.4951], [-0.2547, 0.8017, 0.8527], [-0.0069, 0.3385, -0.0089]]  ).to(SCREAMING_SNAKE_CASE__  )
						self.assertTrue(
						    torch.allclose(
						        outputs.encoder_last_hidden_state[0, 0, :3, :3]							,			SCREAMING_SNAKE_CASE__							,			atol=SCREAMING_SNAKE_CASE__  )  )

						__lowerCamelCase    					=	torch.tensor(
						    [[-0.8422, -0.8434, -0.9718], [-1.0144, -0.5565, -0.4195], [-1.0038, -0.4484, -0.1961]]  ).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__  )  )

						__lowerCamelCase    					=	torch.tensor(
						    [[0.2852, -0.0159, 0.9735], [0.6254, 0.1858, 0.8529], [-0.0680, -0.4116, 1.8413]]  ).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		:							Optional[int]  )     ->							Any:
						__lowerCamelCase    					=	(
						    MaskFormerForInstanceSegmentation.from_pretrained('''facebook/maskformer-swin-small-coco'''  )
						    .to(SCREAMING_SNAKE_CASE__  )
						    .eval()
						)
						__lowerCamelCase    					=	self.default_image_processor
						__lowerCamelCase    					=	prepare_img()
						__lowerCamelCase    					=	image_processor(SCREAMING_SNAKE_CASE__							,			return_tensors='''pt'''  ).to(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	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, 8_00, 10_88)  )

						with torch.no_grad():
										__lowerCamelCase    					=	model(**SCREAMING_SNAKE_CASE__  )
						# masks_queries_logits
						__lowerCamelCase    					=	outputs.masks_queries_logits
						self.assertEqual(
						    masks_queries_logits.shape							,			(1, model.config.decoder_config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)							,			)
						__lowerCamelCase    					=	[
						    [-1.3737124, -1.7724937, -1.9364233],
						    [-1.5977281, -1.9867939, -2.1523695],
						    [-1.5795398, -1.9269832, -2.093942],
						]
						__lowerCamelCase    					=	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
						__lowerCamelCase    					=	outputs.class_queries_logits
						self.assertEqual(
						    class_queries_logits.shape							,			(1, model.config.decoder_config.num_queries, model.config.num_labels + 1)  )
						__lowerCamelCase    					=	torch.tensor(
						    [
						        [1.6_512e00, -5.2_572e00, -3.3_519e00],
						        [3.6_169e-02, -5.9_025e00, -2.9_313e00],
						        [1.0_766e-04, -7.7_630e00, -5.1_263e00],
						    ]  ).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		:							List[str]  )     ->							Dict:
						__lowerCamelCase    					=	(
						    MaskFormerForInstanceSegmentation.from_pretrained('''facebook/maskformer-resnet101-coco-stuff'''  )
						    .to(SCREAMING_SNAKE_CASE__  )
						    .eval()
						)
						__lowerCamelCase    					=	self.default_image_processor
						__lowerCamelCase    					=	prepare_img()
						__lowerCamelCase    					=	image_processor(SCREAMING_SNAKE_CASE__							,			return_tensors='''pt'''  ).to(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	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, 8_00, 10_88)  )

						with torch.no_grad():
										__lowerCamelCase    					=	model(**SCREAMING_SNAKE_CASE__  )
						# masks_queries_logits
						__lowerCamelCase    					=	outputs.masks_queries_logits
						self.assertEqual(
						    masks_queries_logits.shape							,			(1, model.config.decoder_config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)							,			)
						__lowerCamelCase    					=	[[-0.9046, -2.6366, -4.6062], [-3.4179, -5.7890, -8.8057], [-4.9179, -7.6560, -10.7711]]
						__lowerCamelCase    					=	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
						__lowerCamelCase    					=	outputs.class_queries_logits
						self.assertEqual(
						    class_queries_logits.shape							,			(1, model.config.decoder_config.num_queries, model.config.num_labels + 1)  )
						__lowerCamelCase    					=	torch.tensor(
						    [[4.7188, -3.2585, -2.8857], [6.6871, -2.9181, -1.2487], [7.2449, -2.2764, -2.1874]]  ).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		:							Optional[int]  )     ->							Any:
						__lowerCamelCase    					=	(
						    MaskFormerForInstanceSegmentation.from_pretrained('''facebook/maskformer-swin-small-coco'''  )
						    .to(SCREAMING_SNAKE_CASE__  )
						    .eval()
						)
						__lowerCamelCase    					=	self.default_image_processor

						__lowerCamelCase    					=	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'''							,			)

						__lowerCamelCase    					=	inputs['''pixel_values'''].to(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	[el.to(SCREAMING_SNAKE_CASE__  ) for el in inputs['''mask_labels''']]
						__lowerCamelCase    					=	[el.to(SCREAMING_SNAKE_CASE__  ) for el in inputs['''class_labels''']]

						with torch.no_grad():
										__lowerCamelCase    					=	model(**SCREAMING_SNAKE_CASE__  )

						self.assertTrue(outputs.loss is not None  )





import math
import os
import re
import sys
import unittest
from pathlib import Path
from typing import Tuple
from unittest.mock import patch

from parameterized import parameterized

from transformers.testing_utils import (
    CaptureStderr,
    ExtendSysPath,
    TestCasePlus,
    execute_subprocess_async,
    get_gpu_count,
    get_torch_dist_unique_port,
    require_apex,
    require_bitsandbytes,
    require_fairscale,
    require_torch,
    require_torch_gpu,
    require_torch_multi_gpu,
    require_torch_non_multi_gpu,
    slow,
)
from transformers.trainer_callback import TrainerState
from transformers.trainer_utils import set_seed


SCREAMING_SNAKE_CASE__   :			Optional[int]												=    os.path.abspath(os.path.dirname(__file__))
with ExtendSysPath(F'{bindir}/../../examples/pytorch/translation'):
					from run_translation import main  # noqa


set_seed(42)
SCREAMING_SNAKE_CASE__   :			Any												=    "sshleifer/student_marian_en_ro_6_1"
SCREAMING_SNAKE_CASE__   :			Tuple												=    "sshleifer/tiny-mbart"

@require_torch
class    lowerCAmelCase__     (	__lowercase			):






		def  __A       (		self		:							List[Any]							,			SCREAMING_SNAKE_CASE__		:							Tuple=False							,			SCREAMING_SNAKE_CASE__		:							Optional[Any]=None							,			SCREAMING_SNAKE_CASE__		:							Optional[Any]=True							,			SCREAMING_SNAKE_CASE__		:							str=True							,			SCREAMING_SNAKE_CASE__		:							Optional[Any]=True							,			SCREAMING_SNAKE_CASE__		:							Optional[int]=True							,			)     ->							Optional[int]:
						__lowerCamelCase    					=	self.run_trainer(
						    eval_steps=1							,			max_len=12							,			model_name=SCREAMING_SNAKE_CASE__							,			num_train_epochs=1							,			distributed=SCREAMING_SNAKE_CASE__							,			extra_args_str=SCREAMING_SNAKE_CASE__							,			predict_with_generate=SCREAMING_SNAKE_CASE__							,			do_train=SCREAMING_SNAKE_CASE__							,			do_eval=SCREAMING_SNAKE_CASE__							,			do_predict=SCREAMING_SNAKE_CASE__							,			)
						__lowerCamelCase    					=	TrainerState.load_from_json(os.path.join(SCREAMING_SNAKE_CASE__							,			'''trainer_state.json'''  )  ).log_history

						if not do_eval:
										return

						__lowerCamelCase    					=	[log for log in logs if '''eval_loss''' in log.keys()]

						__lowerCamelCase    					=	eval_metrics[0]
						if predict_with_generate:
										assert "eval_bleu" in first_step_stats

										__lowerCamelCase    					=	eval_metrics[-1]
										assert isinstance(last_step_stats['''eval_bleu''']							,			SCREAMING_SNAKE_CASE__  )
										assert not math.isnan(float(last_step_stats['''eval_loss''']  )  ), "eval_loss must not be `nan`"






		@require_torch_non_multi_gpu
		def  __A       (		self		:							Optional[int]  )     ->							int:
						self.run_seqaseq_quick()






		@require_torch_multi_gpu
		def  __A       (		self		:							int  )     ->							List[str]:
						self.run_seqaseq_quick(distributed=SCREAMING_SNAKE_CASE__  )






		@require_torch_multi_gpu
		def  __A       (		self		:							Optional[Any]  )     ->							Tuple:
						self.run_seqaseq_quick(distributed=SCREAMING_SNAKE_CASE__  )






		@unittest.skip('''Requires an update of the env running those tests'''  )
		@require_torch_multi_gpu
		@require_fairscale
		def  __A       (		self		:							Dict  )     ->							Tuple:
						self.run_seqaseq_quick(distributed=SCREAMING_SNAKE_CASE__							,			extra_args_str='''--sharded_ddp simple'''  )






		@unittest.skip('''Requires an update of the env running those tests'''  )
		@require_torch_multi_gpu
		@require_fairscale
		def  __A       (		self		:							Optional[int]  )     ->							List[str]:
						self.run_seqaseq_quick(distributed=SCREAMING_SNAKE_CASE__							,			extra_args_str='''--sharded_ddp simple --fp16'''  )






		@unittest.skip('''Requires an update of the env running those tests'''  )
		@require_torch_multi_gpu
		@require_fairscale
		def  __A       (		self		:							Tuple  )     ->							Any:
						self.run_seqaseq_quick(distributed=SCREAMING_SNAKE_CASE__							,			extra_args_str='''--sharded_ddp zero_dp_2'''							,			predict_with_generate=SCREAMING_SNAKE_CASE__  )






		@unittest.skip('''Requires an update of the env running those tests'''  )
		@require_torch_multi_gpu
		@require_fairscale
		def  __A       (		self		:							Dict  )     ->							Tuple:
						self.run_seqaseq_quick(
						    distributed=SCREAMING_SNAKE_CASE__							,			extra_args_str='''--sharded_ddp zero_dp_2 --fp16'''							,			predict_with_generate=SCREAMING_SNAKE_CASE__  )






		@require_apex
		@require_torch_gpu
		def  __A       (		self		:							Union[str, Any]  )     ->							List[str]:
						# XXX: apex breaks the trainer if it's run twice e.g. run_seq2seq.main() from the same
						# program and it breaks other tests that run from the same pytest worker, therefore until this is
						# sorted out it must be run only in an external program, that is distributed=True in this
						# test and only under one or more gpus - if we want cpu will need to make a special test
						#
						# specifically to the problem traced it to self.optimizer.step() - if it's run 2nd time via
						# 2nd main() call it botches the future eval.
						#
						self.run_seqaseq_quick(distributed=SCREAMING_SNAKE_CASE__							,			extra_args_str='''--fp16 --fp16_backend=apex'''  )
						# test 2nd time - was getting eval_loss': nan'
						# to reproduce the problem set distributed=False
						self.run_seqaseq_quick(distributed=SCREAMING_SNAKE_CASE__							,			extra_args_str='''--fp16 --fp16_backend=apex'''  )






		@parameterized.expand(['''base''', '''low''', '''high''', '''mixed''']  )
		@require_torch_multi_gpu
		def  __A       (		self		:							Any							,			SCREAMING_SNAKE_CASE__		:							List[str]  )     ->							Optional[Any]:
						# as each sub-test is slow-ish split into multiple sub-tests to avoid CI timeout
						__lowerCamelCase    					=	{
						    # test with the default log_level - should be info and thus log info once
						    '''base''': {'''extra_args_str''': '''''', '''n_matches''': 1},
						    # test with low log_level and log_level_replica - should be noisy on all processes
						    # now the info string should appear twice on 2 processes
						    '''low''': {'''extra_args_str''': '''--log_level debug --log_level_replica debug''', '''n_matches''': 2},
						    # test with high log_level and low log_level_replica
						    # now the info string should appear once only on the replica
						    '''high''': {'''extra_args_str''': '''--log_level error --log_level_replica debug''', '''n_matches''': 1},
						    # test with high log_level and log_level_replica - should be quiet on all processes
						    '''mixed''': {'''extra_args_str''': '''--log_level error --log_level_replica error''', '''n_matches''': 0},
						}

						__lowerCamelCase    					=	experiments[experiment_id]
						__lowerCamelCase    					=	{'''distributed''': True, '''predict_with_generate''': False, '''do_eval''': False, '''do_predict''': False}
						__lowerCamelCase    					=	'''Running training'''
						with CaptureStderr() as cl:
										self.run_seqaseq_quick(**SCREAMING_SNAKE_CASE__							,			extra_args_str=data['''extra_args_str''']  )
						__lowerCamelCase    					=	len(re.findall(SCREAMING_SNAKE_CASE__							,			cl.err  )  )
						self.assertEqual(SCREAMING_SNAKE_CASE__							,			data['''n_matches''']  )






		@slow
		def  __A       (		self		:							Any  )     ->							Optional[Any]:
						__lowerCamelCase    					=	self.run_trainer(
						    eval_steps=2							,			max_len=1_28							,			model_name=SCREAMING_SNAKE_CASE__							,			learning_rate=3e-4							,			num_train_epochs=10							,			distributed=SCREAMING_SNAKE_CASE__							,			)

						# Check metrics
						__lowerCamelCase    					=	TrainerState.load_from_json(os.path.join(SCREAMING_SNAKE_CASE__							,			'''trainer_state.json'''  )  ).log_history
						__lowerCamelCase    					=	[log for log in logs if '''eval_loss''' in log.keys()]
						__lowerCamelCase    					=	eval_metrics[0]
						__lowerCamelCase    					=	eval_metrics[-1]

						assert first_step_stats["eval_loss"] > last_step_stats["eval_loss"], "model learned nothing"
						assert isinstance(last_step_stats['''eval_bleu''']							,			SCREAMING_SNAKE_CASE__  )

						# test if do_predict saves generations and metrics
						__lowerCamelCase    					=	os.listdir(SCREAMING_SNAKE_CASE__  )
						__lowerCamelCase    					=	{os.path.basename(SCREAMING_SNAKE_CASE__  ) for p in contents}
						assert "generated_predictions.txt" in contents
						assert "predict_results.json" in contents






		@slow
		@require_bitsandbytes
		def  __A       (		self		:							Optional[int]  )     ->							str:
						from transformers.training_args import OptimizerNames

						def train_and_return_metrics(SCREAMING_SNAKE_CASE__		:							str  ) -> Tuple[int, float]:
										__lowerCamelCase    					=	'''--skip_memory_metrics 0'''

										__lowerCamelCase    					=	self.run_trainer(
										    max_len=1_28							,			model_name=SCREAMING_SNAKE_CASE__							,			learning_rate=3e-4							,			num_train_epochs=1							,			optim=SCREAMING_SNAKE_CASE__							,			distributed=SCREAMING_SNAKE_CASE__							,			extra_args_str=SCREAMING_SNAKE_CASE__							,			do_eval=SCREAMING_SNAKE_CASE__							,			do_predict=SCREAMING_SNAKE_CASE__							,			n_gpus_to_use=1							,			)

										# Check metrics
										__lowerCamelCase    					=	TrainerState.load_from_json(Path(SCREAMING_SNAKE_CASE__							,			'''trainer_state.json'''  )  ).log_history
										__lowerCamelCase    					=	int(logs[0]['''train_mem_gpu_peaked_delta'''] / 2**20  )
										__lowerCamelCase    					=	int(logs[0]['''train_mem_gpu_alloc_delta'''] / 2**20  )

										__lowerCamelCase    					=	logs[0]['''train_loss''']
										return gpu_peak_mem_mb, gpu_alloc_mem_mb, loss

						__lowerCamelCase					,    __lowerCamelCase					,    __lowerCamelCase    					=	train_and_return_metrics(OptimizerNames.ADAMW_TORCH.value  )
						__lowerCamelCase					,    __lowerCamelCase					,    __lowerCamelCase    					=	train_and_return_metrics(OptimizerNames.ADAMW_BNB.value  )

						__lowerCamelCase    					=	gpu_alloc_mem_orig - gpu_alloc_mem_bnb

						__lowerCamelCase    					=	gpu_peak_mem_orig + gpu_alloc_mem_orig
						__lowerCamelCase    					=	gpu_peak_mem_bnb + gpu_alloc_mem_bnb
						__lowerCamelCase    					=	gpu_total_mem_orig - gpu_total_mem_bnb

						# sshleifer/student_marian_en_ro_6_1 has 54M parameter, 29M of which is `nn.Embedding` which
						# doesn't get quantized and remains in fp32. Therefore we only have 25M parameters quantized
						# in 2 bytes and the diff in optim memory usage is derived as so:
						#
						# - normal 25*8=~200MB (8 bytes per param)
						# - bnb    25*2= ~50MB (2 bytes per param)
						#
						# Thus we should expect ~150MB total memory saved.
						#
						# Peak memory should be the same - the total should be different by about that same margin
						#
						# After leaving a small margin to accommodate for differences between gpus let's check
						# that we have at least 120MB in savings
						__lowerCamelCase    					=	1_20

						# uncomment the following if this test starts failing - requires py38 for a new print feature
						# gpu_peak_mem_diff = gpu_peak_mem_orig - gpu_peak_mem_bnb
						# print(f"{gpu_alloc_mem_orig=}MB {gpu_peak_mem_orig=}MB {gpu_alloc_mem_orig+gpu_peak_mem_orig=}MB")
						# print(f" {gpu_alloc_mem_bnb=}MB  {gpu_peak_mem_bnb=}MB  {gpu_alloc_mem_bnb+gpu_peak_mem_bnb=}MB")
						# print(f"{gpu_alloc_mem_diff=}MB")
						# print(f"{gpu_peak_mem_diff=}MB")
						# print(f"{gpu_total_mem_orig=}MB, {gpu_total_mem_bnb=}MB")
						# print(f"{gpu_total_mem_diff=}MB, {gpu_total_mem_diff=}MB")

						self.assertGreater(
						    SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			'''should use ~150MB less alloc gpu memory with BNB, compared to without it for this model but got'''
						    f''' a difference of {gpu_alloc_mem_diff}MB, with gpu_alloc_mem_orig={gpu_alloc_mem_orig}MB and'''
						    f''' gpu_alloc_mem_bnb={gpu_alloc_mem_bnb}MB'''							,			)

						self.assertGreater(
						    SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			'''should use ~150MB less total gpu memory with BNB, compared to without it for this model but got'''
						    f''' a difference of {gpu_total_mem_diff}MB, with gpu_total_mem_orig={gpu_total_mem_orig}MB and'''
						    f''' gpu_total_mem_bnb={gpu_total_mem_bnb}MB'''							,			)

						self.assertEqual(
						    SCREAMING_SNAKE_CASE__							,			SCREAMING_SNAKE_CASE__							,			f'''loss should be the same, but got loss_orig={loss_orig}, loss_bnb={loss_bnb}'''  )







		def  __A       (		self		:							Dict							,			SCREAMING_SNAKE_CASE__		:							int							,			SCREAMING_SNAKE_CASE__		:							str							,			SCREAMING_SNAKE_CASE__		:							int							,			SCREAMING_SNAKE_CASE__		:							float = 3e-3							,			SCREAMING_SNAKE_CASE__		:							str = "adafactor"							,			SCREAMING_SNAKE_CASE__		:							bool = False							,			SCREAMING_SNAKE_CASE__		:							str = None							,			SCREAMING_SNAKE_CASE__		:							int = 0							,			SCREAMING_SNAKE_CASE__		:							bool = True							,			SCREAMING_SNAKE_CASE__		:							bool = True							,			SCREAMING_SNAKE_CASE__		:							bool = True							,			SCREAMING_SNAKE_CASE__		:							bool = True							,			SCREAMING_SNAKE_CASE__		:							int = None							,			)     ->							List[Any]:
						__lowerCamelCase    					=	self.test_file_dir / '''../fixtures/tests_samples/wmt_en_ro'''
						__lowerCamelCase    					=	self.get_auto_remove_tmp_dir()
						__lowerCamelCase    					=	f'''
            --model_name_or_path {model_name}
            --train_file {data_dir}/train.json
            --validation_file {data_dir}/val.json
            --test_file {data_dir}/test.json
            --output_dir {output_dir}
            --overwrite_output_dir
            --max_train_samples 8
            --max_source_length {max_len}
            --max_target_length {max_len}
            --do_train
            --num_train_epochs {str(SCREAMING_SNAKE_CASE__  )}
            --per_device_train_batch_size 4
            --learning_rate {learning_rate}
            --warmup_steps 8
            --logging_steps 0
            --logging_strategy no
            --save_steps {str(SCREAMING_SNAKE_CASE__  )}
            --group_by_length
            --label_smoothing_factor 0.1
            --target_lang ro_RO
            --source_lang en_XX
        '''.split()

						__lowerCamelCase    					=	f'''
            --do_eval
            --per_device_eval_batch_size 4
            --max_eval_samples 8
            --val_max_target_length {max_len}
            --evaluation_strategy steps
            --eval_steps {str(SCREAMING_SNAKE_CASE__  )}
        '''.split()

						__lowerCamelCase    					=	'''
            --do_predict
        '''.split()

						__lowerCamelCase    					=	[]
						if do_train:
										args += args_train

						if do_eval:
										args += args_eval

						if do_predict:
										args += args_predict

						if predict_with_generate:
										args += "--predict_with_generate".split()

						if do_train:
										if optim == "adafactor":
														args += "--adafactor".split()
										else:
														args += f'''--optim {optim}'''.split()

						if extra_args_str is not None:
										args += extra_args_str.split()

						if distributed:
										if n_gpus_to_use is None:
														__lowerCamelCase    					=	get_gpu_count()
										__lowerCamelCase    					=	get_torch_dist_unique_port()
										__lowerCamelCase    					=	f'''
                -m torch.distributed.run
                --nproc_per_node={n_gpus_to_use}
                --master_port={master_port}
                {self.examples_dir_str}/pytorch/translation/run_translation.py
            '''.split()
										__lowerCamelCase    					=	[sys.executable] + distributed_args + args
										# keep for quick debug
										# print(" ".join([f"\nPYTHONPATH={self.src_dir_str}"] +cmd)); die
										execute_subprocess_async(SCREAMING_SNAKE_CASE__							,			env=self.get_env()  )
						else:
										__lowerCamelCase    					=	['''run_translation.py'''] + args
										with patch.object(SCREAMING_SNAKE_CASE__							,			'''argv'''							,			SCREAMING_SNAKE_CASE__  ):
														main()

						return output_dir



