


'''simple docstring'''


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  :     int	,  lowerCamelCase__  :     int	,  lowerCamelCase__  :     Dict=2	,  lowerCamelCase__  :     int=True	,  lowerCamelCase__  :     Tuple=False	,  lowerCamelCase__  :     Any=10	,  lowerCamelCase__  :     Tuple=3	,  lowerCamelCase__  :     Optional[Any]=32 * 4	,  lowerCamelCase__  :     int=32 * 6	,  lowerCamelCase__  :     Union[str, Any]=4	,  lowerCamelCase__  :     List[str]=32	,  )   ->Any:



			'''simple docstring'''







			_UpperCAmelCase		:       Optional[int]								=      parent
			_UpperCAmelCase		:       Any								=      batch_size
			_UpperCAmelCase		:       Dict								=      is_training
			_UpperCAmelCase		:       List[Any]								=      use_auxiliary_loss
			_UpperCAmelCase		:       List[Any]								=      num_queries
			_UpperCAmelCase		:       List[str]								=      num_channels
			_UpperCAmelCase		:       List[str]								=      min_size
			_UpperCAmelCase		:       str								=      max_size
			_UpperCAmelCase		:       Tuple								=      num_labels
			_UpperCAmelCase		:       Optional[int]								=      mask_feature_size




	def 							lowerCAmelCase__    (		self  :     Dict )   ->Any:



			'''simple docstring'''







			_UpperCAmelCase		:       Union[str, Any]								=      floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size] ).to(
			    snake_case_ )

			_UpperCAmelCase		:       Optional[int]								=      torch.ones([self.batch_size, self.min_size, self.max_size]	,  device=snake_case_ )

			_UpperCAmelCase		:       Optional[Any]								=      (
			    torch.rand([self.batch_size, self.num_labels, self.min_size, self.max_size]	,  device=snake_case_ ) > 0.5
			).float()
			_UpperCAmelCase		:       Optional[int]								=      (torch.rand((self.batch_size, self.num_labels)	,  device=snake_case_ ) > 0.5).long()

			_UpperCAmelCase		:       List[str]								=      self.get_config()
			return config, pixel_values, pixel_mask, mask_labels, class_labels




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



			'''simple docstring'''







			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 							lowerCAmelCase__    (		self  :     Optional[Any] )   ->Dict:



			'''simple docstring'''







			_UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase		:       Optional[int]								=      self.prepare_config_and_inputs()
			_UpperCAmelCase		:       Tuple								=      {"pixel_values": pixel_values, "pixel_mask": pixel_mask}
			return config, inputs_dict




	def 							lowerCAmelCase__    (		self  :     Tuple	,  lowerCamelCase__  :     Dict	,  lowerCamelCase__  :     int )   ->int:



			'''simple docstring'''







			_UpperCAmelCase		:       str								=      output.encoder_hidden_states
			_UpperCAmelCase		:       Any								=      output.pixel_decoder_hidden_states
			_UpperCAmelCase		:       Union[str, Any]								=      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_config.decoder_layers )




	def 							lowerCAmelCase__    (		self  :     Tuple	,  lowerCamelCase__  :     List[Any]	,  lowerCamelCase__  :     List[str]	,  lowerCamelCase__  :     Tuple	,  lowerCamelCase__  :     Dict=False )   ->Optional[Any]:



			'''simple docstring'''







			with torch.no_grad():
					_UpperCAmelCase		:       int								=      MaskFormerModel(config=snake_case_ )
					model.to(snake_case_ )
					model.eval()

					_UpperCAmelCase		:       str								=      model(pixel_values=snake_case_	,  pixel_mask=snake_case_ )
					_UpperCAmelCase		:       Dict								=      model(snake_case_	,  output_hidden_states=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(snake_case_	,  snake_case_ )






	def 							lowerCAmelCase__    (		self  :     int	,  lowerCamelCase__  :     List[Any]	,  lowerCamelCase__  :     Tuple	,  lowerCamelCase__  :     Optional[Any]	,  lowerCamelCase__  :     Tuple	,  lowerCamelCase__  :     str )   ->str:



			'''simple docstring'''







			_UpperCAmelCase		:       Any								=      MaskFormerForInstanceSegmentation(config=snake_case_ )
			model.to(snake_case_ )
			model.eval()

			def comm_check_on_output(lowerCamelCase__  :     Optional[int] ):
					# 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		:       Optional[Any]								=      model(pixel_values=snake_case_	,  pixel_mask=snake_case_ )
					_UpperCAmelCase		:       int								=      model(snake_case_ )

					comm_check_on_output(snake_case_ )

					_UpperCAmelCase		:       str								=      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 	lowerCAmelCase__			(			_a    ,       _a    ,       unittest.TestCase	):
	lowerCAmelCase :	Union[str, Any]          =     (MaskFormerModel, MaskFormerForInstanceSegmentation) if is_torch_available() else ()
	lowerCAmelCase :	List[Any]          =     (
	    {"""feature-extraction""": MaskFormerModel, """image-segmentation""": MaskFormerForInstanceSegmentation}
	    if is_torch_available()
	    else {}
	)

	lowerCAmelCase :	Union[str, Any]          =     False
	lowerCAmelCase :	Optional[int]          =     False
	lowerCAmelCase :	int          =     False
	lowerCAmelCase :	List[Any]          =     False




	def 							lowerCAmelCase__    (		self  :     Optional[Any] )   ->Union[str, Any]:



			'''simple docstring'''







			_UpperCAmelCase		:       Union[str, Any]								=      MaskFormerModelTester(self )
			_UpperCAmelCase		:       Union[str, Any]								=      ConfigTester(self	,  config_class=snake_case_	,  has_text_modality=snake_case_ )




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



			'''simple docstring'''







			self.config_tester.run_common_tests()




	def 							lowerCAmelCase__    (		self  :     List[Any] )   ->Optional[int]:



			'''simple docstring'''







			_UpperCAmelCase					,    _UpperCAmelCase		:       List[str]								=      self.model_tester.prepare_config_and_inputs_for_common()
			self.model_tester.create_and_check_maskformer_model(snake_case_	,  **snake_case_	,  output_hidden_states=snake_case_ )




	def 							lowerCAmelCase__    (		self  :     str )   ->str:



			'''simple docstring'''







			_UpperCAmelCase		:       List[Any]								=      self.model_tester.prepare_config_and_inputs()
			self.model_tester.create_and_check_maskformer_instance_segmentation_head_model(*snake_case_ )




	@unittest.skip(reason="MaskFormer does not use inputs_embeds" )
	def 							lowerCAmelCase__    (		self  :     Union[str, Any] )   ->Union[str, Any]:



			'''simple docstring'''







			pass




	@unittest.skip(reason="MaskFormer does not have a get_input_embeddings method" )
	def 							lowerCAmelCase__    (		self  :     Dict )   ->str:



			'''simple docstring'''







			pass




	@unittest.skip(reason="MaskFormer is not a generative model" )
	def 							lowerCAmelCase__    (		self  :     Optional[Any] )   ->int:



			'''simple docstring'''







			pass




	@unittest.skip(reason="MaskFormer does not use token embeddings" )
	def 							lowerCAmelCase__    (		self  :     Optional[Any] )   ->int:



			'''simple docstring'''







			pass




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



			'''simple docstring'''







			pass




	@unittest.skip("Will be fixed soon by reducing the size of the model used for common tests." )
	def 							lowerCAmelCase__    (		self  :     Union[str, Any] )   ->Any:



			'''simple docstring'''







			pass




	def 							lowerCAmelCase__    (		self  :     Optional[int] )   ->List[Any]:



			'''simple docstring'''







			_UpperCAmelCase					,    _UpperCAmelCase		:       List[Any]								=      self.model_tester.prepare_config_and_inputs_for_common()

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

					_UpperCAmelCase		:       Optional[int]								=      ["pixel_values"]
					self.assertListEqual(arg_names[:1]	,  snake_case_ )




	@slow
	def 							lowerCAmelCase__    (		self  :     Optional[int] )   ->Dict:



			'''simple docstring'''







			for model_name in ["facebook/maskformer-swin-small-coco"]:
					_UpperCAmelCase		:       List[str]								=      MaskFormerModel.from_pretrained(snake_case_ )
					self.assertIsNotNone(snake_case_ )




	def 							lowerCAmelCase__    (		self  :     List[str] )   ->int:



			'''simple docstring'''







			_UpperCAmelCase		:       Tuple								=      (self.model_tester.min_size,) * 2
			_UpperCAmelCase		:       List[str]								=      {
			    "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(),
			}

			_UpperCAmelCase		:       Dict								=      MaskFormerForInstanceSegmentation(MaskFormerConfig() ).to(snake_case_ )
			_UpperCAmelCase		:       Union[str, Any]								=      model(**snake_case_ )
			self.assertTrue(outputs.loss is not None )




	def 							lowerCAmelCase__    (		self  :     int )   ->Dict:



			'''simple docstring'''







			_UpperCAmelCase					,    _UpperCAmelCase		:       int								=      self.model_tester.prepare_config_and_inputs_for_common()
			self.model_tester.create_and_check_maskformer_model(snake_case_	,  **snake_case_	,  output_hidden_states=snake_case_ )




	def 							lowerCAmelCase__    (		self  :     Any )   ->List[str]:



			'''simple docstring'''







			_UpperCAmelCase					,    _UpperCAmelCase		:       List[str]								=      self.model_tester.prepare_config_and_inputs_for_common()

			for model_class in self.all_model_classes:
					_UpperCAmelCase		:       str								=      model_class(snake_case_ ).to(snake_case_ )
					_UpperCAmelCase		:       Any								=      model(**snake_case_	,  output_attentions=snake_case_ )
					self.assertTrue(outputs.attentions is not None )




	def 							lowerCAmelCase__    (		self  :     Optional[Any] )   ->str:



			'''simple docstring'''







			if not self.model_tester.is_training:
					return
			# only MaskFormerForInstanceSegmentation has the loss
			_UpperCAmelCase		:       Optional[int]								=      self.all_model_classes[1]
			_UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase		:       str								=      self.model_tester.prepare_config_and_inputs()

			_UpperCAmelCase		:       Dict								=      model_class(snake_case_ )
			model.to(snake_case_ )
			model.train()

			_UpperCAmelCase		:       Tuple								=      model(snake_case_	,  mask_labels=snake_case_	,  class_labels=snake_case_ ).loss
			loss.backward()






	def 							lowerCAmelCase__    (		self  :     Dict )   ->int:



			'''simple docstring'''







			_UpperCAmelCase		:       Union[str, Any]								=      self.all_model_classes[1]
			_UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase					,    _UpperCAmelCase		:       Optional[int]								=      self.model_tester.prepare_config_and_inputs()
			_UpperCAmelCase		:       Optional[int]								=      True
			_UpperCAmelCase		:       Union[str, Any]								=      True

			_UpperCAmelCase		:       List[str]								=      model_class(snake_case_ )
			model.to(snake_case_ )
			model.train()

			_UpperCAmelCase		:       Optional[int]								=      model(snake_case_	,  mask_labels=snake_case_	,  class_labels=snake_case_ )

			_UpperCAmelCase		:       Tuple								=      outputs.encoder_hidden_states[0]
			encoder_hidden_states.retain_grad()

			_UpperCAmelCase		:       List[str]								=      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
			_UpperCAmelCase		:       str								=      outputs.transformer_decoder_hidden_states[0]
			transformer_decoder_hidden_states.retain_grad()

			_UpperCAmelCase		:       Union[str, Any]								=      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 )




