


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





     def __init__(self		,     UpperCAmelCase		,     UpperCAmelCase=2		,     UpperCAmelCase=True		,     UpperCAmelCase=False		,     UpperCAmelCase=1_0		,     UpperCAmelCase=3		,     UpperCAmelCase=3_2 * 8		,     UpperCAmelCase=3_2 * 8		,     UpperCAmelCase=4		,     UpperCAmelCase=6_4		,     )     ->   List[Any]:
      _lowercase							      =parent
      _lowercase							      =batch_size
      _lowercase							      =is_training
      _lowercase							      =use_auxiliary_loss
      _lowercase							      =num_queries
      _lowercase							      =num_channels
      _lowercase							      =min_size
      _lowercase							      =max_size
      _lowercase							      =num_labels
      _lowercase							      =hidden_dim
      _lowercase							      =hidden_dim





     def      __A					(self       )     ->   int:
      _lowercase							      =floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]       ).to(
          UpperCAmelCase       )

      _lowercase							      =torch.ones([self.batch_size, self.min_size, self.max_size]		,     device=UpperCAmelCase       )

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

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





     def      __A					(self       )     ->   int:
      _lowercase							      =MaskaFormerConfig(
          hidden_size=self.hidden_dim		,     )
      _lowercase							      =self.num_queries
      _lowercase							      =self.num_labels

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

      _lowercase							      =6_4
      _lowercase							      =1_2_8
      _lowercase							      =self.hidden_dim
      _lowercase							      =self.hidden_dim
      _lowercase							      =self.hidden_dim
      return config





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





     def      __A					(self		,     UpperCAmelCase		,     UpperCAmelCase       )     ->   str:
      _lowercase							      =output.encoder_hidden_states
      _lowercase							      =output.pixel_decoder_hidden_states
      _lowercase							      =output.transformer_decoder_hidden_states

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





     def      __A					(self		,     UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase=False       )     ->   Any:
      with torch.no_grad():
       _lowercase							      =MaskaFormerModel(config=UpperCAmelCase       )
       model.to(UpperCAmelCase       )
       model.eval()

       _lowercase							      =model(pixel_values=UpperCAmelCase		,     pixel_mask=UpperCAmelCase       )
       _lowercase							      =model(UpperCAmelCase		,     output_hidden_states=UpperCAmelCase       )

      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(UpperCAmelCase		,     UpperCAmelCase       )





     def      __A					(self		,     UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase		,     UpperCAmelCase       )     ->   Any:
      _lowercase							      =MaskaFormerForUniversalSegmentation(config=UpperCAmelCase       )
      model.to(UpperCAmelCase       )
      model.eval()

      def comm_check_on_output(UpperCAmelCase       ):
       # 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():
       _lowercase							      =model(pixel_values=UpperCAmelCase		,     pixel_mask=UpperCAmelCase       )
       _lowercase							      =model(UpperCAmelCase       )

       comm_check_on_output(UpperCAmelCase       )

       _lowercase							      =model(
           pixel_values=UpperCAmelCase		,     pixel_mask=UpperCAmelCase		,     mask_labels=UpperCAmelCase		,     class_labels=UpperCAmelCase       )

      comm_check_on_output(UpperCAmelCase       )

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







