





import inspect
import unittest

from transformers import YolosConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
       import torch
       from torch import nn

       from transformers import YolosForObjectDetection, YolosModel
       from transformers.models.yolos.modeling_yolos import YOLOS_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
       from PIL import Image

       from transformers import AutoImageProcessor




class        SCREAMING_SNAKE_CASE__					:



 def __init__(				self,__lowerCamelCase,__lowerCamelCase=13,__lowerCamelCase=[30, 30],__lowerCamelCase=2,__lowerCamelCase=3,__lowerCamelCase=True,__lowerCamelCase=True,__lowerCamelCase=32,__lowerCamelCase=5,__lowerCamelCase=4,__lowerCamelCase=37,__lowerCamelCase="gelu",__lowerCamelCase=0.1,__lowerCamelCase=0.1,__lowerCamelCase=10,__lowerCamelCase=0.02,__lowerCamelCase=3,__lowerCamelCase=None,__lowerCamelCase=8,__lowerCamelCase=10,):
        A__		   =		parent
        A__		   =		batch_size
        A__		   =		image_size
        A__		   =		patch_size
        A__		   =		num_channels
        A__		   =		is_training
        A__		   =		use_labels
        A__		   =		hidden_size
        A__		   =		num_hidden_layers
        A__		   =		num_attention_heads
        A__		   =		intermediate_size
        A__		   =		hidden_act
        A__		   =		hidden_dropout_prob
        A__		   =		attention_probs_dropout_prob
        A__		   =		type_sequence_label_size
        A__		   =		initializer_range
        A__		   =		num_labels
        A__		   =		scope
        A__		   =		n_targets
        A__		   =		num_detection_tokens
        # we set the expected sequence length (which is used in several tests)
        # expected sequence length = num_patches + 1 (we add 1 for the [CLS] token) + num_detection_tokens
        A__		   =		(image_size[1] // patch_size) * (image_size[0] // patch_size)
        A__		   =		num_patches + 1 + self.num_detection_tokens



 def 							UpperCamelCase     (				self					):
        A__		   =		floats_tensor([self.batch_size, self.num_channels, self.image_size[0], self.image_size[1]]					)

        A__		   =		None
        if self.use_labels:
               # labels is a list of Dict (each Dict being the labels for a given example in the batch)
               A__		   =		[]
               for i in range(self.batch_size					):
                      A__		   =		{}
                      A__		   =		torch.randint(
                          high=self.num_labels,size=(self.n_targets,),device=__lowerCamelCase					)
                      A__		   =		torch.rand(self.n_targets,4,device=__lowerCamelCase					)
                      labels.append(__lowerCamelCase					)

        A__		   =		self.get_config()

        return config, pixel_values, labels



 def 							UpperCamelCase     (				self					):
        return YolosConfig(
            image_size=self.image_size,patch_size=self.patch_size,num_channels=self.num_channels,hidden_size=self.hidden_size,num_hidden_layers=self.num_hidden_layers,num_attention_heads=self.num_attention_heads,intermediate_size=self.intermediate_size,hidden_act=self.hidden_act,hidden_dropout_prob=self.hidden_dropout_prob,attention_probs_dropout_prob=self.attention_probs_dropout_prob,is_decoder=__lowerCamelCase,initializer_range=self.initializer_range,num_detection_tokens=self.num_detection_tokens,num_labels=self.num_labels,)



 def 							UpperCamelCase     (				self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase					):
        A__		   =		YolosModel(config=__lowerCamelCase					)
        model.to(__lowerCamelCase					)
        model.eval()
        A__		   =		model(__lowerCamelCase					)
        self.parent.assertEqual(
            result.last_hidden_state.shape,(self.batch_size, self.expected_seq_len, self.hidden_size)					)



 def 							UpperCamelCase     (				self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase					):
        A__		   =		YolosForObjectDetection(__lowerCamelCase					)
        model.to(__lowerCamelCase					)
        model.eval()

        A__		   =		model(pixel_values=__lowerCamelCase					)
        A__		   =		model(__lowerCamelCase					)

        self.parent.assertEqual(result.logits.shape,(self.batch_size, self.num_detection_tokens, self.num_labels + 1)					)
        self.parent.assertEqual(result.pred_boxes.shape,(self.batch_size, self.num_detection_tokens, 4)					)

        A__		   =		model(pixel_values=__lowerCamelCase,labels=__lowerCamelCase					)

        self.parent.assertEqual(result.loss.shape,()					)
        self.parent.assertEqual(result.logits.shape,(self.batch_size, self.num_detection_tokens, self.num_labels + 1)					)
        self.parent.assertEqual(result.pred_boxes.shape,(self.batch_size, self.num_detection_tokens, 4)					)




 def 							UpperCamelCase     (				self					):
        A__		   =		self.prepare_config_and_inputs()
        A__   ,      A__   ,      A__		   =		config_and_inputs
        A__		   =		{'''pixel_values''': pixel_values}
        return config, inputs_dict




@require_torch
class        SCREAMING_SNAKE_CASE__					(		UpperCamelCase__							,   UpperCamelCase__							,   unittest.TestCase ):
 __SCREAMING_SNAKE_CASE							=      (YolosModel, YolosForObjectDetection) if is_torch_available() else ()
 __SCREAMING_SNAKE_CASE							=      (
     {'''feature-extraction''': YolosModel, '''object-detection''': YolosForObjectDetection} if is_torch_available() else {}
 )

 __SCREAMING_SNAKE_CASE							=      False
 __SCREAMING_SNAKE_CASE							=      False
 __SCREAMING_SNAKE_CASE							=      False
 __SCREAMING_SNAKE_CASE							=      False



 def 							UpperCamelCase     (				self,__lowerCamelCase,__lowerCamelCase,__lowerCamelCase=False					):
        A__		   =		super()._prepare_for_class(__lowerCamelCase,__lowerCamelCase,return_labels=__lowerCamelCase					)

        if return_labels:
               if model_class.__name__ == "YolosForObjectDetection":
                      A__		   =		[]
                      for i in range(self.model_tester.batch_size					):
                             A__		   =		{}
                             A__		   =		torch.ones(
                                 size=(self.model_tester.n_targets,),device=__lowerCamelCase,dtype=torch.long					)
                             A__		   =		torch.ones(
                                 self.model_tester.n_targets,4,device=__lowerCamelCase,dtype=torch.float					)
                             labels.append(__lowerCamelCase					)
                      A__		   =		labels

        return inputs_dict



 def 							UpperCamelCase     (				self					):
        A__		   =		YolosModelTester(self					)
        A__		   =		ConfigTester(self,config_class=__lowerCamelCase,has_text_modality=__lowerCamelCase,hidden_size=37					)



 def 							UpperCamelCase     (				self					):
        self.config_tester.run_common_tests()



 def 							UpperCamelCase     (				self					):
        # YOLOS does not use inputs_embeds
        pass



 def 							UpperCamelCase     (				self					):
        A__   ,      A__		   =		self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
               A__		   =		model_class(__lowerCamelCase					)
               self.assertIsInstance(model.get_input_embeddings(),(nn.Module)					)
               A__		   =		model.get_output_embeddings()
               self.assertTrue(x is None or isinstance(__lowerCamelCase,nn.Linear					)					)



 def 							UpperCamelCase     (				self					):
        A__   ,      A__		   =		self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
               A__		   =		model_class(__lowerCamelCase					)
               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],__lowerCamelCase					)



 def 							UpperCamelCase     (				self					):
        A__		   =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*__lowerCamelCase					)



 def 							UpperCamelCase     (				self					):
        A__   ,      A__		   =		self.model_tester.prepare_config_and_inputs_for_common()
        A__		   =		True

        # in YOLOS, the seq_len is different
        A__		   =		self.model_tester.expected_seq_len
        for model_class in self.all_model_classes:
               A__		   =		True
               A__		   =		False
               A__		   =		True
               A__		   =		model_class(__lowerCamelCase					)
               model.to(__lowerCamelCase					)
               model.eval()
               with torch.no_grad():
                      A__		   =		model(**self._prepare_for_class(__lowerCamelCase,__lowerCamelCase					)					)
               A__		   =		outputs.attentions
               self.assertEqual(len(__lowerCamelCase					),self.model_tester.num_hidden_layers					)

               # check that output_attentions also work using config
               del inputs_dict["output_attentions"]
               A__		   =		True
               A__		   =		model_class(__lowerCamelCase					)
               model.to(__lowerCamelCase					)
               model.eval()
               with torch.no_grad():
                      A__		   =		model(**self._prepare_for_class(__lowerCamelCase,__lowerCamelCase					)					)
               A__		   =		outputs.attentions
               self.assertEqual(len(__lowerCamelCase					),self.model_tester.num_hidden_layers					)

               self.assertListEqual(
                   list(attentions[0].shape[-3:]					),[self.model_tester.num_attention_heads, seq_len, seq_len],)
               A__		   =		len(__lowerCamelCase					)

               # Check attention is always last and order is fine
               A__		   =		True
               A__		   =		True
               A__		   =		model_class(__lowerCamelCase					)
               model.to(__lowerCamelCase					)
               model.eval()
               with torch.no_grad():
                      A__		   =		model(**self._prepare_for_class(__lowerCamelCase,__lowerCamelCase					)					)

               A__		   =		1
               self.assertEqual(out_len + added_hidden_states,len(__lowerCamelCase					)					)

               A__		   =		outputs.attentions

               self.assertEqual(len(__lowerCamelCase					),self.model_tester.num_hidden_layers					)
               self.assertListEqual(
                   list(self_attentions[0].shape[-3:]					),[self.model_tester.num_attention_heads, seq_len, seq_len],)



 def 							UpperCamelCase     (				self					):
        def check_hidden_states_output(__lowerCamelCase,__lowerCamelCase,__lowerCamelCase					):
               A__		   =		model_class(__lowerCamelCase					)
               model.to(__lowerCamelCase					)
               model.eval()

               with torch.no_grad():
                      A__		   =		model(**self._prepare_for_class(__lowerCamelCase,__lowerCamelCase					)					)

               A__		   =		outputs.hidden_states

               A__		   =		getattr(
                   self.model_tester,'''expected_num_hidden_layers''',self.model_tester.num_hidden_layers + 1					)
               self.assertEqual(len(__lowerCamelCase					),__lowerCamelCase					)

               # YOLOS has a different seq_length
               A__		   =		self.model_tester.expected_seq_len

               self.assertListEqual(
                   list(hidden_states[0].shape[-2:]					),[seq_length, self.model_tester.hidden_size],)

        A__   ,      A__		   =		self.model_tester.prepare_config_and_inputs_for_common()

        for model_class in self.all_model_classes:
               A__		   =		True
               check_hidden_states_output(__lowerCamelCase,__lowerCamelCase,__lowerCamelCase					)

               # check that output_hidden_states also work using config
               del inputs_dict["output_hidden_states"]
               A__		   =		True

               check_hidden_states_output(__lowerCamelCase,__lowerCamelCase,__lowerCamelCase					)



 def 							UpperCamelCase     (				self					):
        A__		   =		self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_object_detection(*__lowerCamelCase					)




 @slow
 def 							UpperCamelCase     (				self					):
        for model_name in YOLOS_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               A__		   =		YolosModel.from_pretrained(__lowerCamelCase					)
               self.assertIsNotNone(__lowerCamelCase					)