lowerCamelCase__    	=     1e-4






def  __lowerCAmelCase ():
		_UpperCAmelCase		:       Dict								=      Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"   )
		return image






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




	@cached_property
	def 							lowerCAmelCase__    (		self  :     Tuple )   ->Optional[Any]:



			'''simple docstring'''







			return (
			    MaskFormerImageProcessor.from_pretrained("facebook/maskformer-swin-small-coco" )
			    if is_vision_available()
			    else None
			)




	def 							lowerCAmelCase__    (		self  :     Tuple )   ->Optional[Any]:



			'''simple docstring'''







			_UpperCAmelCase		:       List[Any]								=      MaskFormerModel.from_pretrained("facebook/maskformer-swin-small-coco" ).to(snake_case_ )
			_UpperCAmelCase		:       Union[str, Any]								=      self.default_image_processor
			_UpperCAmelCase		:       int								=      prepare_img()
			_UpperCAmelCase		:       Optional[Any]								=      image_processor(snake_case_	,  return_tensors="pt" ).to(snake_case_ )
			_UpperCAmelCase		:       List[str]								=      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, 8_00, 10_88) )

			with torch.no_grad():
					_UpperCAmelCase		:       List[Any]								=      model(**snake_case_ )

			_UpperCAmelCase		:       int								=      torch.tensor(
			    [[-0.0_4_8_2, 0.9_2_2_8, 0.4_9_5_1], [-0.2_5_4_7, 0.8_0_1_7, 0.8_5_2_7], [-0.0_0_6_9, 0.3_3_8_5, -0.0_0_8_9]] ).to(snake_case_ )
			self.assertTrue(
			    torch.allclose(
			        outputs.encoder_last_hidden_state[0, 0, :3, :3]	,  snake_case_	,  atol=snake_case_ ) )

			_UpperCAmelCase		:       int								=      torch.tensor(
			    [[-0.8_4_2_2, -0.8_4_3_4, -0.9_7_1_8], [-1.0_1_4_4, -0.5_5_6_5, -0.4_1_9_5], [-1.0_0_3_8, -0.4_4_8_4, -0.1_9_6_1]] ).to(snake_case_ )
			self.assertTrue(
			    torch.allclose(
			        outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]	,  snake_case_	,  atol=snake_case_ ) )

			_UpperCAmelCase		:       Dict								=      torch.tensor(
			    [[0.2_8_5_2, -0.0_1_5_9, 0.9_7_3_5], [0.6_2_5_4, 0.1_8_5_8, 0.8_5_2_9], [-0.0_6_8_0, -0.4_1_1_6, 1.8_4_1_3]] ).to(snake_case_ )
			self.assertTrue(
			    torch.allclose(
			        outputs.transformer_decoder_last_hidden_state[0, :3, :3]	,  snake_case_	,  atol=snake_case_ ) )




	def 							lowerCAmelCase__    (		self  :     Optional[Any] )   ->Optional[int]:



			'''simple docstring'''







			_UpperCAmelCase		:       Dict								=      (
			    MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-small-coco" )
			    .to(snake_case_ )
			    .eval()
			)
			_UpperCAmelCase		:       List[str]								=      self.default_image_processor
			_UpperCAmelCase		:       Optional[Any]								=      prepare_img()
			_UpperCAmelCase		:       Optional[int]								=      image_processor(snake_case_	,  return_tensors="pt" ).to(snake_case_ )
			_UpperCAmelCase		:       int								=      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, 8_00, 10_88) )

			with torch.no_grad():
					_UpperCAmelCase		:       Any								=      model(**snake_case_ )
			# masks_queries_logits
			_UpperCAmelCase		:       Union[str, Any]								=      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)	,  )
			_UpperCAmelCase		:       List[Any]								=      [
			    [-1.3_7_3_7_1_2_4, -1.7_7_2_4_9_3_7, -1.9_3_6_4_2_3_3],
			    [-1.5_9_7_7_2_8_1, -1.9_8_6_7_9_3_9, -2.1_5_2_3_6_9_5],
			    [-1.5_7_9_5_3_9_8, -1.9_2_6_9_8_3_2, -2.0_9_3_9_4_2],
			]
			_UpperCAmelCase		:       Optional[int]								=      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
			_UpperCAmelCase		:       Any								=      outputs.class_queries_logits
			self.assertEqual(
			    class_queries_logits.shape	,  (1, model.config.decoder_config.num_queries, model.config.num_labels + 1) )
			_UpperCAmelCase		:       Tuple								=      torch.tensor(
			    [
			        [1.6512E00, -5.2572E00, -3.3519E00],
			        [3.6169E-02, -5.9025E00, -2.9313E00],
			        [1.0766E-04, -7.7630E00, -5.1263E00],
			    ] ).to(snake_case_ )
			self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]	,  snake_case_	,  atol=snake_case_ ) )




	def 							lowerCAmelCase__    (		self  :     int )   ->List[str]:



			'''simple docstring'''







			_UpperCAmelCase		:       List[str]								=      (
			    MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-resnet101-coco-stuff" )
			    .to(snake_case_ )
			    .eval()
			)
			_UpperCAmelCase		:       Any								=      self.default_image_processor
			_UpperCAmelCase		:       List[Any]								=      prepare_img()
			_UpperCAmelCase		:       List[str]								=      image_processor(snake_case_	,  return_tensors="pt" ).to(snake_case_ )
			_UpperCAmelCase		:       Optional[int]								=      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, 8_00, 10_88) )

			with torch.no_grad():
					_UpperCAmelCase		:       List[Any]								=      model(**snake_case_ )
			# masks_queries_logits
			_UpperCAmelCase		:       int								=      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)	,  )
			_UpperCAmelCase		:       List[str]								=      [[-0.9_0_4_6, -2.6_3_6_6, -4.6_0_6_2], [-3.4_1_7_9, -5.7_8_9_0, -8.8_0_5_7], [-4.9_1_7_9, -7.6_5_6_0, -1_0.7_7_1_1]]
			_UpperCAmelCase		:       Union[str, Any]								=      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
			_UpperCAmelCase		:       Any								=      outputs.class_queries_logits
			self.assertEqual(
			    class_queries_logits.shape	,  (1, model.config.decoder_config.num_queries, model.config.num_labels + 1) )
			_UpperCAmelCase		:       Optional[Any]								=      torch.tensor(
			    [[4.7_1_8_8, -3.2_5_8_5, -2.8_8_5_7], [6.6_8_7_1, -2.9_1_8_1, -1.2_4_8_7], [7.2_4_4_9, -2.2_7_6_4, -2.1_8_7_4]] ).to(snake_case_ )
			self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]	,  snake_case_	,  atol=snake_case_ ) )






	def 							lowerCAmelCase__    (		self  :     Union[str, Any] )   ->Any:



			'''simple docstring'''







			_UpperCAmelCase		:       List[Any]								=      (
			    MaskFormerForInstanceSegmentation.from_pretrained("facebook/maskformer-swin-small-coco" )
			    .to(snake_case_ )
			    .eval()
			)
			_UpperCAmelCase		:       Union[str, Any]								=      self.default_image_processor

			_UpperCAmelCase		:       Dict								=      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"	,  )

			_UpperCAmelCase		:       Union[str, Any]								=      inputs["pixel_values"].to(snake_case_ )
			_UpperCAmelCase		:       str								=      [el.to(snake_case_ ) for el in inputs["mask_labels"]]
			_UpperCAmelCase		:       Any								=      [el.to(snake_case_ ) for el in inputs["class_labels"]]

			with torch.no_grad():
					_UpperCAmelCase		:       str								=      model(**snake_case_ )

			self.assertTrue(outputs.loss is not None )