@require_torch
class       lowerCamelCase__   (       lowerCAmelCase   ,							lowerCAmelCase   ,							unittest.TestCase):
     SCREAMING_SNAKE_CASE__      			=			(MaskaFormerModel, MaskaFormerForUniversalSegmentation) if is_torch_available() else ()
     SCREAMING_SNAKE_CASE__      			=			{'''feature-extraction''': MaskaFormerModel} if is_torch_available() else {}

     SCREAMING_SNAKE_CASE__      			=			False
     SCREAMING_SNAKE_CASE__      			=			False
     SCREAMING_SNAKE_CASE__      			=			False
     SCREAMING_SNAKE_CASE__      			=			False





     def      __A					(self       )     ->   Dict:
      _lowercase							      =MaskaFormerModelTester(self       )
      _lowercase							      =ConfigTester(self		,     config_class=UpperCAmelCase		,     has_text_modality=UpperCAmelCase       )





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





     def      __A					(self       )     ->   Union[str, Any]:
      _lowercase      ,    _lowercase							      =self.model_tester.prepare_config_and_inputs_for_common()
      self.model_tester.create_and_check_maskaformer_model(UpperCAmelCase		,     **UpperCAmelCase		,     output_hidden_states=UpperCAmelCase       )





     def      __A					(self       )     ->   str:
      _lowercase							      =self.model_tester.prepare_config_and_inputs()
      self.model_tester.create_and_check_maskaformer_instance_segmentation_head_model(*UpperCAmelCase       )





     @unittest.skip(reason='''Mask2Former does not use inputs_embeds'''       )
     def      __A					(self       )     ->   Dict:
      pass





     @unittest.skip(reason='''Mask2Former does not have a get_input_embeddings method'''       )
     def      __A					(self       )     ->   int:
      pass





     @unittest.skip(reason='''Mask2Former is not a generative model'''       )
     def      __A					(self       )     ->   Tuple:
      pass





     @unittest.skip(reason='''Mask2Former does not use token embeddings'''       )
     def      __A					(self       )     ->   Any:
      pass





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





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





     def      __A					(self       )     ->   List[str]:
      _lowercase      ,    _lowercase							      =self.model_tester.prepare_config_and_inputs_for_common()

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

       _lowercase							      =['''pixel_values''']
       self.assertListEqual(arg_names[:1]		,     UpperCAmelCase       )





     @slow
     def      __A					(self       )     ->   Tuple:
      for model_name in ["facebook/mask2former-swin-small-coco-instance"]:
       _lowercase							      =MaskaFormerModel.from_pretrained(UpperCAmelCase       )
       self.assertIsNotNone(UpperCAmelCase       )





     def      __A					(self       )     ->   Optional[Any]:
      _lowercase							      =(self.model_tester.min_size,) * 2
      _lowercase							      ={
          '''pixel_values''': torch.randn((2, 3, *size)		,     device=UpperCAmelCase       ),
          '''mask_labels''': torch.randn((2, 1_0, *size)		,     device=UpperCAmelCase       ),
          '''class_labels''': torch.zeros(2		,     1_0		,     device=UpperCAmelCase       ).long(),
      }
      _lowercase							      =self.model_tester.get_config()

      _lowercase							      =MaskaFormerForUniversalSegmentation(UpperCAmelCase       ).to(UpperCAmelCase       )
      _lowercase							      =model(**UpperCAmelCase       )
      self.assertTrue(outputs.loss is not None       )





     def      __A					(self       )     ->   Optional[Any]:
      _lowercase      ,    _lowercase							      =self.model_tester.prepare_config_and_inputs_for_common()
      self.model_tester.create_and_check_maskaformer_model(UpperCAmelCase		,     **UpperCAmelCase		,     output_hidden_states=UpperCAmelCase       )





     def      __A					(self       )     ->   Dict:
      _lowercase      ,    _lowercase							      =self.model_tester.prepare_config_and_inputs_for_common()

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





     def      __A					(self       )     ->   Any:
      if not self.model_tester.is_training:
       return

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

      _lowercase							      =model_class(UpperCAmelCase       )
      model.to(UpperCAmelCase       )
      model.train()

      _lowercase							      =model(UpperCAmelCase		,     mask_labels=UpperCAmelCase		,     class_labels=UpperCAmelCase       ).loss
      loss.backward()





     def      __A					(self       )     ->   Dict:
      _lowercase							      =self.all_model_classes[1]
      _lowercase      ,    _lowercase      ,    _lowercase      ,    _lowercase      ,    _lowercase							      =self.model_tester.prepare_config_and_inputs()
      _lowercase							      =True
      _lowercase							      =True

      _lowercase							      =model_class(UpperCAmelCase       ).to(UpperCAmelCase       )
      model.train()

      _lowercase							      =model(UpperCAmelCase		,     mask_labels=UpperCAmelCase		,     class_labels=UpperCAmelCase       )

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

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

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

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

      outputs.loss.backward(retain_graph=UpperCAmelCase       )

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