def       UpperCamelCase__(			)->Optional[Any]:
       A__		   =		Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''    )
       return image




@require_torch
@require_vision
class        SCREAMING_SNAKE_CASE__					(		unittest.TestCase ):



 @cached_property
 def 							UpperCamelCase     (				self					):
        return AutoImageProcessor.from_pretrained('''hustvl/yolos-small'''					) if is_vision_available() else None




 @slow
 def 							UpperCamelCase     (				self					):
        A__		   =		YolosForObjectDetection.from_pretrained('''hustvl/yolos-small'''					).to(__lowerCamelCase					)

        A__		   =		self.default_image_processor
        A__		   =		prepare_img()
        A__		   =		image_processor(images=__lowerCamelCase,return_tensors='''pt'''					).to(__lowerCamelCase					)

        # forward pass
        with torch.no_grad():
               A__		   =		model(inputs.pixel_values					)

        # verify outputs
        A__		   =		torch.Size((1, 100, 92)					)
        self.assertEqual(outputs.logits.shape,__lowerCamelCase					)

        A__		   =		torch.tensor(
            [[-24.0248, -10.3024, -14.8290], [-42.0392, -16.8200, -27.4334], [-27.2743, -11.8154, -18.7148]],device=__lowerCamelCase,)
        A__		   =		torch.tensor(
            [[0.2559, 0.5455, 0.4706], [0.2989, 0.7279, 0.1875], [0.7732, 0.4017, 0.4462]],device=__lowerCamelCase					)
        self.assertTrue(torch.allclose(outputs.logits[0, :3, :3],__lowerCamelCase,atol=1E-4					)					)
        self.assertTrue(torch.allclose(outputs.pred_boxes[0, :3, :3],__lowerCamelCase,atol=1E-4					)					)

        # verify postprocessing
        A__		   =		image_processor.post_process_object_detection(
            __lowerCamelCase,threshold=0.3,target_sizes=[image.size[::-1]]					)[0]
        A__		   =		torch.tensor([0.9994, 0.9790, 0.9964, 0.9972, 0.9861]					).to(__lowerCamelCase					)
        A__		   =		[75, 75, 17, 63, 17]
        A__		   =		torch.tensor([335.0609, 79.3848, 375.4216, 187.2495]					).to(__lowerCamelCase					)

        self.assertEqual(len(results['''scores''']					),5					)
        self.assertTrue(torch.allclose(results['''scores'''],__lowerCamelCase,atol=1E-4					)					)
        self.assertSequenceEqual(results['''labels'''].tolist(),__lowerCamelCase					)
        self.assertTrue(torch.allclose(results['''boxes'''][0, :],__lowerCamelCase					)					)












import argparse

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


logging.set_verbosity_info()


def       UpperCamelCase__(			UpperCamelCase__     :			Any					,	UpperCamelCase__     :			Optional[Any]					,	UpperCamelCase__     :			Dict					,	UpperCamelCase__     :			int    )->Union[str, Any]:
       # 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__":
       a__: Optional[int]							   =						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.'
       )
       a__: Tuple							   =						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
       )