'''simple docstring'''


import json
import multiprocessing
import os
import re
from collections import defaultdict

import torch
from accelerate import Accelerator
from accelerate.utils import set_seed
from arguments import HumanEvalArguments
from datasets import load_dataset, load_metric
from torch.utils.data import IterableDataset
from torch.utils.data.dataloader import DataLoader
from tqdm import tqdm

import transformers
from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser, StoppingCriteria, StoppingCriteriaList


lowerCamelCase__    	=     ['\nclass', '\ndef', '\n#', '\n@', '\nprint', '\nif']






class 	lowerCAmelCase__			(			UpperCAmelCase__	):




			def __init__(		self  :     Tuple	,  lowerCamelCase__  :     List[Any]	,  lowerCamelCase__  :     Optional[int]	,  lowerCamelCase__  :     Optional[Any]=None	,  lowerCamelCase__  :     int=1 )   ->Optional[int]:



					'''simple docstring'''







					_UpperCAmelCase		:       List[Any]								=      tokenizer
					_UpperCAmelCase		:       Tuple								=      dataset
					_UpperCAmelCase		:       Union[str, Any]								=      len(lowerCamelCase__ ) if n_tasks is None else n_tasks
					_UpperCAmelCase		:       Any								=      n_copies






			def __iter__(		self  :     Any )   ->Dict:



					'''simple docstring'''







					_UpperCAmelCase		:       Optional[int]								=      []
					for task in range(self.n_tasks ):
							# without strip, the model generate commented codes ...
							prompts.append(self.tokenizer.eos_token + self.dataset[task]["prompt"].strip() )
					_UpperCAmelCase		:       Optional[Any]								=      self.tokenizer(lowerCamelCase__	,  padding=lowerCamelCase__	,  return_tensors="pt" )
					for task in range(self.n_tasks ):
							for _ in range(self.n_copies ):
									yield {
									    "ids": outputs.input_ids[task],
									    "task_id": task,
									    "input_len": outputs.attention_mask[task].sum(),
									}