UpperCAmelCase__   	=							1e-4
def 				UpperCAmelCase_					(							)     ->			Optional[Any]:







 """simple docstring"""


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







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





     @cached_property
     def      __A					(self       )     ->   List[str]:
      return "facebook/mask2former-swin-small-coco-instance"





     @cached_property
     def      __A					(self       )     ->   Tuple:
      return MaskaFormerImageProcessor.from_pretrained(self.model_checkpoints       ) if is_vision_available() else None





     def      __A					(self       )     ->   str:
      _lowercase							      =MaskaFormerModel.from_pretrained(self.model_checkpoints       ).to(UpperCAmelCase       )
      _lowercase							      =self.default_image_processor
      _lowercase							      =prepare_img()
      _lowercase							      =image_processor(UpperCAmelCase		,     return_tensors='''pt'''       ).to(UpperCAmelCase       )
      _lowercase							      =inputs['''pixel_values'''].shape
      # check size is divisible by 32
      self.assertTrue((inputs_shape[-1] % 3_2) == 0 and (inputs_shape[-2] % 3_2) == 0       )
      # check size
      self.assertEqual(UpperCAmelCase		,     (1, 3, 3_8_4, 3_8_4)       )

      with torch.no_grad():
       _lowercase							      =model(**UpperCAmelCase       )

      _lowercase							      =torch.tensor(
          [[-0.2790, -1.0717, -1.1668], [-0.5128, -0.3128, -0.4987], [-0.5832, 0.1971, -0.0197]]       ).to(UpperCAmelCase       )
      self.assertTrue(
          torch.allclose(
              outputs.encoder_last_hidden_state[0, 0, :3, :3]		,     UpperCAmelCase		,     atol=UpperCAmelCase       )       )

      _lowercase							      =torch.tensor(
          [[0.8973, 1.1847, 1.1776], [1.1934, 1.5040, 1.5128], [1.1153, 1.4486, 1.4951]]       ).to(UpperCAmelCase       )
      self.assertTrue(
          torch.allclose(
              outputs.pixel_decoder_last_hidden_state[0, 0, :3, :3]		,     UpperCAmelCase		,     atol=UpperCAmelCase       )       )

      _lowercase							      =torch.tensor(
          [[2.1152, 1.7000, -0.8603], [1.5808, 1.8004, -0.9353], [1.6043, 1.7495, -0.5999]]       ).to(UpperCAmelCase       )
      self.assertTrue(
          torch.allclose(
              outputs.transformer_decoder_last_hidden_state[0, :3, :3]		,     UpperCAmelCase		,     atol=UpperCAmelCase       )       )





     def      __A					(self       )     ->   Tuple:
      _lowercase							      =MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints       ).to(UpperCAmelCase       ).eval()
      _lowercase							      =self.default_image_processor
      _lowercase							      =prepare_img()
      _lowercase							      =image_processor(UpperCAmelCase		,     return_tensors='''pt'''       ).to(UpperCAmelCase       )
      _lowercase							      =inputs['''pixel_values'''].shape
      # check size is divisible by 32
      self.assertTrue((inputs_shape[-1] % 3_2) == 0 and (inputs_shape[-2] % 3_2) == 0       )
      # check size
      self.assertEqual(UpperCAmelCase		,     (1, 3, 3_8_4, 3_8_4)       )

      with torch.no_grad():
       _lowercase							      =model(**UpperCAmelCase       )
      # masks_queries_logits
      _lowercase							      =outputs.masks_queries_logits
      self.assertEqual(
          masks_queries_logits.shape		,     (1, model.config.num_queries, inputs_shape[-2] // 4, inputs_shape[-1] // 4)       )
      _lowercase							      =[
          [-8.7839, -9.0056, -8.8121],
          [-7.4104, -7.0313, -6.5401],
          [-6.6105, -6.3427, -6.4675],
      ]
      _lowercase							      =torch.tensor(UpperCAmelCase       ).to(UpperCAmelCase       )
      self.assertTrue(torch.allclose(masks_queries_logits[0, 0, :3, :3]		,     UpperCAmelCase		,     atol=UpperCAmelCase       )       )
      # class_queries_logits
      _lowercase							      =outputs.class_queries_logits
      self.assertEqual(class_queries_logits.shape		,     (1, model.config.num_queries, model.config.num_labels + 1)       )
      _lowercase							      =torch.tensor(
          [
              [1.8324, -8.0835, -4.1922],
              [0.8450, -9.0050, -3.6053],
              [0.3045, -7.7293, -3.0275],
          ]       ).to(UpperCAmelCase       )
      self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]		,     UpperCAmelCase		,     atol=UpperCAmelCase       )       )





     def      __A					(self       )     ->   Optional[Any]:
      _lowercase							      =MaskaFormerForUniversalSegmentation.from_pretrained(self.model_checkpoints       ).to(UpperCAmelCase       ).eval()
      _lowercase							      =self.default_image_processor

      _lowercase							      =image_processor(
          [np.zeros((3, 8_0_0, 1_3_3_3)       ), np.zeros((3, 8_0_0, 1_3_3_3)       )]		,     segmentation_maps=[np.zeros((3_8_4, 3_8_4)       ).astype(np.floataa       ), np.zeros((3_8_4, 3_8_4)       ).astype(np.floataa       )]		,     return_tensors='''pt'''		,     )

      _lowercase							      =inputs['''pixel_values'''].to(UpperCAmelCase       )
      _lowercase							      =[el.to(UpperCAmelCase       ) for el in inputs['''mask_labels''']]
      _lowercase							      =[el.to(UpperCAmelCase       ) for el in inputs['''class_labels''']]

      with torch.no_grad():
       _lowercase							      =model(**UpperCAmelCase       )

      self.assertTrue(outputs.loss is not None       )





