





import collections
import inspect
import unittest

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

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch
  from torch import nn

  from transformers import DinatBackbone, DinatForImageClassification, DinatModel
  from transformers.models.dinat.modeling_dinat import DINAT_PRETRAINED_MODEL_ARCHIVE_LIST

if is_vision_available():
  from PIL import Image

  from transformers import AutoImageProcessor



class 				__snake_case  :


    '''simple docstring'''







    def __init__(					self      ,							A_      ,							A_=13      ,							A_=64      ,							A_=4      ,							A_=3      ,							A_=16      ,							A_=[1, 2, 1]      ,							A_=[2, 4, 8]      ,							A_=3      ,							A_=[[3], [1, 2], [1]]      ,							A_=2.0      ,							A_=True      ,							A_=0.0      ,							A_=0.0      ,							A_=0.1      ,							A_="gelu"      ,							A_=True      ,							A_=0.02      ,							A_=1E-5      ,							A_=True      ,							A_=None      ,							A_=True      ,							A_=10      ,							A_=["stage1", "stage2"]      ,							A_=[1, 2]      ,							):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= parent
        SCREAMING_SNAKE_CASE__       							= batch_size
        SCREAMING_SNAKE_CASE__       							= image_size
        SCREAMING_SNAKE_CASE__       							= patch_size
        SCREAMING_SNAKE_CASE__       							= num_channels
        SCREAMING_SNAKE_CASE__       							= embed_dim
        SCREAMING_SNAKE_CASE__       							= depths
        SCREAMING_SNAKE_CASE__       							= num_heads
        SCREAMING_SNAKE_CASE__       							= kernel_size
        SCREAMING_SNAKE_CASE__       							= dilations
        SCREAMING_SNAKE_CASE__       							= mlp_ratio
        SCREAMING_SNAKE_CASE__       							= qkv_bias
        SCREAMING_SNAKE_CASE__       							= hidden_dropout_prob
        SCREAMING_SNAKE_CASE__       							= attention_probs_dropout_prob
        SCREAMING_SNAKE_CASE__       							= drop_path_rate
        SCREAMING_SNAKE_CASE__       							= hidden_act
        SCREAMING_SNAKE_CASE__       							= patch_norm
        SCREAMING_SNAKE_CASE__       							= layer_norm_eps
        SCREAMING_SNAKE_CASE__       							= initializer_range
        SCREAMING_SNAKE_CASE__       							= is_training
        SCREAMING_SNAKE_CASE__       							= scope
        SCREAMING_SNAKE_CASE__       							= use_labels
        SCREAMING_SNAKE_CASE__       							= num_labels
        SCREAMING_SNAKE_CASE__       							= out_features
        SCREAMING_SNAKE_CASE__       							= out_indices



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]    )

        SCREAMING_SNAKE_CASE__       							= None
        if self.use_labels:
            SCREAMING_SNAKE_CASE__       							= ids_tensor([self.batch_size]      ,							self.num_labels    )

        SCREAMING_SNAKE_CASE__       							= self.get_config()

        return config, pixel_values, labels



    def    lowercase_				(					self    ):





        '''simple docstring'''

        return DinatConfig(
            num_labels=self.num_labels      ,							image_size=self.image_size      ,							patch_size=self.patch_size      ,							num_channels=self.num_channels      ,							embed_dim=self.embed_dim      ,							depths=self.depths      ,							num_heads=self.num_heads      ,							kernel_size=self.kernel_size      ,							dilations=self.dilations      ,							mlp_ratio=self.mlp_ratio      ,							qkv_bias=self.qkv_bias      ,							hidden_dropout_prob=self.hidden_dropout_prob      ,							attention_probs_dropout_prob=self.attention_probs_dropout_prob      ,							drop_path_rate=self.drop_path_rate      ,							hidden_act=self.hidden_act      ,							patch_norm=self.patch_norm      ,							layer_norm_eps=self.layer_norm_eps      ,							initializer_range=self.initializer_range      ,							out_features=self.out_features      ,							out_indices=self.out_indices      ,							)



    def    lowercase_				(					self      ,							A_      ,							A_      ,							A_    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= DinatModel(config=A_    )
        model.to(A_    )
        model.eval()
        SCREAMING_SNAKE_CASE__       							= model(A_    )

        SCREAMING_SNAKE_CASE__       							= SCREAMING_SNAKE_CASE__       							= (config.image_size // config.patch_size) // (2 ** (len(config.depths    ) - 1))
        SCREAMING_SNAKE_CASE__       							= int(config.embed_dim * 2 ** (len(config.depths    ) - 1)    )

        self.parent.assertEqual(
            result.last_hidden_state.shape      ,							(self.batch_size, expected_height, expected_width, expected_dim)    )



    def    lowercase_				(					self      ,							A_      ,							A_      ,							A_    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= DinatForImageClassification(A_    )
        model.to(A_    )
        model.eval()
        SCREAMING_SNAKE_CASE__       							= model(A_      ,							labels=A_    )
        self.parent.assertEqual(result.logits.shape      ,							(self.batch_size, self.num_labels)    )

        # test greyscale images
        SCREAMING_SNAKE_CASE__       							= 1
        SCREAMING_SNAKE_CASE__       							= DinatForImageClassification(A_    )
        model.to(A_    )
        model.eval()

        SCREAMING_SNAKE_CASE__       							= floats_tensor([self.batch_size, 1, self.image_size, self.image_size]    )
        SCREAMING_SNAKE_CASE__       							= model(A_    )
        self.parent.assertEqual(result.logits.shape      ,							(self.batch_size, self.num_labels)    )



    def    lowercase_				(					self      ,							A_      ,							A_      ,							A_    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= DinatBackbone(config=A_    )
        model.to(A_    )
        model.eval()
        SCREAMING_SNAKE_CASE__       							= model(A_    )

        # verify hidden states
        self.parent.assertEqual(len(result.feature_maps    )      ,							len(config.out_features    )    )
        self.parent.assertListEqual(list(result.feature_maps[0].shape    )      ,							[self.batch_size, model.channels[0], 16, 16]    )

        # verify channels
        self.parent.assertEqual(len(model.channels    )      ,							len(config.out_features    )    )

        # verify backbone works with out_features=None
        SCREAMING_SNAKE_CASE__       							= None
        SCREAMING_SNAKE_CASE__       							= DinatBackbone(config=A_    )
        model.to(A_    )
        model.eval()
        SCREAMING_SNAKE_CASE__       							= model(A_    )

        # verify feature maps
        self.parent.assertEqual(len(result.feature_maps    )      ,							1    )
        self.parent.assertListEqual(list(result.feature_maps[0].shape    )      ,							[self.batch_size, model.channels[-1], 4, 4]    )

        # verify channels
        self.parent.assertEqual(len(model.channels    )      ,							1    )




    def    lowercase_				(					self    ):





        '''simple docstring'''

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





@require_natten
@require_torch
class 				__snake_case  (     __SCREAMING_SNAKE_CASE ,    __SCREAMING_SNAKE_CASE ,    unittest.TestCase						):


    '''simple docstring'''







    lowerCamelCase__						:		Union[str, Any]							=  (
        (
            DinatModel,
            DinatForImageClassification,
            DinatBackbone,
        )
        if is_torch_available()
        else ()
    )
    lowerCamelCase__						:		str							=  (
        {"""feature-extraction""": DinatModel, """image-classification""": DinatForImageClassification}
        if is_torch_available()
        else {}
    )
    lowerCamelCase__						:		Optional[int]							=  False

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



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= DinatModelTester(self    )
        SCREAMING_SNAKE_CASE__       							= ConfigTester(self      ,							config_class=A_      ,							embed_dim=37    )



    def    lowercase_				(					self    ):





        '''simple docstring'''

        self.create_and_test_config_common_properties()
        self.config_tester.create_and_test_config_to_json_string()
        self.config_tester.create_and_test_config_to_json_file()
        self.config_tester.create_and_test_config_from_and_save_pretrained()
        self.config_tester.create_and_test_config_with_num_labels()
        self.config_tester.check_config_can_be_init_without_params()
        self.config_tester.check_config_arguments_init()



    def    lowercase_				(					self    ):





        '''simple docstring'''

        return



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_model(*A_    )



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_for_image_classification(*A_    )



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= self.model_tester.prepare_config_and_inputs()
        self.model_tester.create_and_check_backbone(*A_    )



    @unittest.skip(reason='''Dinat does not use inputs_embeds'''    )
    def    lowercase_				(					self    ):





        '''simple docstring'''

        pass



    @unittest.skip(reason='''Dinat does not use feedforward chunking'''    )
    def    lowercase_				(					self    ):





        '''simple docstring'''

        pass



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__							,		SCREAMING_SNAKE_CASE__       							= self.model_tester.prepare_config_and_inputs_for_common()

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



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__							,		SCREAMING_SNAKE_CASE__       							= self.model_tester.prepare_config_and_inputs_for_common()

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

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



    def    lowercase_				(					self    ):





        '''simple docstring'''

        self.skipTest('''Dinat\'s attention operation is handled entirely by NATTEN.'''    )



    def    lowercase_				(					self      ,							A_      ,							A_      ,							A_      ,							A_    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= model_class(A_    )
        model.to(A_    )
        model.eval()

        with torch.no_grad():
            SCREAMING_SNAKE_CASE__       							= model(**self._prepare_for_class(A_      ,							A_    )    )

        SCREAMING_SNAKE_CASE__       							= outputs.hidden_states

        SCREAMING_SNAKE_CASE__       							= getattr(
            self.model_tester      ,							'''expected_num_hidden_layers'''      ,							len(self.model_tester.depths    ) + 1    )
        self.assertEqual(len(A_    )      ,							A_    )

        # Dinat has a different seq_length
        SCREAMING_SNAKE_CASE__       							= (
            config.patch_size
            if isinstance(config.patch_size      ,							collections.abc.Iterable    )
            else (config.patch_size, config.patch_size)
        )

        SCREAMING_SNAKE_CASE__       							= image_size[0] // patch_size[0]
        SCREAMING_SNAKE_CASE__       							= image_size[1] // patch_size[1]

        self.assertListEqual(
            list(hidden_states[0].shape[-3:]    )      ,							[height, width, self.model_tester.embed_dim]      ,							)

        if model_class.__name__ != "DinatBackbone":
            SCREAMING_SNAKE_CASE__       							= outputs.reshaped_hidden_states
            self.assertEqual(len(A_    )      ,							A_    )

            SCREAMING_SNAKE_CASE__							,		SCREAMING_SNAKE_CASE__							,		SCREAMING_SNAKE_CASE__							,		SCREAMING_SNAKE_CASE__       							= reshaped_hidden_states[0].shape
            SCREAMING_SNAKE_CASE__       							= (
                reshaped_hidden_states[0].view(A_      ,							A_      ,							A_      ,							A_    ).permute(0      ,							2      ,							3      ,							1    )
            )
            self.assertListEqual(
                list(reshaped_hidden_states.shape[-3:]    )      ,							[height, width, self.model_tester.embed_dim]      ,							)



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__							,		SCREAMING_SNAKE_CASE__       							= self.model_tester.prepare_config_and_inputs_for_common()

        SCREAMING_SNAKE_CASE__       							= (
            self.model_tester.image_size
            if isinstance(self.model_tester.image_size      ,							collections.abc.Iterable    )
            else (self.model_tester.image_size, self.model_tester.image_size)
        )

        for model_class in self.all_model_classes:
            SCREAMING_SNAKE_CASE__       							= True
            self.check_hidden_states_output(A_      ,							A_      ,							A_      ,							A_    )

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

            self.check_hidden_states_output(A_      ,							A_      ,							A_      ,							A_    )



    @slow
    def    lowercase_				(					self    ):





        '''simple docstring'''

        for model_name in DINAT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
            SCREAMING_SNAKE_CASE__       							= DinatModel.from_pretrained(A_    )
            self.assertIsNotNone(A_    )




    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__							,		SCREAMING_SNAKE_CASE__       							= self.model_tester.prepare_config_and_inputs_for_common()

        SCREAMING_SNAKE_CASE__       							= _config_zero_init(A_    )
        for model_class in self.all_model_classes:
            SCREAMING_SNAKE_CASE__       							= model_class(config=A_    )
            for name, param in model.named_parameters():
                if "embeddings" not in name and param.requires_grad:
                    self.assertIn(
                        ((param.data.mean() * 1E9).round() / 1E9).item()      ,							[0.0, 1.0]      ,							msg=f'''Parameter {name} of model {model_class} seems not properly initialized'''      ,							)





@require_natten
@require_vision
@require_torch
class 				__snake_case  (     unittest.TestCase						):


    '''simple docstring'''







    @cached_property
    def    lowercase_				(					self    ):





        '''simple docstring'''

        return AutoImageProcessor.from_pretrained('''shi-labs/dinat-mini-in1k-224'''    ) if is_vision_available() else None




    @slow
    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= DinatForImageClassification.from_pretrained('''shi-labs/dinat-mini-in1k-224'''    ).to(A_    )
        SCREAMING_SNAKE_CASE__       							= self.default_image_processor

        SCREAMING_SNAKE_CASE__       							= Image.open('''./tests/fixtures/tests_samples/COCO/000000039769.png'''    )
        SCREAMING_SNAKE_CASE__       							= image_processor(images=A_      ,							return_tensors='''pt'''    ).to(A_    )

        # forward pass
        with torch.no_grad():
            SCREAMING_SNAKE_CASE__       							= model(**A_    )

        # verify the logits
        SCREAMING_SNAKE_CASE__       							= torch.Size((1, 10_00)    )
        self.assertEqual(outputs.logits.shape      ,							A_    )
        SCREAMING_SNAKE_CASE__       							= torch.tensor([-0.1545, -0.7667, 0.4642]    ).to(A_    )
        self.assertTrue(torch.allclose(outputs.logits[0, :3]      ,							A_      ,							atol=1E-4    )    )





@require_torch
@require_natten
class 				__snake_case  (     unittest.TestCase ,    __SCREAMING_SNAKE_CASE						):


    '''simple docstring'''







    lowerCamelCase__						:		Union[str, Any]							=  (DinatBackbone,) if is_torch_available() else ()
    lowerCamelCase__						:		Union[str, Any]							=  DinatConfig



    def    lowercase_				(					self    ):





        '''simple docstring'''

        SCREAMING_SNAKE_CASE__       							= DinatModelTester(self    )







import os
from pathlib import Path
from typing import Dict, List

import fire
import torch

from transformers import AutoModelForSeqaSeqLM, AutoTokenizer
from transformers.utils.logging import get_logger


_A				:							Any							= get_logger(__name__)


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_  )							->	str:
    if text.startswith(lowerCAmelCase_  ):
        return text[len(lowerCAmelCase_  ) :]
    return text  # or whatever


def        __snake_case       (						lowerCAmelCase_  )							->	Union[str, Any]:
    return {remove_prefix(lowerCAmelCase_   ,       '''model.'''  ): v for k, v in sd.items()}


def        __snake_case       (						lowerCAmelCase_  )							->	Optional[int]:
    SCREAMING_SNAKE_CASE__       							= {}
    for k in state_dicts[0].keys():
        SCREAMING_SNAKE_CASE__       							= [sd[k] for sd in state_dicts]
        SCREAMING_SNAKE_CASE__       							= sum(lowerCAmelCase_  ) / len(lowerCAmelCase_  )
        assert isinstance(lowerCAmelCase_   ,       torch.Tensor  )
        SCREAMING_SNAKE_CASE__       							= new_t
    return new_sd


def        __snake_case       (						lowerCAmelCase_   ,       lowerCAmelCase_   ,       lowerCAmelCase_  )							->	None:
    SCREAMING_SNAKE_CASE__       							= AutoModelForSeqaSeqLM.from_pretrained(lowerCAmelCase_  )
    if os.path.isfile(lowerCAmelCase_  ):
        SCREAMING_SNAKE_CASE__       							= [pl_ckpt_path]
    else:
        assert os.path.isdir(lowerCAmelCase_  )
        SCREAMING_SNAKE_CASE__       							= list(Path(lowerCAmelCase_  ).glob('''*.ckpt'''  )  )
        assert ckpt_files, f'''could not find any ckpt files inside the {pl_ckpt_path} directory'''

    if len(lowerCAmelCase_  ) > 1:
        logger.info(f'''averaging the weights of {ckpt_files}'''  )

    SCREAMING_SNAKE_CASE__       							= [sanitize(torch.load(lowerCAmelCase_   ,       map_location='''cpu'''  )['''state_dict''']  ) for x in ckpt_files]
    SCREAMING_SNAKE_CASE__       							= average_state_dicts(lowerCAmelCase_  )

    SCREAMING_SNAKE_CASE__							,		SCREAMING_SNAKE_CASE__       							= hf_model.load_state_dict(lowerCAmelCase_   ,       strict=lowerCAmelCase_  )
    assert not missing, f'''missing keys: {missing}'''
    hf_model.save_pretrained(lowerCAmelCase_  )
    try:
        SCREAMING_SNAKE_CASE__       							= AutoTokenizer.from_pretrained(lowerCAmelCase_  )
        tok.save_pretrained(lowerCAmelCase_  )
    except Exception:
        pass
        # dont copy tokenizer if cant


if __name__ == "__main__":
  fire.Fire(convert_pl_to_hf)