class 	lowerCAmelCase__			(			UpperCAmelCase__	):




			def __init__(		self  :     Optional[int]	,  lowerCamelCase__  :     Tuple	,  lowerCamelCase__  :     Any	,  lowerCamelCase__  :     Dict )   ->Any:



					'''simple docstring'''







					_UpperCAmelCase		:       List[str]								=      start_length
					_UpperCAmelCase		:       Union[str, Any]								=      eof_strings
					_UpperCAmelCase		:       Union[str, Any]								=      tokenizer






			def __call__(		self  :     Tuple	,  lowerCamelCase__  :     Union[str, Any]	,  lowerCamelCase__  :     List[str]	,  **lowerCamelCase__  :     Optional[int] )   ->Union[str, Any]:



					'''simple docstring'''







					_UpperCAmelCase		:       Tuple								=      self.tokenizer.batch_decode(input_ids[:, self.start_length :] )
					_UpperCAmelCase		:       Optional[int]								=      []
					for decoded_generation in decoded_generations:
							done.append(any(stop_string in decoded_generation for stop_string in self.eof_strings ) )
					return all(lowerCamelCase__ )






def  __lowerCAmelCase (__lowerCAmelCase   ):
		_UpperCAmelCase		:       Tuple								=      re.split("(%s)" % "|".join(__lowerCAmelCase   )      , __lowerCAmelCase   )
		# last string should be ""
		return "".join(string_list[:-2]   )