from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_tf_available, is_torch_available


UpperCAmelCase__   	=							{
    '''configuration_xlm''': ['''XLM_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''XLMConfig''', '''XLMOnnxConfig'''],
    '''tokenization_xlm''': ['''XLMTokenizer'''],
}

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase__   	=							[
      '''XLM_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''XLMForMultipleChoice''',
      '''XLMForQuestionAnswering''',
      '''XLMForQuestionAnsweringSimple''',
      '''XLMForSequenceClassification''',
      '''XLMForTokenClassification''',
      '''XLMModel''',
      '''XLMPreTrainedModel''',
      '''XLMWithLMHeadModel''',
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  UpperCAmelCase__   	=							[
      '''TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST''',
      '''TFXLMForMultipleChoice''',
      '''TFXLMForQuestionAnsweringSimple''',
      '''TFXLMForSequenceClassification''',
      '''TFXLMForTokenClassification''',
      '''TFXLMMainLayer''',
      '''TFXLMModel''',
      '''TFXLMPreTrainedModel''',
      '''TFXLMWithLMHeadModel''',
  ]


if TYPE_CHECKING:
  from .configuration_xlm import XLM_PRETRAINED_CONFIG_ARCHIVE_MAP, XLMConfig, XLMOnnxConfig
  from .tokenization_xlm import XLMTokenizer

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_xlm import (
        XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
        XLMForMultipleChoice,
        XLMForQuestionAnswering,
        XLMForQuestionAnsweringSimple,
        XLMForSequenceClassification,
        XLMForTokenClassification,
        XLMModel,
        XLMPreTrainedModel,
        XLMWithLMHeadModel,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_xlm import (
        TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFXLMForMultipleChoice,
        TFXLMForQuestionAnsweringSimple,
        TFXLMForSequenceClassification,
        TFXLMForTokenClassification,
        TFXLMMainLayer,
        TFXLMModel,
        TFXLMPreTrainedModel,
        TFXLMWithLMHeadModel,
    )

else:
  import sys

  UpperCAmelCase__   	=							_LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)

