



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



       def __init__(		self							:     Optional[int]     ,		__SCREAMING_SNAKE_CASE							:     Tuple     ,		__SCREAMING_SNAKE_CASE							:     int=2     ,		__SCREAMING_SNAKE_CASE							:     Optional[int]=True     ,		__SCREAMING_SNAKE_CASE							:     Optional[int]=False     ,		__SCREAMING_SNAKE_CASE							:     Dict=10     ,		__SCREAMING_SNAKE_CASE							:     str=3     ,		__SCREAMING_SNAKE_CASE							:     Union[str, Any]=32 * 4     ,		__SCREAMING_SNAKE_CASE							:     int=32 * 6     ,		__SCREAMING_SNAKE_CASE							:     List[Any]=4     ,		__SCREAMING_SNAKE_CASE							:     Tuple=32     ,		):


           '''simple docstring'''

           __a						       =				parent
           __a						       =				batch_size
           __a						       =				is_training
           __a						       =				use_auxiliary_loss
           __a						       =				num_queries
           __a						       =				num_channels
           __a						       =				min_size
           __a						       =				max_size
           __a						       =				num_labels
           __a						       =				mask_feature_size



       def 		_lowerCamelCase			(		self							:     List[Any]):


           '''simple docstring'''

           __a						       =				floats_tensor([self.batch_size, self.num_channels, self.min_size, self.max_size]).to(
               __SCREAMING_SNAKE_CASE)

           __a						       =				torch.ones([self.batch_size, self.min_size, self.max_size]     ,		device=__SCREAMING_SNAKE_CASE)

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

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



       def 		_lowerCamelCase			(		self							:     Any):


           '''simple docstring'''

           return MaskFormerConfig.from_backbone_and_decoder_configs(
               backbone_config=SwinConfig(
                   depths=[1, 1, 1, 1]     ,		)     ,		decoder_config=DetrConfig(
                   decoder_ffn_dim=128     ,		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							:     Union[str, Any]):


           '''simple docstring'''

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



       def 		_lowerCamelCase			(		self							:     Optional[int]     ,		__SCREAMING_SNAKE_CASE							:     Optional[int]     ,		__SCREAMING_SNAKE_CASE							:     str):


           '''simple docstring'''

           __a						       =				output.encoder_hidden_states
           __a						       =				output.pixel_decoder_hidden_states
           __a						       =				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 		_lowerCamelCase			(		self							:     int     ,		__SCREAMING_SNAKE_CASE							:     Optional[Any]     ,		__SCREAMING_SNAKE_CASE							:     int     ,		__SCREAMING_SNAKE_CASE							:     Tuple     ,		__SCREAMING_SNAKE_CASE							:     List[str]=False):


           '''simple docstring'''

           with torch.no_grad():
               __a						       =				MaskFormerModel(config=__SCREAMING_SNAKE_CASE)
               model.to(__SCREAMING_SNAKE_CASE)
               model.eval()

               __a						       =				model(pixel_values=__SCREAMING_SNAKE_CASE     ,		pixel_mask=__SCREAMING_SNAKE_CASE)
               __a						       =				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 		_lowerCamelCase			(		self							:     int     ,		__SCREAMING_SNAKE_CASE							:     List[Any]     ,		__SCREAMING_SNAKE_CASE							:     Any     ,		__SCREAMING_SNAKE_CASE							:     Any     ,		__SCREAMING_SNAKE_CASE							:     Optional[Any]     ,		__SCREAMING_SNAKE_CASE							:     Union[str, Any]):


           '''simple docstring'''

           __a						       =				MaskFormerForInstanceSegmentation(config=__SCREAMING_SNAKE_CASE)
           model.to(__SCREAMING_SNAKE_CASE)
           model.eval()

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

           with torch.no_grad():
               __a						       =				model(pixel_values=__SCREAMING_SNAKE_CASE     ,		pixel_mask=__SCREAMING_SNAKE_CASE)
               __a						       =				model(__SCREAMING_SNAKE_CASE)

               comm_check_on_output(__SCREAMING_SNAKE_CASE)

               __a						       =				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     _A       (	__UpperCAmelCase      ,__UpperCAmelCase      ,unittest.TestCase    ):
       UpperCamelCase__  :    Union[str, Any]											=						(MaskFormerModel, MaskFormerForInstanceSegmentation) if is_torch_available() else ()
       UpperCamelCase__  :    Any											=						(
           {'''feature-extraction''': MaskFormerModel, '''image-segmentation''': MaskFormerForInstanceSegmentation}
           if is_torch_available()
           else {}
       )

       UpperCamelCase__  :    Optional[Any]											=						False
       UpperCamelCase__  :    Any											=						False
       UpperCamelCase__  :    Union[str, Any]											=						False
       UpperCamelCase__  :    Dict											=						False



       def 		_lowerCamelCase			(		self							:     Tuple):


           '''simple docstring'''

           __a						       =				MaskFormerModelTester(self)
           __a						       =				ConfigTester(self     ,		config_class=__SCREAMING_SNAKE_CASE     ,		has_text_modality=__SCREAMING_SNAKE_CASE)



       def 		_lowerCamelCase			(		self							:     Union[str, Any]):


           '''simple docstring'''

           self.config_tester.run_common_tests()



       def 		_lowerCamelCase			(		self							:     int):


           '''simple docstring'''

           __a       ,   __a						       =				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 		_lowerCamelCase			(		self							:     List[Any]):


           '''simple docstring'''

           __a						       =				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 		_lowerCamelCase			(		self							:     List[Any]):


           '''simple docstring'''

           pass



       @unittest.skip(reason='''MaskFormer does not have a get_input_embeddings method''')
       def 		_lowerCamelCase			(		self							:     List[str]):


           '''simple docstring'''

           pass



       @unittest.skip(reason='''MaskFormer is not a generative model''')
       def 		_lowerCamelCase			(		self							:     int):


           '''simple docstring'''

           pass



       @unittest.skip(reason='''MaskFormer does not use token embeddings''')
       def 		_lowerCamelCase			(		self							:     List[Any]):


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


           '''simple docstring'''

           pass



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


           '''simple docstring'''

           pass



       def 		_lowerCamelCase			(		self							:     Optional[int]):


           '''simple docstring'''

           __a       ,   __a						       =				self.model_tester.prepare_config_and_inputs_for_common()

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

               __a						       =				['''pixel_values''']
               self.assertListEqual(arg_names[:1]     ,		__SCREAMING_SNAKE_CASE)



       @slow
       def 		_lowerCamelCase			(		self							:     List[Any]):


           '''simple docstring'''

           for model_name in ["facebook/maskformer-swin-small-coco"]:
               __a						       =				MaskFormerModel.from_pretrained(__SCREAMING_SNAKE_CASE)
               self.assertIsNotNone(__SCREAMING_SNAKE_CASE)



       def 		_lowerCamelCase			(		self							:     Union[str, Any]):


           '''simple docstring'''

           __a						       =				(self.model_tester.min_size,) * 2
           __a						       =				{
               '''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(),
           }

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



       def 		_lowerCamelCase			(		self							:     Optional[Any]):


           '''simple docstring'''

           __a       ,   __a						       =				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 		_lowerCamelCase			(		self							:     Any):


           '''simple docstring'''

           __a       ,   __a						       =				self.model_tester.prepare_config_and_inputs_for_common()

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



       def 		_lowerCamelCase			(		self							:     str):


           '''simple docstring'''

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

           __a						       =				model_class(__SCREAMING_SNAKE_CASE)
           model.to(__SCREAMING_SNAKE_CASE)
           model.train()

           __a						       =				model(__SCREAMING_SNAKE_CASE     ,		mask_labels=__SCREAMING_SNAKE_CASE     ,		class_labels=__SCREAMING_SNAKE_CASE).loss
           loss.backward()





       def 		_lowerCamelCase			(		self							:     Tuple):


           '''simple docstring'''

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

           __a						       =				model_class(__SCREAMING_SNAKE_CASE)
           model.to(__SCREAMING_SNAKE_CASE)
           model.train()

           __a						       =				model(__SCREAMING_SNAKE_CASE     ,		mask_labels=__SCREAMING_SNAKE_CASE     ,		class_labels=__SCREAMING_SNAKE_CASE)

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

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

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

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


__snake_case						:int      		=     1E-4





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

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



       @cached_property
       def 		_lowerCamelCase			(		self							:     Union[str, Any]):


           '''simple docstring'''

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



       def 		_lowerCamelCase			(		self							:     Union[str, Any]):


           '''simple docstring'''

           __a						       =				MaskFormerModel.from_pretrained('''facebook/maskformer-swin-small-coco''').to(__SCREAMING_SNAKE_CASE)
           __a						       =				self.default_image_processor
           __a						       =				prepare_img()
           __a						       =				image_processor(__SCREAMING_SNAKE_CASE     ,		return_tensors='''pt''').to(__SCREAMING_SNAKE_CASE)
           __a						       =				inputs['''pixel_values'''].shape
           # check size is divisible by 32
           self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0)
           # check size
           self.assertEqual(__SCREAMING_SNAKE_CASE     ,		(1, 3, 800, 1_088))

           with torch.no_grad():
               __a						       =				model(**__SCREAMING_SNAKE_CASE)

           __a						       =				torch.tensor(
               [[-0.04_82, 0.92_28, 0.49_51], [-0.25_47, 0.80_17, 0.85_27], [-0.00_69, 0.33_85, -0.00_89]]).to(__SCREAMING_SNAKE_CASE)
           self.assertTrue(
               torch.allclose(
                   outputs.encoder_last_hidden_state[0, 0, :3, :3]     ,		__SCREAMING_SNAKE_CASE     ,		atol=__SCREAMING_SNAKE_CASE))

           __a						       =				torch.tensor(
               [[-0.84_22, -0.84_34, -0.97_18], [-1.01_44, -0.55_65, -0.41_95], [-1.00_38, -0.44_84, -0.19_61]]).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))

           __a						       =				torch.tensor(
               [[0.28_52, -0.01_59, 0.97_35], [0.62_54, 0.18_58, 0.85_29], [-0.06_80, -0.41_16, 1.84_13]]).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 		_lowerCamelCase			(		self							:     List[str]):


           '''simple docstring'''

           __a						       =				(
               MaskFormerForInstanceSegmentation.from_pretrained('''facebook/maskformer-swin-small-coco''')
               .to(__SCREAMING_SNAKE_CASE)
               .eval()
           )
           __a						       =				self.default_image_processor
           __a						       =				prepare_img()
           __a						       =				image_processor(__SCREAMING_SNAKE_CASE     ,		return_tensors='''pt''').to(__SCREAMING_SNAKE_CASE)
           __a						       =				inputs['''pixel_values'''].shape
           # check size is divisible by 32
           self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0)
           # check size
           self.assertEqual(__SCREAMING_SNAKE_CASE     ,		(1, 3, 800, 1_088))

           with torch.no_grad():
               __a						       =				model(**__SCREAMING_SNAKE_CASE)
           # masks_queries_logits
           __a						       =				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)     ,		)
           __a						       =				[
               [-1.3_73_71_24, -1.7_72_49_37, -1.9_36_42_33],
               [-1.5_97_72_81, -1.9_86_79_39, -2.1_52_36_95],
               [-1.5_79_53_98, -1.9_26_98_32, -2.09_39_42],
           ]
           __a						       =				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
           __a						       =				outputs.class_queries_logits
           self.assertEqual(
               class_queries_logits.shape     ,		(1, model.config.decoder_config.num_queries, model.config.num_labels + 1))
           __a						       =				torch.tensor(
               [
                   [1.6512E00, -5.2572E00, -3.3519E00],
                   [3.6169E-02, -5.9025E00, -2.9313E00],
                   [1.0766E-04, -7.7630E00, -5.1263E00],
               ]).to(__SCREAMING_SNAKE_CASE)
           self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]     ,		__SCREAMING_SNAKE_CASE     ,		atol=__SCREAMING_SNAKE_CASE))



       def 		_lowerCamelCase			(		self							:     str):


           '''simple docstring'''

           __a						       =				(
               MaskFormerForInstanceSegmentation.from_pretrained('''facebook/maskformer-resnet101-coco-stuff''')
               .to(__SCREAMING_SNAKE_CASE)
               .eval()
           )
           __a						       =				self.default_image_processor
           __a						       =				prepare_img()
           __a						       =				image_processor(__SCREAMING_SNAKE_CASE     ,		return_tensors='''pt''').to(__SCREAMING_SNAKE_CASE)
           __a						       =				inputs['''pixel_values'''].shape
           # check size is divisible by 32
           self.assertTrue((inputs_shape[-1] % 32) == 0 and (inputs_shape[-2] % 32) == 0)
           # check size
           self.assertEqual(__SCREAMING_SNAKE_CASE     ,		(1, 3, 800, 1_088))

           with torch.no_grad():
               __a						       =				model(**__SCREAMING_SNAKE_CASE)
           # masks_queries_logits
           __a						       =				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)     ,		)
           __a						       =				[[-0.90_46, -2.63_66, -4.60_62], [-3.41_79, -5.78_90, -8.80_57], [-4.91_79, -7.65_60, -10.77_11]]
           __a						       =				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
           __a						       =				outputs.class_queries_logits
           self.assertEqual(
               class_queries_logits.shape     ,		(1, model.config.decoder_config.num_queries, model.config.num_labels + 1))
           __a						       =				torch.tensor(
               [[4.71_88, -3.25_85, -2.88_57], [6.68_71, -2.91_81, -1.24_87], [7.24_49, -2.27_64, -2.18_74]]).to(__SCREAMING_SNAKE_CASE)
           self.assertTrue(torch.allclose(outputs.class_queries_logits[0, :3, :3]     ,		__SCREAMING_SNAKE_CASE     ,		atol=__SCREAMING_SNAKE_CASE))





       def 		_lowerCamelCase			(		self							:     Tuple):


           '''simple docstring'''

           __a						       =				(
               MaskFormerForInstanceSegmentation.from_pretrained('''facebook/maskformer-swin-small-coco''')
               .to(__SCREAMING_SNAKE_CASE)
               .eval()
           )
           __a						       =				self.default_image_processor

           __a						       =				image_processor(
               [np.zeros((3, 800, 1_333)), np.zeros((3, 800, 1_333))]     ,		segmentation_maps=[np.zeros((384, 384)).astype(np.floataa), np.zeros((384, 384)).astype(np.floataa)]     ,		return_tensors='''pt'''     ,		)

           __a						       =				inputs['''pixel_values'''].to(__SCREAMING_SNAKE_CASE)
           __a						       =				[el.to(__SCREAMING_SNAKE_CASE) for el in inputs['''mask_labels''']]
           __a						       =				[el.to(__SCREAMING_SNAKE_CASE) for el in inputs['''class_labels''']]

           with torch.no_grad():
               __a						       =				model(**__SCREAMING_SNAKE_CASE)

           self.assertTrue(outputs.loss is not None)










import argparse

from transformers import BigBirdConfig, BigBirdForPreTraining, BigBirdForQuestionAnswering, load_tf_weights_in_big_bird
from transformers.utils import logging


logging.set_verbosity_info()





def      __snake_case	(			_UpperCAmelCase   , _UpperCAmelCase   , _UpperCAmelCase   , _UpperCAmelCase				):
    # Initialise PyTorch model
    __a						       =				BigBirdConfig.from_json_file(_UpperCAmelCase				)
    print(f'Building PyTorch model from configuration: {config}'				)

    if is_trivia_qa:
        __a						       =				BigBirdForQuestionAnswering(_UpperCAmelCase				)
    else:
        __a						       =				BigBirdForPreTraining(_UpperCAmelCase				)

    # Load weights from tf checkpoint
    load_tf_weights_in_big_bird(_UpperCAmelCase   , _UpperCAmelCase   , is_trivia_qa=_UpperCAmelCase				)

    # Save pytorch-model
    print(f'Save PyTorch model to {pytorch_dump_path}'				)
    model.save_pretrained(_UpperCAmelCase				)


if __name__ == "__main__":
 __snake_case						:Tuple      		=     argparse.ArgumentParser()
 # Required parameters
 parser.add_argument(
     '''--tf_checkpoint_path''', default=None, type=str, required=True, help='''Path to the TensorFlow checkpoint path.'''
 )
 parser.add_argument(
     '''--big_bird_config_file''',
     default=None,
     type=str,
     required=True,
     help=(
         '''The config json file corresponding to the pre-trained BERT model. \n'''
         '''This specifies the model architecture.'''
     ),
 )
 parser.add_argument(
     '''--pytorch_dump_path''', default=None, type=str, required=True, help='''Path to the output PyTorch model.'''
 )
 parser.add_argument(
     '''--is_trivia_qa''', action='''store_true''', help='''Whether to convert a model with a trivia_qa head.'''
 )
 __snake_case						:Any      		=     parser.parse_args()
 convert_tf_checkpoint_to_pytorch(
     args.tf_checkpoint_path, args.big_bird_config_file, args.pytorch_dump_path, args.is_trivia_qa
 )