def  __lowerCAmelCase (__lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase=20      , **__lowerCAmelCase   ):
		_UpperCAmelCase		:       Tuple								=      defaultdict(__lowerCAmelCase   )  # dict of list of generated tokens
		for step, batch in tqdm(enumerate(__lowerCAmelCase   )   ):
				with torch.no_grad():
						_UpperCAmelCase		:       Tuple								=      batch["ids"].shape[-1]
						_UpperCAmelCase		:       Optional[int]								=      accelerator.unwrap_model(__lowerCAmelCase   ).generate(
						    input_ids=batch["ids"][:, : batch["input_len"]]      , num_return_sequences=__lowerCAmelCase      , **__lowerCAmelCase   )
						# each task is generated batch_size times
						_UpperCAmelCase		:       str								=      batch["task_id"].repeat(__lowerCAmelCase   )
						_UpperCAmelCase		:       str								=      accelerator.pad_across_processes(
						    __lowerCAmelCase      , dim=1      , pad_index=tokenizer.pad_token_id   )

						_UpperCAmelCase					,    _UpperCAmelCase		:       int								=      accelerator.gather((generated_tokens, generated_tasks)   )
						_UpperCAmelCase		:       Dict								=      generated_tokens.cpu().numpy()
						_UpperCAmelCase		:       Dict								=      generated_tasks.cpu().numpy()

						for task, generated_tokens in zip(__lowerCAmelCase      , __lowerCAmelCase   ):
								gen_token_dict[task].append(__lowerCAmelCase   )

		_UpperCAmelCase		:       int								=      [[] for _ in range(__lowerCAmelCase   )]
		for task, generated_tokens in gen_token_dict.items():
				for s in generated_tokens:
						_UpperCAmelCase		:       List[Any]								=      tokenizer.decode(__lowerCAmelCase      , skip_special_tokens=__lowerCAmelCase      , clean_up_tokenization_spaces=__lowerCAmelCase   )
						code_gens[task].append(remove_last_block(__lowerCAmelCase   )   )
		return code_gens






def  __lowerCAmelCase ():
		# Setup configuration
		_UpperCAmelCase		:       List[str]								=      HfArgumentParser(__lowerCAmelCase   )
		_UpperCAmelCase		:       Tuple								=      parser.parse_args()

		transformers.logging.set_verbosity_error()
		# enables code execution in code_eval metric
		_UpperCAmelCase		:       Any								=      args.HF_ALLOW_CODE_EVAL
		# make sure tokenizer plays nice with multiprocessing
		_UpperCAmelCase		:       List[str]								=      "false"

		if args.num_workers is None:
				_UpperCAmelCase		:       List[str]								=      multiprocessing.cpu_count()

		# Use dataset load to feed to accelerate
		_UpperCAmelCase		:       List[Any]								=      Accelerator()
		set_seed(args.seed      , device_specific=__lowerCAmelCase   )

		# Load model and tokenizer
		_UpperCAmelCase		:       Tuple								=      AutoTokenizer.from_pretrained(args.model_ckpt   )
		_UpperCAmelCase		:       List[str]								=      tokenizer.eos_token
		_UpperCAmelCase		:       str								=      AutoModelForCausalLM.from_pretrained(args.model_ckpt   )

		# Generation settings
		_UpperCAmelCase		:       Tuple								=      {
		    "do_sample": args.do_sample,
		    "temperature": args.temperature,
		    "max_new_tokens": args.max_new_tokens,
		    "top_p": args.top_p,
		    "top_k": args.top_k,
		    "stopping_criteria": StoppingCriteriaList([EndOfFunctionCriteria(0      , __lowerCAmelCase      , __lowerCAmelCase   )]   ),
		}

		# Load evaluation dataset and metric
		_UpperCAmelCase		:       Union[str, Any]								=      load_dataset("openai_humaneval"   )
		_UpperCAmelCase		:       List[Any]								=      load_metric("code_eval"   )

		_UpperCAmelCase		:       Optional[Any]								=      args.num_tasks if args.num_tasks is not None else len(human_eval["test"]   )
		_UpperCAmelCase		:       Any								=      args.n_samples // args.batch_size

		_UpperCAmelCase		:       Tuple								=      TokenizedDataset(__lowerCAmelCase      , human_eval["test"]      , n_copies=__lowerCAmelCase      , n_tasks=__lowerCAmelCase   )
		# do not confuse args.batch_size, which is actually the num_return_sequences
		_UpperCAmelCase		:       List[str]								=      DataLoader(__lowerCAmelCase      , batch_size=1   )

		# Run a quick test to see if code evaluation is enabled
		try:
				_UpperCAmelCase		:       Optional[int]								=      code_eval_metric.compute(references=[""]      , predictions=[[""]]   )
		except ValueError as exception:
				print(
				    "Code evaluation not enabled. Read the warning below carefully and then use `--HF_ALLOW_CODE_EVAL=\"1\"`"
				    " flag to enable code evaluation."   )
				raise exception

		_UpperCAmelCase					,    _UpperCAmelCase		:       Any								=      accelerator.prepare(__lowerCAmelCase      , __lowerCAmelCase   )

		_UpperCAmelCase		:       Dict								=      complete_code(
		    __lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase      , __lowerCAmelCase      , n_tasks=__lowerCAmelCase      , batch_size=args.batch_size      , **__lowerCAmelCase      , )

		if accelerator.is_main_process:
				_UpperCAmelCase		:       List[Any]								=      []

				for task in tqdm(range(__lowerCAmelCase   )   ):
						_UpperCAmelCase		:       str								=      human_eval["test"][task]["test"]
						_UpperCAmelCase		:       Union[str, Any]								=      F"""check({human_eval['test'][task]['entry_point']})"""
						references.append("\n" + test_func + "\n" + entry_point   )

				# Evaluate completions with "code_eval" metric
				_UpperCAmelCase					,    _UpperCAmelCase		:       str								=      code_eval_metric.compute(
				    references=__lowerCAmelCase      , predictions=__lowerCAmelCase      , num_workers=args.num_workers   )
				print(F"""Results: {pass_at_k}"""   )

				# Save results to json file
				with open(args.output_file      , "w"   ) as fp:
						json.dump(__lowerCAmelCase      , __lowerCAmelCase   )


# For some reason the folliwng seems to be necessary sometimes for code_eval to work nice with multiprocessing
# https://stackoverflow.com/questions/60804599/python-multiprocessing-keeps-spawning-the-whole-script
if __name__ == "__main__":
						main()



